Spaces:
Sleeping
Sleeping
File size: 3,549 Bytes
76022ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | import time
import uuid
import asyncio
from typing import Any, Dict, List, Literal, Optional, Set
from pydantic import BaseModel, Field
# -----------------------------------------------------------------------------
# 1. THE CONTRACT: UNIVERSAL DEBUG EVENT SCHEMA
# -----------------------------------------------------------------------------
DebugEventType = Literal[
"run_started",
"input_received",
"prompt_built",
"retrieval_started",
"retrieval_result",
"tool_called",
"tool_result",
"llm_request",
"llm_response",
"policy_check",
"error",
"run_finished"
]
class DebugEvent(BaseModel):
runId: str
sessionId: Optional[str] = None
userId: Optional[str] = None
timestamp: int = Field(default_factory=lambda: int(time.time() * 1000))
type: DebugEventType
title: str
payload: Dict[str, Any] = Field(default_factory=dict)
# -----------------------------------------------------------------------------
# 2. THE BUS: PORTABLE TRANSPORT & STORAGE
# -----------------------------------------------------------------------------
class DebugBus:
def __init__(self):
self.listeners: Set[asyncio.Queue] = set()
self.session_subscriptions: Dict[str, Set[asyncio.Queue]] = {}
self.runs: Dict[str, List[DebugEvent]] = {} # Memory storage for replay
def subscribe(self, session_id: str, queue: asyncio.Queue):
if session_id not in self.session_subscriptions:
self.session_subscriptions[session_id] = set()
self.session_subscriptions[session_id].add(queue)
self.listeners.add(queue)
def unsubscribe(self, session_id: str, queue: asyncio.Queue):
if session_id in self.session_subscriptions:
self.session_subscriptions[session_id].remove(queue)
if not self.session_subscriptions[session_id]:
del self.session_subscriptions[session_id]
self.listeners.discard(queue)
async def emit(self, event: DebugEvent):
# 1. Store locally for replay
if event.runId not in self.runs:
self.runs[event.runId] = []
self.runs[event.runId].append(event)
# 2. Push to active subscribers
# Broadcasters listen to specific session_id streams
session_id = event.sessionId
data = event.model_dump()
if session_id and session_id in self.session_subscriptions:
for q in self.session_subscriptions[session_id]:
await q.put(data)
# Global listeners (if any)
# for q in self.listeners:
# await q.put(data)
def get_run(self, run_id: str) -> List[DebugEvent]:
return self.runs.get(run_id, [])
def get_all_runs(self) -> List[Dict[str, Any]]:
return [{"runId": rid, "events": [e.model_dump() for e in evs]} for rid, evs in self.runs.items()]
# Singleton instance
debug_bus = DebugBus()
# -----------------------------------------------------------------------------
# 3. THE EMITTER: CONVENIENCE WRAPPER
# -----------------------------------------------------------------------------
async def emit_debug_event(
runId: str,
title: str,
type: DebugEventType,
sessionId: Optional[str] = None,
userId: Optional[str] = None,
payload: Optional[Dict[str, Any]] = None
):
event = DebugEvent(
runId=runId,
sessionId=sessionId,
userId=userId,
type=type,
title=title,
payload=payload or {}
)
await debug_bus.emit(event)
|