spb/uqo-chat

1"""Internal LLM types: stream events, tool calls, tool results."""23from __future__ import annotations45import json6import re7from dataclasses import dataclass, field8from typing import Any, Literal910StreamEventType = Literal[11    "text_delta", "tool_call_start", "tool_call_delta", "tool_call_end",12    "usage", "done", "error", "reasoning_delta",13]141516@dataclass17class StreamEvent:18    type: StreamEventType19    data: dict[str, Any] = field(default_factory=dict)202122def repair_json(text: str) -> dict[str, Any] | None:23    """Best-effort repair of a truncated/malformed JSON object (model output).2425    Strategy: strip code fences, then walk backwards over plausible cut points and close the26    open brackets. Returns None when nothing parses.27    """28    s = text.strip()29    s = re.sub(r"^```(?:json)?\s*|\s*```$", "", s).strip()30    if not s:31        return None32    try:33        v = json.loads(s)34        return v if isinstance(v, dict) else None35    except json.JSONDecodeError:36        pass37    start = s.find("{")38    if start < 0:39        return None40    s = s[start:]4142    def closers_at(cut: int) -> str | None:43        stack: list[str] = []44        in_str = False45        esc = False46        for ch in s[:cut]:47            if in_str:48                if esc:49                    esc = False50                elif ch == "\\":51                    esc = True52                elif ch == '"':53                    in_str = False54                continue55            if ch == '"':56                in_str = True57            elif ch in "{[":58                stack.append("}" if ch == "{" else "]")59            elif ch in "}]":60                if stack:61                    stack.pop()62        if in_str:63            return None64        return "".join(reversed(stack))6566    n = len(s)67    for cut in range(n, max(0, n - 4000), -1):68        tail = s[cut - 1] if cut > 0 else ""69        if tail not in '"}]0123456789elsu.':  # plausible end of a value70            continue71        closers = closers_at(cut)72        if closers is None:73            continue74        candidate = re.sub(r",\s*$", "", s[:cut]) + closers75        try:76            v = json.loads(candidate)77            return v if isinstance(v, dict) else None78        except json.JSONDecodeError:79            continue80    return None818283@dataclass84class ToolCallReq:85    id: str86    name: str87    arguments_json: str = ""88    index: int = 089    truncated: bool = False  # the model hit its output limit while emitting this call9091    def arguments(self) -> dict[str, Any]:92        if not self.arguments_json.strip():93            return {}94        try:95            v = json.loads(self.arguments_json)96            return v if isinstance(v, dict) else {}97        except json.JSONDecodeError:98            repaired = repair_json(self.arguments_json)99            if repaired is not None:100                repaired["__repaired__"] = True101                return repaired102            return {"__invalid_json__": self.arguments_json[:2000]}103104105@dataclass106class Artifact:107    type: str  # xlsx | image | csv | file | quiz | sources | search108    file_id: str | None = None109    filename: str | None = None110    preview: Any = None111    url: str | None = None112113    def to_dict(self) -> dict[str, Any]:114        return {115            "type": self.type,116            "file_id": self.file_id,117            "filename": self.filename,118            "preview": self.preview,119            "url": self.url,120        }121122123@dataclass124class ToolResult:125    content: str  # for the model126    artifacts: list[Artifact] = field(default_factory=list)127    meta: dict[str, Any] = field(default_factory=dict)128    payload: dict[str, Any] = field(default_factory=dict)  # for the UI card129    error: bool = False130131    def summary(self, limit: int = 240) -> str:132        s = self.meta.get("summary") or self.content133        return s if len(s) <= limit else s[: limit - 1] + "…"134135136@dataclass137class Usage:138    input_tokens: int = 0139    output_tokens: int = 0140    cost_usd: float = 0.0141    model: str = ""142    provider: str = ""143    generation_id: str = ""144