api/app/core/deep_agent.py

"""Deep Agent — LangGraph hierarchical supervisors for home and floating modes.

Two supervisor graphs (both ``create_react_agent``):
  * **HomeSupervisor**     — gathers data from multiple domains, presents
                             structured overview with tool-result blocks.
  * **FloatingSupervisor** — focused, scoped assistant for a single entity/domain.

Each supervisor delegates to four sub-agent tools, each a compiled
``create_react_agent`` wrapping the domain CRUD tools (task, project, note,
timeline).  The sub-agents talk to Electron via ``execute_on_client``.

Streaming uses ``astream(stream_mode=["messages", "updates"])`` so that
callers can sniff:
  * ``("messages", (token, metadata))``  — text tokens for streaming
  * ``("updates", ...)``                 — tool call results for mutations

An ``update_core_memory`` tool is available to both supervisors for
persisting user preferences mid-conversation (MemGPT-style).
"""

from __future__ import annotations

import json
import logging
from typing import Any, AsyncGenerator

from langchain_core.messages import AIMessageChunk, HumanMessage
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent

from app.core.llm import get_llm
from app.core.ws_context import (
    clear_tool_result_collector,
    set_tool_result_collector,
)

logger = logging.getLogger(__name__)

# ── Sub-agent tool imports ────────────────────────────────────────────

from app.agents.task_agent import (  # noqa: E402
    add_task_comment,
    create_task,
    delete_task,
    delete_task_comment,
    list_task_comments,
    list_tasks,
    list_tasks_due_today,
    update_task,
)
from app.agents.note_agent import (  # noqa: E402
    create_note,
    delete_note,
    get_note,
    list_notes,
    update_note,
)
from app.agents.project_agent import (  # noqa: E402
    create_project,
    delete_project,
    get_project,
    list_all_projects,
    list_projects,
    update_project,
)
from app.agents.timeline_agent import (  # noqa: E402
    create_timeline,
    delete_timeline,
    list_timelines,
    update_timeline,
)

# ── Sub-agent definitions ─────────────────────────────────────────────

_TASK_TOOLS = [
    list_tasks,
    create_task,
    update_task,
    delete_task,
    list_tasks_due_today,
    list_task_comments,
    add_task_comment,
    delete_task_comment,
]

_NOTE_TOOLS = [list_notes, get_note, create_note, update_note, delete_note]

_PROJECT_TOOLS = [
    list_projects,
    list_all_projects,
    get_project,
    create_project,
    update_project,
    delete_project,
]

_TIMELINE_TOOLS = [list_timelines, create_timeline, update_timeline, delete_timeline]


def _build_subagent_tool(
    name: str,
    description: str,
    system_prompt: str,
    tools: list,
):
    """Build a compiled sub-agent graph and wrap it as a LangChain tool."""
    subgraph = create_react_agent(
        model=get_llm(),
        tools=tools,
        prompt=system_prompt,
        name=name,
    )

    @tool(name=name, description=description)
    async def _run(query: str) -> str:
        result = await subgraph.ainvoke(
            {"messages": [HumanMessage(content=query)]}
        )
        messages = result["messages"]
        # Return the last AI message content
        for msg in reversed(messages):
            if hasattr(msg, "content") and msg.content and not getattr(msg, "tool_calls", None):
                return str(msg.content)
        return "No response from sub-agent."

    return _run


