"""Multi-backend chat router.""" import json import os import subprocess import sys import tempfile import time import uuid import urllib.error import urllib.request from tai_pkg.claude import Claude from tai_pkg.codex import Codex from tai_pkg.config import Config from tai_pkg.consts import Consts from tai_pkg.errors import BackendError from tai_pkg.gemini import Gemini from tai_pkg.grok import Grok from tai_pkg.models import ModelHealth from tai_pkg.models import ModelScoring from tai_pkg.muse import Muse from tai_pkg.opencode import Opencode from tai_pkg.render import Ansi from tai_pkg.render import Render from tai_pkg.sandbox import Sandbox from tai_pkg.text import Text __all__ = ["ChatClient"] class ChatClient: def __init__(self, config): self.config = config self.scope = ("default", "main") self.fixed_key = None self.health = ModelHealth(getattr(config, "home", "") or "") def muse_key(self): if self.fixed_key is not None: return self.fixed_key return "%s/%s" % (self.scope[0], self.scope[1]) def reset_muse(self): home = getattr(self.config, "home", None) if home is None: return key = self.muse_key() def drop(data): data.pop(key, None) Muse.muse_transact(home, drop) def openrouter_models(self): models = [item for item in getattr(self.config, "openrouter_models", []) if item] explicit = getattr(self.config, "explicit_model", "") if explicit and explicit not in models: models.insert(0, explicit) free = self.health.cached_list("openrouter:free") if free is None: fetched = ModelScoring.fetch_model_ids(Consts.OPENROUTER_BASE + "/models", self.config.openrouter_key) free = [item for item in fetched if item.endswith(":free")] if fetched else [] if free: self.health.store_list("openrouter:free", free) free = ModelScoring.prefer_muse(free or []) models = [item for item in free if item not in models] + models if not models: models = [Consts.OPENROUTER_MODEL] return models def opencode_models(self): models = list(getattr(self.config, "opencode_models", []) or []) if not models: models = self.health.cached_list("opencode:free") if models is None: models = Opencode.opencode_free_models() or [item for item in Consts.ZEN_FREE_MODELS if Text.is_free_model(item)] self.health.store_list("opencode:free", models) return ModelScoring.prefer_muse(models) def ollama_models(self): if getattr(self.config, "ollama_enabled", True) is False: return [] models = list(getattr(self.config, "ollama_models", []) or []) if not models: models = self.health.cached_list("ollama:local", Consts.OLLAMA_LIST_TTL) if models is None: models = ModelScoring.fetch_model_ids(Consts.OLLAMA_BASE + "/models", "", timeout=3) or [] if models: self.health.store_list("ollama:local", models) return models def devplace_models(self): models = list(getattr(self.config, "devplace_models", []) or []) if models: return ModelScoring.prefer_muse(models) free = self.health.cached_list("devplace:free") if free is None: fetched = ModelScoring.fetch_model_ids(self.config.devplace_base + "/models", self.config.devplace_key) if fetched is None: return list(Consts.DEVPLACE_FREE_ROUTES) free = [item for item in fetched if item in Consts.DEVPLACE_FREE_ROUTES or (Consts.MUSE_LABEL in item.lower() and Text.is_free_model(item))] if free: self.health.store_list("devplace:free", free) return ModelScoring.prefer_muse(free or []) def candidates(self): found = [] if getattr(self.config, "opencode_enabled", True) and Opencode.opencode_binary(): for model in self.opencode_models(): found.append({"provider": Consts.OPENCODE_LABEL, "chat_url": Consts.OPENCODE_CHAT_URL, "models_url": Consts.OPENCODE_MODELS_URL, "key": "", "model": model, "profile": ""}) for model in self.ollama_models(): found.append({"provider": Consts.OLLAMA_LABEL, "chat_url": Consts.OLLAMA_BASE + "/chat/completions", "models_url": Consts.OLLAMA_BASE + "/models", "key": "", "model": model, "profile": ""}) if getattr(self.config, "devplace_key", ""): for model in self.devplace_models(): found.append({"provider": Consts.DEVPLACE_LABEL, "chat_url": self.config.devplace_base + "/chat/completions", "models_url": self.config.devplace_base + "/models", "key": self.config.devplace_key, "model": model, "profile": ""}) for entry in self.config.backends: for model in Config.split_models(entry.get("model")) or [self.config.model]: found.append({"provider": entry["label"], "chat_url": entry["base"] + "/chat/completions", "models_url": entry["base"] + "/models", "key": entry["key"], "model": model, "profile": entry.get("profile") or ""}) if getattr(self.config, "openrouter_enabled", True): for model in self.openrouter_models(): found.append({"provider": Consts.OPENROUTER_LABEL, "chat_url": Consts.OPENROUTER_BASE + "/chat/completions", "models_url": Consts.OPENROUTER_BASE + "/models", "key": self.config.openrouter_key, "model": model, "profile": ""}) if getattr(self.config, "pollinations_enabled", True): for model in getattr(self.config, "pollinations_models", []) or list(Consts.POLLINATIONS_MODELS): found.append({"provider": Consts.POLLINATIONS_LABEL, "chat_url": Consts.POLLINATIONS_BASE + "/chat/completions", "models_url": Consts.POLLINATIONS_BASE + "/models", "key": "", "model": model, "profile": ""}) if getattr(self.config, "zen_enabled", False): for model in getattr(self.config, "zen_models", []) or list(Consts.ZEN_FREE_MODELS): found.append({"provider": Consts.ZEN_LABEL, "chat_url": Consts.ZEN_BASE + "/chat/completions", "models_url": Consts.ZEN_BASE + "/models", "key": "", "model": model, "profile": "opencode"}) found.append({"provider": Consts.MUSE_LABEL, "chat_url": Consts.MUSE_CHAT_URL, "models_url": Consts.MUSE_MODELS_URL, "key": "", "model": self.config.model, "profile": ""}) if Grok.grok_binary(): found.append({"provider": Consts.GROK_LABEL, "chat_url": Consts.GROK_CHAT_URL, "models_url": Consts.GROK_MODELS_URL, "key": "", "model": self.config.model, "profile": ""}) if Claude.claude_binary(): found.append({"provider": Consts.CLAUDE_LABEL, "chat_url": Consts.CLAUDE_CHAT_URL, "models_url": Consts.CLAUDE_MODELS_URL, "key": "", "model": self.config.model, "profile": ""}) if Codex.codex_binary(): found.append({"provider": Consts.CODEX_LABEL, "chat_url": Consts.CODEX_CHAT_URL, "models_url": Consts.CODEX_MODELS_URL, "key": "", "model": self.config.model, "profile": ""}) if Gemini.gemini_binary(): found.append({"provider": Consts.GEMINI_LABEL, "chat_url": Consts.GEMINI_CHAT_URL, "models_url": Consts.GEMINI_MODELS_URL, "key": "", "model": self.config.model, "profile": ""}) return found def ordered_candidates(self): http = [cand for cand in self.candidates() if cand["chat_url"] not in Consts.CLI_CHAT_URLS] cli = [cand for cand in self.candidates() if cand["chat_url"] in Consts.CLI_CHAT_URLS] ranked = sorted(enumerate(http), key=lambda pair: (ModelScoring.openrouter_free_tier(pair[1]), -self.health.weight(pair[1]["provider"], pair[1]["model"]), pair[0])) usable = [cand for _, cand in ranked if not self.health.circuit_open(cand["provider"], cand["model"])] return (usable or [cand for _, cand in ranked]) + cli def backends(self): rows = [] seen = set() for cand in self.candidates(): key = (cand["provider"], cand["chat_url"]) if key in seen: continue seen.add(key) rows.append((cand["provider"], cand["chat_url"], cand["models_url"], cand["key"], cand["model"])) return rows def headers(self, key, extra=None): headers = {"Content-Type": "application/json"} if key: headers["Authorization"] = "Bearer " + key if extra: headers.update(extra) return headers def complete(self, messages, tools=None, stream_sink=None, ephemeral=False, deadline=None): payload = {"model": self.config.model, "messages": messages} if tools: payload["tools"] = tools payload["tool_choice"] = "auto" limit = getattr(self.config, "retry_seconds", 0) give_up = time.time() + limit if limit > 0 else None if deadline is not None: give_up = deadline if give_up is None else min(give_up, deadline) round_no = 0 while True: errors = [] cands = self.ordered_candidates() if round_no == 0 else self.all_candidates() for cand in cands: reply = self.try_candidate(cand, messages, tools, payload, stream_sink, ephemeral, errors) if reply is not None: return reply pause = Consts.RETRY_BACKOFF[min(round_no, len(Consts.RETRY_BACKOFF) - 1)] round_no += 1 if give_up is not None and time.time() + pause > give_up: raise BackendError("all backends failed after %d rounds (%s)" % (round_no, "; ".join(errors[-6:]))) print(Render.paint("all %d backends failed (round %d), retrying in %ds: %s" % (len(cands), round_no, pause, (errors[-1] if errors else "")[:160]), Ansi.DIM), file=sys.stderr) if round_no == 1 and not Opencode.opencode_binary(): engine = Sandbox.container_engine() if engine and Sandbox.box_state(engine) == "missing": print(Render.paint( "no local opencode CLI, and no sandbox box set up yet either -- that box " "would also work as a last-resort backend (see box_opencode_candidates()). " "Run `tai.py --env sandbox` once to build it: a small image, podman or " "docker, and it only ever touches that one container, never the rest of " "this machine.", Ansi.DIM ), file=sys.stderr) time.sleep(pause) def all_candidates(self): http = [cand for cand in self.candidates() if cand["chat_url"] not in Consts.CLI_CHAT_URLS] cli = [cand for cand in self.candidates() if cand["chat_url"] in Consts.CLI_CHAT_URLS] ranked = sorted(enumerate(http), key=lambda pair: (ModelScoring.openrouter_free_tier(pair[1]), -self.health.weight(pair[1]["provider"], pair[1]["model"]), pair[0])) return [cand for _, cand in ranked] + cli + self.box_opencode_candidates() def box_opencode_candidates(self): """Last-resort only -- complete() reaches this via all_candidates(), never via ordered_candidates() (the round-0 fast path), so a normal turn that succeeds through a real backend never touches the box at all. If there's already a local opencode CLI, self.candidates() covers it directly and this stays empty; it only fires when that's missing AND the sandbox box is already running (never auto-created here -- see ensure_box(), which stays an explicit, approval-gated action via --env sandbox or the install tool).""" if not getattr(self.config, "opencode_enabled", True) or Opencode.opencode_binary(): return [] engine = Sandbox.container_engine() if not engine or Sandbox.box_state(engine) != "running": return [] if not Sandbox.box_opencode_binary(engine): return [] models = self.box_opencode_models(engine) return [ {"provider": Consts.OPENCODE_LABEL, "chat_url": Consts.OPENCODE_CHAT_URL, "models_url": Consts.OPENCODE_MODELS_URL, "key": "", "model": model, "profile": "", "runner": "box", "engine": engine} for model in models ] def box_opencode_models(self, engine): cached = self.health.cached_list("opencode:box:free") if cached is not None: return ModelScoring.prefer_muse(cached) try: done = Sandbox.box_exec(engine, ["opencode", "models", "opencode"], timeout=Consts.MODEL_LIST_TIMEOUT * 4) except (OSError, subprocess.SubprocessError): return [] text = done.stdout.decode("utf-8", "replace") models = [line.strip().split("/", 1)[1] for line in text.splitlines() if line.strip().startswith("opencode/")] models = [item for item in models if Text.is_free_model(item)] if models: self.health.store_list("opencode:box:free", models) return ModelScoring.prefer_muse(models) def call_candidate(self, cand, messages, tools, attempt, stream_sink, ephemeral): label, chat_url, key, model = cand["provider"], cand["chat_url"], cand["key"], cand["model"] extra = Opencode.opencode_headers() if cand["profile"] == "opencode" else None timeout = Consts.STREAM_TIMEOUT if stream_sink is not None else Consts.HTTP_TIMEOUT if chat_url == Consts.MUSE_CHAT_URL: return self.muse_turn(messages, tools, stream_sink, label, timeout=timeout, ephemeral=ephemeral) if chat_url == Consts.GROK_CHAT_URL: return self.grok_turn(messages, tools, stream_sink, label, timeout=timeout) if chat_url == Consts.CLAUDE_CHAT_URL: return self.claude_turn(messages, tools, stream_sink, label, timeout=timeout) if chat_url == Consts.CODEX_CHAT_URL: return self.codex_turn(messages, tools, stream_sink, label, timeout=timeout) if chat_url == Consts.GEMINI_CHAT_URL: return self.gemini_turn(messages, tools, stream_sink, label, timeout=timeout) if chat_url == Consts.OPENCODE_CHAT_URL: if cand.get("runner") == "box": return self.box_opencode_turn(messages, tools, stream_sink, label, model, cand["engine"], timeout=timeout) return self.opencode_turn(messages, tools, stream_sink, label, model, timeout=timeout) if stream_sink is None: return self.single_shot(chat_url, key, attempt, label, extra) return self.streaming(chat_url, key, attempt, label, stream_sink, extra) def try_candidate(self, cand, messages, tools, payload, stream_sink, ephemeral, errors): label, model = cand["provider"], cand["model"] attempt = dict(payload, model=model) tries = [(tools, attempt)] if attempt.get("tools"): tries.append((None, {name: value for name, value in attempt.items() if name not in ("tools", "tool_choice")})) for pos, (use_tools, body) in enumerate(tries): started = time.time() try: reply = self.call_candidate(cand, messages, use_tools, body, stream_sink, ephemeral) if not isinstance(reply, dict): raise BackendError("backend returned no reply") self.health.record(label, model, True, (time.time() - started) * 1000, max(1, len(reply.get("content") or "") // 4) / max(time.time() - started, 0.001)) return reply except BackendError as exc: if pos == 0 and exc.status == 400 and len(tries) > 1: continue self.health.record(label, model, False) errors.append("%s/%s: %s" % (label, model, exc)) return None except Exception as exc: self.health.record(label, model, False) errors.append("%s/%s: %s" % (label, model, Text.short_error(exc))) return None return None def single_shot(self, url, key, payload, label, extra=None): body = json.dumps(payload).encode("utf-8") request = urllib.request.Request(url, data=body, headers=self.headers(key, extra), method="POST") try: with urllib.request.urlopen(request, timeout=Consts.HTTP_TIMEOUT) as response: data = json.loads(response.read().decode("utf-8", "replace")) except urllib.error.HTTPError as exc: raise BackendError("HTTP %d: %s" % (exc.code, exc.read().decode("utf-8", "replace")[:300]), exc.code) except urllib.error.URLError as exc: raise BackendError(Text.short_error(exc)) failure = ModelScoring.embedded_failure_message(data) if failure: raise BackendError("upstream error: %s" % failure[:300]) try: message = data["choices"][0]["message"] except (KeyError, IndexError, TypeError): raise BackendError("upstream returned no choices: %s" % json.dumps(data)[:300]) return self.normalize(message, label) def normalize(self, message, label): calls = [] for call in message.get("tool_calls") or []: func = call.get("function") or {} calls.append({"id": call.get("id") or uuid.uuid4().hex, "name": func.get("name") or "", "arguments": Text.tool_args_text(func.get("arguments"))}) return {"role": "assistant", "content": message.get("content") or "", "reasoning": message.get("reasoning") or message.get("reasoning_content") or "", "tool_calls": calls, "backend": label} def streaming(self, url, key, payload, label, sink, extra=None): body = json.dumps(dict(payload, stream=True)).encode("utf-8") request = urllib.request.Request(url, data=body, headers=self.headers(key, extra), method="POST") try: response = urllib.request.urlopen(request, timeout=Consts.STREAM_TIMEOUT) except urllib.error.HTTPError as exc: raise BackendError("HTTP %d: %s" % (exc.code, exc.read().decode("utf-8", "replace")[:300]), exc.code) except urllib.error.URLError as exc: raise BackendError(Text.short_error(exc)) content_parts = [] reasoning_parts = [] calls = {} try: for raw in response: line = raw.decode("utf-8", "replace").strip() if not line.startswith("data:"): continue data = line[5:].strip() if data == "[DONE]": break try: event = json.loads(data) except ValueError: continue choices = event.get("choices") or [] if not choices: continue delta = choices[0].get("delta") or {} piece = delta.get("content") if piece: content_parts.append(piece) sink(piece) think = delta.get("reasoning") or delta.get("reasoning_content") if think: reasoning_parts.append(think) for call in delta.get("tool_calls") or []: slot = calls.setdefault(call.get("index", 0), {"id": "", "name": "", "arguments": ""}) if call.get("id"): slot["id"] = call["id"] func = call.get("function") or {} if func.get("name"): slot["name"] += func["name"] if func.get("arguments"): slot["arguments"] += Text.tool_args_text(func["arguments"]) except OSError as exc: if not content_parts and not calls: raise BackendError(Text.short_error(exc)) finally: response.close() ordered = [calls[pos] for pos in sorted(calls)] for slot in ordered: if not slot["id"]: slot["id"] = uuid.uuid4().hex return {"role": "assistant", "content": "".join(content_parts), "reasoning": "".join(reasoning_parts), "tool_calls": ordered, "backend": label} def muse_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT, ephemeral=False): binary = Muse.muse_binary() if not binary: raise BackendError("muse CLI not found (MUSE_BIN or PATH)") home = getattr(self.config, "home", None) if ephemeral or home is None: with tempfile.TemporaryDirectory(prefix="tai-muse-") as folder: schema_path = os.path.join(folder, "schema.json") prompt_path = os.path.join(folder, "prompt.txt") Muse.muse_write_file(schema_path, json.dumps(Muse.muse_output_schema())) Muse.muse_write_file(prompt_path, Muse.muse_prompt(messages, tools)) answer = Muse.muse_run(Muse.muse_argv(self.config, schema_path, prompt_path), folder, stream_sink, timeout) return self.muse_reply(answer, label) key = self.muse_key() release = Muse.muse_acquire(home, "session:" + key) def on_reap(ok): try: if not ok: self.reset_muse() finally: release() gated = Muse.muse_gate(on_reap) try: return self.muse_persistent(messages, tools, stream_sink, label, timeout, home, key, gated) except BaseException: if not gated.armed.is_set(): release() raise def muse_persistent(self, messages, tools, stream_sink, label, timeout, home, key, gated): def plan(data): entry = data.get(key) pending = None session_id = None if isinstance(entry, dict) and isinstance(entry.get("id"), str) and isinstance(entry.get("hashes"), list): pending = Muse.muse_session_split(messages, entry["hashes"]) if pending is not None and not Muse.muse_session_alive(entry["id"]): pending = None if pending is not None: session_id = entry["id"] if pending is None or not pending: session_id = str(uuid.uuid4()) pending = list(messages) return session_id, pending session_id, pending = Muse.muse_transact(home, plan) try: with tempfile.TemporaryDirectory(prefix="tai-muse-") as folder: schema_path = os.path.join(folder, "schema.json") prompt_path = os.path.join(folder, "prompt.txt") Muse.muse_write_file(schema_path, json.dumps(Muse.muse_output_schema())) Muse.muse_write_file(prompt_path, Muse.muse_prompt(pending, tools)) answer = Muse.muse_run(Muse.muse_argv(self.config, schema_path, prompt_path, session_id), folder, stream_sink, timeout, gated) except BackendError: self.reset_muse() raise def commit(data): data[key] = {"id": session_id, "hashes": [Muse.muse_message_hash(item) for item in messages], "updated": Text.now_iso()} Muse.muse_transact(home, commit) return self.muse_reply(answer, label) def opencode_turn(self, messages, tools, stream_sink, label, model, timeout=Consts.HTTP_TIMEOUT): answer = Opencode.opencode_parse(Opencode.opencode_run(model, Opencode.opencode_prompt(messages, tools), timeout)) if stream_sink is not None and answer.get("content"): stream_sink(answer["content"]) return self.muse_reply(answer, label) def box_opencode_turn(self, messages, tools, stream_sink, label, model, engine, timeout=Consts.HTTP_TIMEOUT): answer = Opencode.opencode_parse(Sandbox.box_opencode_run(engine, model, Opencode.opencode_prompt(messages, tools), timeout)) if stream_sink is not None and answer.get("content"): stream_sink(answer["content"]) return self.muse_reply(answer, label + " (box)") def grok_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT): answer = Grok.grok_parse_output(Grok.grok_run(Opencode.opencode_prompt(messages, tools), timeout)) if stream_sink is not None and answer.get("content"): stream_sink(answer["content"]) return self.muse_reply(answer, label) def claude_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT): answer = Claude.claude_parse_output(Claude.claude_run(Opencode.opencode_prompt(messages, tools), timeout)) if stream_sink is not None and answer.get("content"): stream_sink(answer["content"]) return self.muse_reply(answer, label) def codex_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT): stdout, outfile_text, used_schema = Codex.codex_run(Opencode.opencode_prompt(messages, tools), timeout) answer = Codex.codex_parse_output(stdout, outfile_text, used_schema) if stream_sink is not None and answer.get("content"): stream_sink(answer["content"]) return self.muse_reply(answer, label) def gemini_turn(self, messages, tools, stream_sink, label, timeout=Consts.HTTP_TIMEOUT): answer = Gemini.gemini_parse_output(Gemini.gemini_run(Opencode.opencode_prompt(messages, tools), timeout)) if stream_sink is not None and answer.get("content"): stream_sink(answer["content"]) return self.muse_reply(answer, label) def muse_reply(self, answer, label): calls = [] for call in answer.get("tool_calls") or []: calls.append({"id": uuid.uuid4().hex, "name": call.get("name") or "", "arguments": Text.tool_args_text(call.get("arguments"))}) return {"role": "assistant", "content": answer.get("content") or "", "reasoning": "", "tool_calls": calls, "backend": label}