From 76fd77f904cd7b7a69471d95a4a5dc1aa69aa125 Mon Sep 17 00:00:00 2001 From: Chaim Date: Sat, 25 Apr 2026 12:41:25 +0000 Subject: [PATCH] feat(kb/ask): SSE streaming endpoint with live agent progress MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds POST /kb/ask/stream that returns text/event-stream. The agent runs in a fire-and-forget task; an on_event callback wired into AgentRunner.run pushes {type, label, ...} events into an asyncio.Queue, and the response generator consumes them as SSE data: lines. Event types and Hebrew labels: - thinking — before each LLM call ("קוראת את השאלה" / "ממשיכה לחקור") - tool_start — before each tool. For search_insurance_kb the label includes the actual query string Shira chose, so the user can see WHAT she's searching for ("מחפשת בבסיס הידע: '...'"). Other tools fall back to a generic Hebrew label keyed off the function name. - tool_done / tool_error — after each tool. For search_insurance_kb the label echoes the first non-empty line of the tool result (which the legal_kb tool formats as "## נמצאו N קטעים"). - writing — when the model returns text without tool calls. - answer — final event with {text, sources, conversationId}, mirrors the non-streaming /kb/ask response shape. - error — runner exception; client should close the EventSource. Heartbeats: when the queue is idle for >15s the generator yields a ": keep-alive" comment line (ignored by EventSource) so any proxy in the path doesn't drop the connection during long agent thinks. The original POST /kb/ask is unchanged — the streaming endpoint is additive. Both share _build_ask_runner_context + _build_ask_messages helpers extracted from the original handler. Refs Task Master #1 --- api/routes/kb_public.py | 130 ++++++++++++++++++++++++++++------- api/services/agent_runner.py | 78 ++++++++++++++++++++- 2 files changed, 182 insertions(+), 26 deletions(-) diff --git a/api/routes/kb_public.py b/api/routes/kb_public.py index c8b5df2..f1cc1c3 100644 --- a/api/routes/kb_public.py +++ b/api/routes/kb_public.py @@ -233,16 +233,8 @@ class AskRequest(BaseModel): conversation_history: list[dict] | None = None -@router.post("/ask") -async def ask(body: AskRequest, request: Request): - """Full agent loop: the user's question goes to shira, she calls - search_insurance_kb as needed and returns a summarized answer. - - Thin wrapper around the existing smart-assistant chat endpoint so the - KnowledgeBase UI can expose an ask mode without rebuilding prompt - plumbing. - """ - _verify_auth(request) +def _build_ask_runner_context() -> tuple: + """Common construction shared by /kb/ask and /kb/ask/stream.""" from api.services.agent_runner import AgentRunner from api.services.prompt_builder import build_office_prompt from api.services.skills import list_skills @@ -254,15 +246,12 @@ async def ask(body: AskRequest, request: Request): max_tokens=int(os.environ.get("MAX_OUTPUT_TOKENS", "4096")), max_iterations=int(os.environ.get("MAX_AGENT_ITERATIONS", "10")), ) - context = {"currentUser": {"id": "kb-ui", "name": "KB UI"}} system_prompt = build_office_prompt( context, user_profile="", available_skills=list_skills(), ) - # Pin the question to the insurance KB context so generic-law answers - # (e.g. "תקנה 37 לתקנות סדר הדין האזרחי") don't hijack the response. system_prompt += ( "\n\nContext: The user is asking from the Knowledge Base UI of the " "Israeli National Insurance firm. Every question should be answered " @@ -272,7 +261,10 @@ async def ask(body: AskRequest, request: Request): "Never answer from generic legal training if the KB has a matching " "section — if no KB section matches, say so explicitly." ) + return runner, context, system_prompt + +def _build_ask_messages(body: AskRequest) -> list[dict]: messages: list[dict] = [] for m in (body.conversation_history or []): role = m.get("role", "user") @@ -281,12 +273,29 @@ async def ask(body: AskRequest, request: Request): "content": m.get("content", ""), }) messages.append({"role": "user", "content": body.message}) + return messages - # Collect the sources Shira actually searched through during this turn. - # The legal_kb tool appends to this list on every search_insurance_kb - # call; the UI uses it to render a side PDF viewer next to the answer. + +def _filter_pdf_sources(sources_used: list[dict]) -> list[dict]: + return [ + s for s in sources_used + if (s.get("original_path") or "").lower().endswith(".pdf") + ] + + +@router.post("/ask") +async def ask(body: AskRequest, request: Request): + """Full agent loop: the user's question goes to shira, she calls + search_insurance_kb as needed and returns a summarized answer. + + Thin wrapper around the existing smart-assistant chat endpoint so the + KnowledgeBase UI can expose an ask mode without rebuilding prompt + plumbing. + """ + _verify_auth(request) + runner, context, system_prompt = _build_ask_runner_context() + messages = _build_ask_messages(body) sources_used: list[dict] = [] - try: text = await runner.run( system_prompt=system_prompt, @@ -300,15 +309,86 @@ async def ask(body: AskRequest, request: Request): except Exception as e: logger.exception("[kb.ask] error") raise HTTPException(status_code=500, detail=str(e)) - - # Keep only the ones that actually have a PDF attached. - pdf_sources = [ - s for s in sources_used - if (s.get("original_path") or "").lower().endswith(".pdf") - ] - return { "text": text, "conversationId": body.conversation_id, - "sources": pdf_sources, + "sources": _filter_pdf_sources(sources_used), } + + +@router.post("/ask/stream") +async def ask_stream(body: AskRequest, request: Request): + """Streaming variant of /kb/ask. Returns text/event-stream with one + SSE event per agent step (thinking, tool_start, tool_done) and a + final {type:"answer",text,sources} event. The KnowledgeBase UI + consumes this via EventSource (proxied through an EspoCRM EntryPoint + so the auth token stays on the server). + """ + import asyncio + import json as _json + from fastapi.responses import StreamingResponse + + _verify_auth(request) + runner, context, system_prompt = _build_ask_runner_context() + messages = _build_ask_messages(body) + sources_used: list[dict] = [] + + queue: asyncio.Queue = asyncio.Queue() + + def _on_event(ev: dict) -> None: + # AgentRunner calls this synchronously between awaits; an unbounded + # queue lets us put_nowait without blocking the loop. + try: + queue.put_nowait(ev) + except Exception as e: + logger.warning("[kb.ask/stream] queue put failed: %s", e) + + async def _runner_task(): + try: + text = await runner.run( + system_prompt=system_prompt, + messages=messages, + context=context, + espocrm_url=os.environ.get("ESPOCRM_URL", ""), + espocrm_api_key=os.environ.get("ESPOCRM_API_KEY", ""), + allowed_toolsets=["legal"], + kb_sources_used=sources_used, + on_event=_on_event, + ) + await queue.put({ + "type": "answer", + "text": text, + "sources": _filter_pdf_sources(sources_used), + "conversationId": body.conversation_id, + }) + except Exception as e: + logger.exception("[kb.ask/stream] runner error") + await queue.put({"type": "error", "message": str(e)}) + finally: + await queue.put(None) # sentinel + + asyncio.create_task(_runner_task()) + + async def _gen(): + # Heartbeat every 15s keeps the connection open through any + # timeout-happy proxies; comment lines (`: ...`) are ignored by + # EventSource but keep TCP/HTTP alive. + while True: + try: + ev = await asyncio.wait_for(queue.get(), timeout=15.0) + except asyncio.TimeoutError: + yield ": keep-alive\n\n" + continue + if ev is None: + break + yield "data: " + _json.dumps(ev, ensure_ascii=False) + "\n\n" + + return StreamingResponse( + _gen(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache, no-transform", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) diff --git a/api/services/agent_runner.py b/api/services/agent_runner.py index 96ea90d..bff9db0 100644 --- a/api/services/agent_runner.py +++ b/api/services/agent_runner.py @@ -4,7 +4,7 @@ from __future__ import annotations import json import logging -from typing import Any +from typing import Any, Callable from openai import AsyncOpenAI @@ -21,6 +21,49 @@ from api.services.delegate import register_delegate_tool logger = logging.getLogger("shira.agent") +# Hebrew labels for the on_event callback in run(). These are surfaced +# verbatim in the KnowledgeBase ask UI to give users live progress +# (instead of just an elapsed-time counter). Keep them short — they +# render in a single line of the progress log. + +def _thinking_label(iteration: int) -> str: + if iteration == 0: + return "קוראת את השאלה" + return "ממשיכה לחקור" + + +def _tool_label(name: str, args: dict) -> str: + if name == "search_insurance_kb": + q = (args.get("query") or "").strip() + if q: + qs = q if len(q) <= 60 else q[:60] + "…" + return f'מחפשת בבסיס הידע: "{qs}"' + return "מחפשת בבסיס הידע" + if name == "list_skills" or name == "view_skill": + return "בודקת מיומנויות זמינות" + if name == "save_memory" or name == "view_memory": + return "מסתכלת בזיכרון" + if name.startswith("delegate"): + return "מאצילה משימת משנה לסוכן עזר" + if name.startswith("create_") or name.startswith("update_") or name.startswith("save_"): + return f"כותבת ב-CRM: {name}" + return f"מפעילה כלי: {name}" + + +def _tool_done_label(name: str, result) -> str: + if name == "search_insurance_kb": + # legal_kb_tools formats the result as a Hebrew string starting + # with a "## נמצאו N קטעים" line. Surface that count. + s = str(result or "") + for line in s.splitlines(): + line = line.strip().lstrip("#").strip() + if not line: + continue + return line[:80] if len(line) > 80 else line + return "סיימה לחפש" + return "הושלם" + + class AgentRunner: """ Runs an agentic conversation loop using the OpenAI SDK pointed at ai-gateway. @@ -59,6 +102,7 @@ class AgentRunner: blocked_tools: set[str] | None = None, allowed_toolsets: list[str] | None = None, kb_sources_used: list[dict] | None = None, + on_event: Callable[[dict], None] | None = None, ) -> str: """Run a conversation with tool-calling loop. Returns the final text response. @@ -66,7 +110,20 @@ class AgentRunner: depth: Current delegation depth. 0 = main conversation, 1 = child agent. blocked_tools: Tool names to exclude (used by delegation to prevent recursion). allowed_toolsets: If set, only register tools from these categories (crm, documents, legal). + on_event: Optional sync callback fired between iterations and tool calls. + Receives a dict {type, label, ...} where type ∈ {thinking, + tool_start, tool_done, tool_error}. Used by /kb/ask/stream + to surface live progress to the UI. Exceptions in the + callback are logged and ignored — the agent never fails + because the caller's display logic threw. """ + def emit(event: dict) -> None: + if on_event is None: + return + try: + on_event(event) + except Exception as e: + logger.warning("[agent] on_event callback failed: %s", e) # Build tools from context crm = EspoCrmClient(espocrm_url, espocrm_api_key) @@ -119,6 +176,12 @@ class AgentRunner: for iteration in range(self.max_iterations): logger.info("[agent] depth=%d iteration %d, messages=%d", depth, iteration + 1, len(openai_messages)) + emit({ + "type": "thinking", + "iteration": iteration + 1, + "label": _thinking_label(iteration), + }) + try: response = await self.client.chat.completions.create( model=self.model, @@ -135,6 +198,7 @@ class AgentRunner: # If no tool calls, return the text if not message.tool_calls: + emit({"type": "writing", "label": "מנסחת תשובה"}) return message.content or "לא הצלחתי להבין את הבקשה. נסי שוב." # Process tool calls @@ -148,17 +212,29 @@ class AgentRunner: fn_args = {} logger.info("[agent] tool call: %s(%s)", fn_name, json.dumps(fn_args, ensure_ascii=False)[:100]) + emit({ + "type": "tool_start", + "name": fn_name, + "label": _tool_label(fn_name, fn_args), + }) tool_def = tools_registry.get(fn_name) if not tool_def: result = f"❌ כלי לא ידוע: {fn_name}" + emit({"type": "tool_error", "name": fn_name, "label": f"כלי לא מוכר: {fn_name}"}) else: try: handler = tool_def["handler"] result = await handler(**fn_args) + emit({ + "type": "tool_done", + "name": fn_name, + "label": _tool_done_label(fn_name, result), + }) except Exception as e: logger.error("[agent] tool error %s: %s", fn_name, e) result = f"❌ שגיאה בהפעלת {fn_name}: {e}" + emit({"type": "tool_error", "name": fn_name, "label": f"שגיאה ב-{fn_name}"}) openai_messages.append({ "role": "tool",