feat(kb/ask): SSE streaming endpoint with live agent progress

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
This commit is contained in:
2026-04-25 12:41:25 +00:00
parent b127cafb01
commit 76fd77f904
2 changed files with 182 additions and 26 deletions
+105 -25
View File
@@ -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",
},
)
+77 -1
View File
@@ -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",