def _make_subagent_tools() -> list:
    """Create the four sub-agent tools for the supervisor."""
    return [
        _build_subagent_tool(
            name="task_agent",
            description=(
                "Manages tasks and comments: list, create, update, delete, "
                "due-today, comments. Delegate task-related queries here."
            ),
            system_prompt=(
                "You are a task management assistant. You create, update, list, "
                "and track tasks and their comments.\n\n"
                "Rules:\n"
                "  - status must be one of: todo, in_progress, done\n"
                "  - priority must be one of: high, medium, low\n"
                "  - due_date is a Unix timestamp in milliseconds\n"
                "  - assignees is a JSON-encoded array of strings\n"
                "  - is_approved defaults to 0; set to 1 only when the user confirms\n"
                "  - For update_task, use -1 for integer fields you do not want to change\n"
                "  - Always confirm the action in plain, user-friendly language."
            ),
            tools=_TASK_TOOLS,
        ),
        _build_subagent_tool(
            name="note_agent",
            description=(
                "Manages notes: list, get, create, update, delete. "
                "Delegate note-related queries here."
            ),
            system_prompt=(
                "You are a note-taking assistant. You help users create, retrieve, "
                "update, and delete Markdown notes in their workspace.\n\n"
                "Rules:\n"
                "  - content is always Markdown; preserve formatting when updating\n"
                "  - When updating, call get_note first if you need to read existing "
                "content before appending or replacing sections\n"
                "  - Do not fabricate note content."
            ),
            tools=_NOTE_TOOLS,
        ),
        _build_subagent_tool(
            name="project_agent",
            description=(
                "Manages projects: list, get, create, update, archive, delete. "
                "Delegate project-related queries here."
            ),
            system_prompt=(
                "You are a project management assistant. You help users create, "
                "find, update, and archive projects.\n\n"
                "Rules:\n"
                "  - status must be one of: active, archived\n"
                "  - Prefer archiving over deletion\n"
                "  - ai_summary is populated only when the user asks for a summary."
            ),
            tools=_PROJECT_TOOLS,
        ),
        _build_subagent_tool(
            name="timeline_agent",
            description=(
                "Manages project timelines (milestones): list, create, update, "
                "delete. Delegate timeline/milestone queries here."
            ),
            system_prompt=(
                "You are a project timeline assistant. Timelines are milestone "
                "dates that track progress on a project.\n\n"
                "Rules:\n"
                "  - project_id is REQUIRED for every create\n"
                "  - date is a Unix timestamp in milliseconds\n"
                "  - For update_timeline, use -1 for integer fields you do not "
                "want to change."
            ),
            tools=_TIMELINE_TOOLS,
        ),
    ]


# ── Update core memory tool ──────────────────────────────────────────

def _make_update_core_memory_tool(user_id: str, db_session_factory):
    """Create a tool that persists a key/value preference in core memory."""

    @tool
    async def update_core_memory(key: str, value: str) -> str:
        """Save a user preference or fact to long-term core memory.
        key: short label for the memory (e.g. 'preferred_language', 'timezone')
        value: the value to remember
        Use this when the user states a preference or fact worth remembering.
        """
        from app.core.memory_middleware import MemoryMiddleware

        async with db_session_factory() as db:
            memory = MemoryMiddleware(db)
            await memory.update_core(user_id, key, value)
        return f"Remembered: {key} = {value}"

    return update_core_memory


# ── System prompts ────────────────────────────────────────────────────

_HOME_SYSTEM = (
    "You are Adiuva, a smart workspace assistant on the Home dashboard.\n"
    "Your job is to help the user by gathering data from their workspace and "
    "presenting a comprehensive overview.\n\n"
    "You have sub-agent tools (task_agent, note_agent, project_agent, "
    "timeline_agent) that can query and mutate workspace data. Delegate to "
    "the appropriate sub-agent(s) based on the user's request. You can call "
    "multiple sub-agents if needed.\n\n"
    "You also have an update_core_memory tool — use it when the user states "
    "a preference or important fact worth remembering long-term.\n\n"
    "After gathering data, synthesize a clear, helpful response for the user.\n\n"
    "Memory context:\n{memory_context}"
)

_FLOATING_SYSTEM = (
    "You are Adiuva, a focused workspace assistant in the floating panel.\n"
    "The user is currently working in the '{scope_type}' section"
    "{scope_detail}.\n\n"
    "You have sub-agent tools (task_agent, note_agent, project_agent, "
    "timeline_agent) that can query and mutate workspace data. Focus your "
    "help on the user's current scope, but you can use other sub-agents "
    "if the request requires it.\n\n"
    "You also have an update_core_memory tool — use it when the user states "
    "a preference or important fact worth remembering long-term.\n\n"
    "Provide direct, conversational responses.\n\n"
    "Memory context:\n{memory_context}"
)


def _format_memory_context(memory: dict[str, Any]) -> str:
    """Format the memory dict into a readable string for the system prompt."""
    if not memory:
        return "(no memory available)"
    parts = []
    if memory.get("core_memory"):
        parts.append("Preferences: " + json.dumps(memory["core_memory"]))
    if memory.get("associative_memory"):
        parts.append("Related memories: " + "; ".join(memory["associative_memory"][:3]))
    if memory.get("episodic_memory"):
        parts.append("Recent sessions: " + "; ".join(memory["episodic_memory"][:3]))
    if memory.get("proactive_hints"):
        parts.append("Patterns: " + "; ".join(memory["proactive_hints"][:3]))
    return "\n".join(parts) if parts else "(no memory available)"


# ── Graph builders ────────────────────────────────────────────────────

