socaichat/backend/app/graph/builder.py

"""Build and compile the LangGraph agent.

Phase 1 used a simple single-node StateGraph.
Phase 2 upgrades to create_react_agent (ReAct pattern) with dynamic tool binding.

When no tools are requested, we fall back to a plain single-node graph so the
agent does not produce unnecessary tool-call reasoning.
"""

from __future__ import annotations

from langchain_openai import AzureChatOpenAI
from langgraph.graph import StateGraph
from langgraph.prebuilt import create_react_agent

from app.config import settings
from app.graph.nodes import call_model
from app.graph.state import ChatState
from app.store.memory import get_checkpointer

# Model parameter presets
MODEL_PARAMS: dict[str, dict] = {
    "flash": {"max_tokens": 500, "temperature": 0.2},
    "pro": {"max_tokens": 4096, "temperature": 0.3},
}

# System prompt that instructs the ReAct agent
SYSTEM_PROMPT = (
    "You are SOC Assistant, an enterprise AI assistant. "
    "You help users with knowledge base queries, ticket management, "
    "and general questions. "
    "When the user has enabled specific tools, you may use them if relevant. "
    "If you decide not to use an available tool, briefly explain why. "
    "Always respond in the same language the user uses. "
    "Be concise, accurate, and helpful.\n\n"
    "When web search tools are available:\n"
    "- Use search_web first to find relevant pages\n"
    "- Use read_url to get full content from the most relevant URLs (1-3 max)\n"
    "- Use sort_by_relevance to rank results if you have many documents\n"
    "\n## Tool result presentation rules\n"
    "The frontend renders tool results as structured UI cards — users see the full data visually.\n"
    "After calling a tool, your text response MUST follow these rules:\n"
    "- DO NOT repeat or list the raw data from the tool result\n"
    "- Write only: one brief sentence summarizing what was found + any analytical insight\n"
    "- The tool result JSON contains a 'llm_text' field — use that as your starting point\n"
    "Specifically:\n"
    "- kb_search: State whether the knowledge base answered the question. Do NOT re-list document names or content.\n"
    "- ticket_list/ticket_detail: Give a one-line status distribution insight. Do NOT enumerate each ticket.\n"
    "- generate_document: Confirm document type and that it is ready. The download link is shown in UI.\n"
    "- sandbox_run: State whether execution succeeded. If error, explain the root cause briefly.\n"
    "- web_search: State whether useful sources were found. Do NOT re-list URLs or snippets.\n"
)

# Cache compiled graphs to avoid re-creation on every request.
# Key: (model, frozenset(tool_names))
_graph_cache: dict[tuple, object] = {}


def _get_llm(model: str) -> AzureChatOpenAI:
    """Create an AzureChatOpenAI instance with preset parameters."""
    params = MODEL_PARAMS.get(model, MODEL_PARAMS["flash"])
    return AzureChatOpenAI(
        azure_endpoint=settings.azure_openai_endpoint,
        api_key=settings.azure_openai_api_key,
        api_version=settings.azure_openai_api_version,
        azure_deployment=settings.azure_openai_deployment,
        max_tokens=params["max_tokens"],
        temperature=params["temperature"],
        streaming=True,
    )


# MCP tool names that should not be cached (bound to per-request client)
_MCP_TOOL_NAMES = {"search_web", "read_url", "sort_by_relevance"}


async def get_chat_graph(model: str = "flash", tools: list | None = None):
    """Get or create a compiled graph for the given model and tool set.

    When tools are provided, creates a ReAct agent that can call tools.
    When no tools, falls back to a simple single-node graph.

    MCP tools are bound to a per-request client session, so graphs
    containing them are never cached.
    """
    tools = tools or []
    tool_names = frozenset(t.name for t in tools)
    has_mcp_tools = bool(tool_names & _MCP_TOOL_NAMES)

    cache_key = (model, tool_names)
    if not has_mcp_tools and cache_key in _graph_cache:
        return _graph_cache[cache_key]

    checkpointer = await get_checkpointer()
    llm = _get_llm(model)

    if tools:
        # ReAct agent with tool calling
        graph = create_react_agent(
            llm,
            tools=tools,
            checkpointer=checkpointer,
            prompt=SYSTEM_PROMPT,
        )
    else:
        # Simple graph without tools (Phase 1 style)
        builder = StateGraph(ChatState)
        builder.add_node("agent", call_model)
        builder.set_entry_point("agent")
        builder.set_finish_point("agent")
        graph = builder.compile(checkpointer=checkpointer)

    if not has_mcp_tools:
        _graph_cache[cache_key] = graph
    return graph