Spaces:
Sleeping
Sleeping
File size: 1,770 Bytes
159f5a5 c57ea8c 159f5a5 a8e67fc 159f5a5 a8e67fc 159f5a5 68af3c5 159f5a5 c57ea8c 159f5a5 | 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 | from __future__ import annotations
import asyncio
from typing import Annotated, Any, Literal, Optional
from pydantic import BaseModel, Field
class QueryInput(BaseModel):
query: str
team_id: str
session_id: str
# Populated server-side from the auth session — never trusted from the client body.
# Empty list = fall back to team_id-based filtering (legacy / dev mode).
allowed_channel_ids: list[str] = Field(default_factory=list)
class AgentTask(BaseModel):
agent: Literal["doc_search", "ticket_lookup", "confluence_search", "slack_search", "live_docs", "summariser", "sql_query"]
input: str
depends_on: list[str] = Field(default_factory=list)
class ExecutionPlan(BaseModel):
tasks: list[AgentTask]
reasoning: str
class RetrievedChunk(BaseModel):
chunk_id: str
text: str
source: str
source_type: str
score: float
reranker_score: Optional[float] = None
class AgentResult(BaseModel):
agent: str
chunks: list[RetrievedChunk] = Field(default_factory=list)
retrieval_confidence: Literal["high", "medium", "low"] = "low"
error: Optional[str] = None
class KnowledgeGraphState(BaseModel):
model_config = {"arbitrary_types_allowed": True}
query_input: QueryInput
execution_plan: Optional[ExecutionPlan] = None
agent_results: Annotated[dict[str, AgentResult], lambda x, y: {**x, **y}] = Field(default_factory=dict)
final_answer: Optional[str] = None
citations: list[RetrievedChunk] = Field(default_factory=list)
guardrail_passed: Optional[bool] = None
guardrail_score: Optional[float] = None
escalate: bool = False
sse_queue: Any = None # asyncio.Queue for streaming events
class SSEEvent(BaseModel):
event: str
data: Any
|