def build_home_graph(
    user_id: str,
    memory_context: dict[str, Any],
    db_session_factory,
):
    """Build the Home supervisor graph."""
    subagent_tools = _make_subagent_tools()
    memory_tool = _make_update_core_memory_tool(user_id, db_session_factory)
    all_tools = subagent_tools + [memory_tool]

    prompt = _HOME_SYSTEM.format(
        memory_context=_format_memory_context(memory_context),
    )

    return create_react_agent(
        model=get_llm(),
        tools=all_tools,
        prompt=prompt,
        name="home_supervisor",
    )


def build_floating_graph(
    user_id: str,
    memory_context: dict[str, Any],
    scope: dict[str, Any],
    db_session_factory,
):
    """Build the Floating supervisor graph."""
    subagent_tools = _make_subagent_tools()
    memory_tool = _make_update_core_memory_tool(user_id, db_session_factory)
    all_tools = subagent_tools + [memory_tool]

    scope_type = scope.get("type", "general")
    scope_id = scope.get("id")
    scope_detail = f" (id: {scope_id})" if scope_id else ""

    prompt = _FLOATING_SYSTEM.format(
        scope_type=scope_type,
        scope_detail=scope_detail,
        memory_context=_format_memory_context(memory_context),
    )

    return create_react_agent(
        model=get_llm(),
        tools=all_tools,
        prompt=prompt,
        name="floating_supervisor",
    )


# ── Stream event type ────────────────────────────────────────────────

# Events yielded by run_*_stream:
#   ("token",      str)           — text token for streaming
#   ("tool_start", dict)          — {"name": "task_agent", "args": {...}}
#   ("tool_end",   dict)          — {"name": "task_agent", "result": "..."}


# ── Stream runners ────────────────────────────────────────────────────

async def _run_graph_stream(
    graph,
    message: str,
) -> AsyncGenerator[tuple[str, Any], None]:
    """Run a supervisor graph with streaming, yielding event tuples.

    Uses ``stream_mode=["messages", "updates"]`` to get both token-level
    streaming and update events for tool calls.
    """
    inputs = {"messages": [HumanMessage(content=message)]}

    collector: list[dict] = []
    set_tool_result_collector(collector)
    try:
        async for stream_mode, chunk in graph.astream(
            inputs,
            stream_mode=["messages", "updates"],
        ):
            if stream_mode == "messages":
                msg, metadata = chunk
                # Only yield tokens from the supervisor's final response
                # (not from sub-agent internal LLM calls)
                if (
                    isinstance(msg, AIMessageChunk)
                    and msg.content
                    and not msg.tool_calls
                    and metadata.get("langgraph_node") == "agent"
                ):
                    yield ("token", str(msg.content))

            elif stream_mode == "updates":
                # Updates is a dict of {node_name: state_update}
                if not isinstance(chunk, dict):
                    continue
                for node_name, state_update in chunk.items():
                    if node_name != "tools":
                        continue
                    # Tool node executed — extract tool call results
                    tool_messages = state_update.get("messages", [])
                    for tool_msg in tool_messages:
                        if hasattr(tool_msg, "name") and hasattr(tool_msg, "content"):
                            yield (
                                "tool_end",
                                {"name": tool_msg.name, "result": str(tool_msg.content)},
                            )
    finally:
        clear_tool_result_collector()

    # Yield the collected mutations so callers can attach them to stream_end
    yield ("mutations", collector)


async def run_home_stream(
    user_id: str,
    message: str,
    context: dict[str, Any],
    db_session_factory,
) -> AsyncGenerator[tuple[str, Any], None]:
    """Run the Home supervisor and yield streaming events."""
    graph = build_home_graph(user_id, context, db_session_factory)
    async for event in _run_graph_stream(graph, message):
        yield event


async def run_floating_stream(
    user_id: str,
    message: str,
    context: dict[str, Any],
    scope: dict[str, Any],
    db_session_factory,
) -> AsyncGenerator[tuple[str, Any], None]:
    """Run the Floating supervisor and yield streaming events."""
    graph = build_floating_graph(user_id, context, scope, db_session_factory)
    async for event in _run_graph_stream(graph, message):
        yield event


async def run_home(
    user_id: str,
    message: str,
    context: dict[str, Any],
    db_session_factory,
) -> str:
    """Run the Home supervisor (non-streaming) and return full response text."""
    graph = build_home_graph(user_id, context, db_session_factory)
    result = await graph.ainvoke(
        {"messages": [HumanMessage(content=message)]}
    )
    messages = result["messages"]
    for msg in reversed(messages):
        if hasattr(msg, "content") and msg.content and not getattr(msg, "tool_calls", None):
            return str(msg.content)
    return ""