"""The brain: a streaming, tool-calling chat loop over a swappable LLM vendor. Phase-1 vendor is DeepSeek (OpenAI-compatible chat completions), ported from the pattern in ``voice/deepseek.py`` and extended with SSE streaming + tool calling. The sentence-regrouping approach (start speaking sentence one while the rest still generates) is ported from ``voice/llm.py``. Vendor-swappable: ``Brain`` speaks a neutral event stream and OpenAI-shaped message history. To add Anthropic later, implement ``LLMAdapter.stream`` to translate that history to/from the Anthropic API and yield the same events — nothing else changes. Neutral adapter event contract (async iterator of dicts): {"type": "text", "text": ""} # streamed assistant text {"type": "tool_calls", "calls": [ {id,name,arguments}, ... ]} # 0..1 times {"type": "finish", "reason": "stop"|"tool_calls"|...} # exactly once """ from __future__ import annotations import json import os import re from typing import AsyncIterator, Awaitable, Callable # --- sentence regrouping (ported from voice/llm.py) ------------------------ # Flush a chunk when we hit end-of-sentence punctuation (incl. Arabic) or the # buffer grows long without any. _SENTENCE_END = re.compile(r"[.!?؟…]['\")\]]?\s|\n") _MAX_CHUNK = 220 # chars; hard flush so a run-on sentence can't stall TTS # Strip markdown symbols that sound wrong read aloud (keeps [emotion] tags — # square brackets are intentionally NOT in this class). _MD_STRIP = re.compile(r"[*_`#~|]|^\s*[-•]\s+", re.MULTILINE) _MAX_TOOL_ROUNDS = 5 # hard cap on tool round-trips per user turn # The fast lane emits this sentinel (alone) to escalate a turn to the full # brain — same contract as the desktop's voice/local_brain.py. FAST_HANDOFF = "<>" def clean_for_tts(text: str) -> str: """Remove markdown symbols before a sentence goes to the UI/TTS.""" return _MD_STRIP.sub("", text).strip() class SentenceRegrouper: """Accumulates streamed token deltas and emits complete sentences. ``push(delta)`` returns a list of finished sentences (cleaned, [emotion] tags preserved); ``flush()`` returns whatever remains at end-of-turn. """ def __init__(self) -> None: self._buf = "" def push(self, text: str) -> list[str]: self._buf += text out: list[str] = [] while True: m = _SENTENCE_END.search(self._buf) if m: chunk = self._buf[: m.end()].strip() self._buf = self._buf[m.end():] if chunk: out.append(clean_for_tts(chunk)) elif len(self._buf) > _MAX_CHUNK: cut = self._buf.rfind(" ", 0, _MAX_CHUNK) cut = cut if cut > 40 else _MAX_CHUNK chunk = self._buf[:cut].strip() self._buf = self._buf[cut:] if chunk: out.append(clean_for_tts(chunk)) else: break return [s for s in out if s] def flush(self) -> str: tail = clean_for_tts(self._buf) self._buf = "" return tail # --- adapters -------------------------------------------------------------- class LLMAdapter: """Vendor interface. Implement ``stream`` to yield the neutral events.""" async def stream(self, messages: list[dict], tools: list[dict] | None) -> AsyncIterator[dict]: raise NotImplementedError async def aclose(self) -> None: # pragma: no cover - trivial pass class DeepSeekAdapter(LLMAdapter): """DeepSeek / OpenAI-compatible streaming chat completions.""" def __init__(self, client=None, model: str | None = None) -> None: self.base = os.environ.get( "DEEPSEEK_URL", "https://api.deepseek.com").rstrip("/") self.model = model or os.environ.get("MODEL", "deepseek-v4-pro") self.temperature = float(os.environ.get("BRAIN_TEMPERATURE", "0.6")) self.max_tokens = int(os.environ.get("BRAIN_MAX_TOKENS", "1024")) self._client = client # inject an httpx.AsyncClient in tests self._owns_client = client is None def _ensure_client(self): if self._client is None: import httpx self._client = httpx.AsyncClient(timeout=httpx.Timeout(120.0)) return self._client async def aclose(self) -> None: if self._owns_client and self._client is not None: await self._client.aclose() self._client = None async def stream(self, messages: list[dict], tools: list[dict] | None) -> AsyncIterator[dict]: key = os.environ.get("DEEPSEEK_API_KEY", "") if not key: raise RuntimeError("DEEPSEEK_API_KEY not set") payload: dict = { "model": self.model, "messages": messages, "stream": True, "temperature": self.temperature, "max_tokens": self.max_tokens, } if tools: payload["tools"] = tools payload["tool_choice"] = "auto" client = self._ensure_client() headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"} tool_buf: dict[int, dict] = {} finish = None async with client.stream("POST", self.base + "/chat/completions", json=payload, headers=headers) as resp: resp.raise_for_status() async for line in resp.aiter_lines(): if not line or not line.startswith("data:"): continue data = line[len("data:"):].strip() if data == "[DONE]": break try: obj = json.loads(data) except ValueError: continue choices = obj.get("choices") or [{}] choice = choices[0] delta = choice.get("delta") or {} if delta.get("content"): yield {"type": "text", "text": delta["content"]} for tc in (delta.get("tool_calls") or []): idx = tc.get("index", 0) slot = tool_buf.setdefault( idx, {"id": "", "name": "", "arguments": ""}) if tc.get("id"): slot["id"] = tc["id"] fn = tc.get("function") or {} if fn.get("name"): slot["name"] = fn["name"] if fn.get("arguments"): slot["arguments"] += fn["arguments"] if choice.get("finish_reason"): finish = choice["finish_reason"] if tool_buf: yield {"type": "tool_calls", "calls": [tool_buf[i] for i in sorted(tool_buf)]} yield {"type": "finish", "reason": finish or "stop"} # --- the brain loop -------------------------------------------------------- ToolExecutor = Callable[[str, dict], Awaitable[str]] class Brain: """Runs the streaming tool-calling loop for one user turn. ``messages`` is the full OpenAI-shaped list (system + history + the new user message). It is MUTATED in place: assistant/tool messages produced during the turn are appended, so the caller's history stays current. """ def __init__(self, adapter: LLMAdapter, tool_specs: list[dict] | None, tool_exec: ToolExecutor | None) -> None: self.adapter = adapter self.tool_specs = tool_specs or None self.tool_exec = tool_exec async def run(self, messages: list[dict], on_token: Callable[[], None] | None = None ) -> AsyncIterator[str]: """Yield reply sentences (raw, [emotion] tags intact) as they complete. ``on_token`` (optional) fires once, on the FIRST streamed text token of the turn — the caller uses it for time-to-first-token timing.""" fired = False for _round in range(_MAX_TOOL_ROUNDS): regroup = SentenceRegrouper() text_acc = "" calls: list[dict] = [] async for ev in self.adapter.stream(messages, self.tool_specs): kind = ev.get("type") if kind == "text": if on_token is not None and not fired: fired = True try: on_token() except Exception: # noqa: BLE001 — timing must not break the turn pass text_acc += ev["text"] for sentence in regroup.push(ev["text"]): yield sentence elif kind == "tool_calls": calls = ev.get("calls") or [] # "finish" needs no handling here tail = regroup.flush() if tail: yield tail if not calls: if text_acc.strip(): messages.append({"role": "assistant", "content": text_acc}) return # record the assistant's tool-call turn, then run each tool messages.append({ "role": "assistant", "content": text_acc or None, "tool_calls": [ {"id": c.get("id") or f"call_{i}", "type": "function", "function": {"name": c.get("name", ""), "arguments": c.get("arguments") or "{}"}} for i, c in enumerate(calls) ], }) for i, c in enumerate(calls): try: args = json.loads(c.get("arguments") or "{}") if not isinstance(args, dict): args = {} except ValueError: args = {} name = c.get("name", "") try: result = (await self.tool_exec(name, args) if self.tool_exec else f"tool '{name}' is unavailable") except Exception as e: # noqa: BLE001 — a tool error must not kill the turn result = f"tool '{name}' failed: {str(e)[:200]}" messages.append({"role": "tool", "tool_call_id": c.get("id") or f"call_{i}", "content": result}) # loop again so the model can use the tool results # ran out of tool rounds — stop quietly return # --- the fast lane --------------------------------------------------------- def _fast_history(history: list[dict]) -> list[dict]: """The fast lane's view of the SHARED history: plain user/assistant text turns only. Tool-call scaffolding produced by full-brain turns (assistant messages with ``content=None`` + ``tool_calls``, and ``tool`` results) is dropped — the fast model has no tools and just needs the conversation.""" out: list[dict] = [] for m in history: role = m.get("role") content = m.get("content") if role in ("user", "assistant") and isinstance(content, str) \ and content.strip(): out.append({"role": role, "content": content}) return out class FastLane: """The cheap, tool-less front lane over a fast model (deepseek-v4-flash). Streams spoken sentences from a compact-persona prompt; if the model emits the escalation sentinel ``<>`` (per that prompt) it signals a handoff so the caller silently reruns the turn on the full ``Brain``. Ported from the desktop's ``voice/local_brain.py`` two-lane design. ``run`` yields tuples: ``("say", sentence)`` a spoken sentence (raw, [emotion] tag intact) ``("handoff", None)`` escalate — emitted ALONE and FIRST, then the generator stops (nothing spoken, nothing to commit) ``("final", text)`` LAST on a clean finish — the full assistant text (sentinel stripped) for the caller to commit to the shared history """ def __init__(self, adapter: LLMAdapter, system_prompt: str) -> None: self.adapter = adapter self.system_prompt = system_prompt def build_messages(self, history: list[dict], user_content: str ) -> list[dict]: return ([{"role": "system", "content": self.system_prompt}] + _fast_history(history) + [{"role": "user", "content": user_content}]) async def run(self, history: list[dict], user_content: str, on_token: Callable[[], None] | None = None ) -> AsyncIterator[tuple]: messages = self.build_messages(history, user_content) regroup = SentenceRegrouper() full = "" spoke = False fired = False # hold the inner stream so we can aclose it even on an early handoff # return (otherwise the abandoned generator warns at GC time). stream = self.adapter.stream(messages, None) try: async for ev in stream: if ev.get("type") != "text": continue if on_token is not None and not fired: fired = True try: on_token() except Exception: # noqa: BLE001 pass full += ev["text"] # a sentinel before we've spoken a word -> escalate immediately if not spoke and FAST_HANDOFF in full: yield ("handoff", None) return for sentence in regroup.push(ev["text"]): if FAST_HANDOFF in sentence or not sentence.strip(): continue # never speak a stray sentinel spoke = True yield ("say", sentence) tail = regroup.flush() if not spoke and FAST_HANDOFF in full: yield ("handoff", None) return if tail and FAST_HANDOFF not in tail: spoke = True yield ("say", tail) if not spoke: # empty / sentinel-only reply -> escalate rather than say nothing yield ("handoff", None) return committed = full.replace(FAST_HANDOFF, "").strip() if committed: yield ("final", committed) finally: aclose = getattr(stream, "aclose", None) if aclose is not None: try: await aclose() except Exception: # noqa: BLE001 pass