fix
This commit is contained in:
@@ -25,17 +25,25 @@ class SseEmitter:
|
||||
In a producer task:
|
||||
await emitter.emit("tool_call", {...})
|
||||
await emitter.done({"result": "ok"})
|
||||
|
||||
Production mode: when `debug=False`, raw `tool_call` and `llm_call_*`
|
||||
events are filtered or transformed into friendly `status` events.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, debug: bool = True) -> None:
|
||||
self._queue: asyncio.Queue[tuple[str, str, str] | None] = asyncio.Queue()
|
||||
# (event_type, data_json, event_id)
|
||||
self._event_counter = 0
|
||||
self._closed = False
|
||||
self._debug = debug
|
||||
|
||||
async def emit(self, event_type: str, data: Any) -> None:
|
||||
if self._closed:
|
||||
return
|
||||
# In production mode, transform/fiter debug-only events
|
||||
if not self._debug:
|
||||
event_type, data = self._transform_for_prod(event_type, data)
|
||||
if event_type is None:
|
||||
return # event filtered out
|
||||
self._event_counter += 1
|
||||
event_id = f"evt_{self._event_counter}"
|
||||
try:
|
||||
@@ -44,6 +52,76 @@ class SseEmitter:
|
||||
data_str = json.dumps({"error": "serialization_failed"})
|
||||
await self._queue.put((event_type, data_str, event_id))
|
||||
|
||||
def _transform_for_prod(self, event_type: str, data: Any) -> tuple[str | None, Any]:
|
||||
"""Transform debug events into user-friendly status events for production."""
|
||||
if event_type == "tool_call":
|
||||
# Transform tool_call into a friendly status message
|
||||
tool = data.get("tool", "") if isinstance(data, dict) else ""
|
||||
result = data.get("result", {}) if isinstance(data, dict) else {}
|
||||
is_success = data.get("is_success", True) if isinstance(data, dict) else True
|
||||
# Friendly message based on tool type
|
||||
friendly = self._friendly_tool_message(tool, result, is_success)
|
||||
if friendly:
|
||||
return ("status", {"message": friendly, "type": "tool"})
|
||||
return (None, None) # filter out
|
||||
elif event_type in ("llm_call_start", "llm_call_end"):
|
||||
# Filter out raw LLM call events in production
|
||||
return (None, None)
|
||||
elif event_type == "phase_start":
|
||||
# Keep but with friendly name
|
||||
phase = data.get("phase") if isinstance(data, dict) else None
|
||||
friendly_names = {
|
||||
1: "planning",
|
||||
2: "writing",
|
||||
3: "sending",
|
||||
}
|
||||
name = friendly_names.get(phase, data.get("name", ""))
|
||||
return ("phase_start", {"phase": phase, "name": name, "status": name})
|
||||
elif event_type == "phase_end":
|
||||
return (event_type, data)
|
||||
elif event_type == "warning":
|
||||
return (event_type, data)
|
||||
else:
|
||||
return (event_type, data)
|
||||
|
||||
def _friendly_tool_message(self, tool: str, result: dict, is_success: bool) -> str | None:
|
||||
"""Generate a user-friendly message for a tool call."""
|
||||
if not is_success:
|
||||
return None # hide failed tool calls in production
|
||||
data = result.get("data", {}) if isinstance(result, dict) else {}
|
||||
msg = result.get("message", "") if isinstance(result, dict) else ""
|
||||
if tool == "entity_create":
|
||||
name = data.get("entity_id", "")
|
||||
return f"Added new entity" + (f": {name}" if name else "")
|
||||
elif tool == "entity_update":
|
||||
return "Updated entity"
|
||||
elif tool == "entity_delete":
|
||||
return "Removed entity"
|
||||
elif tool == "env_update":
|
||||
return "Updated game state"
|
||||
elif tool == "update_plot_rails":
|
||||
return "Updated story progress"
|
||||
elif tool == "advance_time":
|
||||
new_time = data.get("new_time", "")
|
||||
return f"Time advanced" + (f" to {new_time}" if new_time else "")
|
||||
elif tool == "schedule_trigger":
|
||||
return "Scheduled future event"
|
||||
elif tool == "rag_query":
|
||||
return None # hide RAG queries in production
|
||||
elif tool == "rag_add":
|
||||
return "Recorded a new fact"
|
||||
elif tool == "calc":
|
||||
return None # hide calculations
|
||||
elif tool == "random_choice":
|
||||
return None
|
||||
elif tool == "submit_plan":
|
||||
return None
|
||||
elif tool == "submit_step":
|
||||
return None
|
||||
elif tool == "suggest_actions":
|
||||
return None
|
||||
return msg if msg else None
|
||||
|
||||
async def ping(self) -> None:
|
||||
await self.emit("ping", {"ts": _now_iso()})
|
||||
|
||||
|
||||
Reference in New Issue
Block a user