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:
+105
-25
@@ -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",
|
||||
},
|
||||
)
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user