"""LangChain tool-calling agent + streaming generator. Built on the LangChain 1.x ``create_agent`` API (which returns a compiled LangGraph). We stream with ``stream_mode=["updates", "messages"]`` so we get both per-node updates (tool calls + tool results) AND token-level streaming of the final answer. Public surface: * ``has_api_key()`` — is the OpenAI key configured? * ``run_turn_streaming(conversation, user, user_text)`` — yields dicts shaped like ``{"type": "", "payload": {...}}`` so the SSE view can serialise them as event-stream frames. When ``OPENAI_API_KEY`` is missing the function emits a single canned stub reply, so dev / CI environments work without an API key (mirroring the pattern in ``apps.repository.analytics.ai_insights``). """ from __future__ import annotations import json import logging from typing import Any, Generator from django.conf import settings from .history import build_history_messages from .prompts import build_system_prompt from .tools import build_tools, extract_surfaced_object_ids logger = logging.getLogger(__name__) STUB_REPLY = ( "The assistant is running in offline mode (no OPENAI_API_KEY configured). " "Once an API key is set, I'll answer questions about the platform and your " "data here — using only what your role lets you see." ) _STUB_CHUNK = 60 # chars per stub-mode token chunk RECURSION_LIMIT = 10 # ≈ 5 tool calls (each = agent node + tools node) def has_api_key() -> bool: return bool(getattr(settings, "OPENAI_API_KEY", "").strip()) def _chunk_string(text: str, n: int = _STUB_CHUNK) -> Generator[str, None, None]: if not text: return for i in range(0, len(text), n): yield text[i : i + n] def _safe_parse_observation(content: Any) -> Any: """Tool observations come back as JSON-serialised strings on the ToolMessage.content. Try to parse so we can pull surfaced IDs.""" if isinstance(content, (dict, list)): return content if isinstance(content, str): try: return json.loads(content) except (ValueError, TypeError): return content return content def run_turn_streaming(conversation, user, user_text: str) -> Generator[dict, None, None]: """Run one agent turn and yield streaming events. Events emitted: * ``{"type": "tool_call", "payload": {"name": ..., "arguments": ...}}`` * ``{"type": "tool_result", "payload": {"name": ..., "preview": ..., "surfaced_object_ids": [...]}}`` * ``{"type": "token", "payload": {"text": "..."}}`` * ``{"type": "final", "payload": {"output": "..."}}`` (full final text — last) * ``{"type": "error", "payload": {"message": "..."}}`` (on failure) The view persists messages and writes the AuditLog after the generator is exhausted, using the collected metadata. """ if not has_api_key(): for chunk in _chunk_string(STUB_REPLY): yield {"type": "token", "payload": {"text": chunk}} yield {"type": "final", "payload": {"output": STUB_REPLY, "stub": True}} return try: from langchain.agents import create_agent from langchain_core.messages import HumanMessage from langchain_openai import ChatOpenAI except Exception as exc: # noqa: BLE001 logger.exception("LangChain imports failed in agent") msg = f"Assistant import failed: {exc}" for chunk in _chunk_string(msg): yield {"type": "token", "payload": {"text": chunk}} yield {"type": "final", "payload": {"output": msg, "error": True}} return tools = build_tools(user) system_prompt = build_system_prompt(user, conversation, available_tools=tools) history = build_history_messages(conversation) llm = ChatOpenAI( model=settings.LANGCHAIN_CHAT_MODEL_FAST, temperature=0, api_key=settings.OPENAI_API_KEY, ) agent = create_agent(model=llm, tools=tools, system_prompt=system_prompt) # Agent input: full message list (history + the new user turn). input_messages = [*history, HumanMessage(content=user_text)] # Track tool-call args by ID so we can correlate them with tool results. pending_args: dict[str, dict] = {} tool_call_records: list[dict] = [] final_text_parts: list[str] = [] try: for stream_mode, chunk in agent.stream( {"messages": input_messages}, stream_mode=["updates", "messages"], config={"recursion_limit": RECURSION_LIMIT}, ): if stream_mode == "messages": # Token-level stream: (AIMessageChunk, metadata). msg_chunk, metadata = chunk # Skip tool node tokens — only the agent (model) node tokens # represent the final answer text. node = (metadata or {}).get("langgraph_node", "") if node and node not in {"agent", "model", "call_model"}: continue content = getattr(msg_chunk, "content", "") if isinstance(content, str) and content: final_text_parts.append(content) yield {"type": "token", "payload": {"text": content}} elif isinstance(content, list): # Some providers stream content as blocks; collect text blocks. for block in content: if isinstance(block, dict) and block.get("type") == "text": text = block.get("text", "") if text: final_text_parts.append(text) yield {"type": "token", "payload": {"text": text}} elif stream_mode == "updates": # Per-node update: {: {"messages": [...]}}. for node_name, update in (chunk or {}).items(): msgs = (update or {}).get("messages", []) if isinstance(update, dict) else [] # Agent node emits AIMessages (possibly with tool_calls). for m in msgs: for tc in getattr(m, "tool_calls", []) or []: tc_id = tc.get("id", "") tc_name = tc.get("name", "") tc_args = tc.get("args", {}) or {} pending_args[tc_id] = tc_args yield { "type": "tool_call", "payload": {"name": tc_name, "arguments": tc_args}, } # Tools node emits ToolMessages with the tool result content. if node_name in {"tools", "tool_node"}: for tm in msgs: tool_name = getattr(tm, "name", "") or "" content = getattr(tm, "content", "") obs = _safe_parse_observation(content) tc_id = getattr(tm, "tool_call_id", "") or "" args = pending_args.get(tc_id, {}) record = { "name": tool_name, "preview": (str(content) if content is not None else "")[:1000], "arguments": args, "surfaced_object_ids": extract_surfaced_object_ids(tool_name, obs), } tool_call_records.append(record) yield { "type": "tool_result", "payload": { "name": record["name"], "preview": record["preview"], "arguments": record["arguments"], "surfaced_object_ids": record["surfaced_object_ids"], }, } except Exception as exc: # noqa: BLE001 logger.exception("Assistant agent run failed") msg = f"Sorry — I hit an error: {exc.__class__.__name__}." for piece in _chunk_string(msg): yield {"type": "token", "payload": {"text": piece}} yield {"type": "final", "payload": {"output": msg, "error": True}} return # Note: tool_call_records is captured here for the view layer but tool # call metadata is also yielded as it streams; the view re-collects from # the SSE event records, so we don't need to re-yield it. final_text = "".join(final_text_parts) yield {"type": "final", "payload": {"output": final_text}}