feat: OpenMesh 基础平台与 MD/PDF 转换技能
Some checks failed
CI / pytest (push) Has been cancelled
CI / gui-unit (push) Has been cancelled
CI / gui-e2e (push) Has been cancelled

- 后端: coworker 智能体框架, WS API, 文件上传, 附件处理
- 前端: Open WebUI, 文件全量走 upload API (含 MD/TXT/JSON 等文本类)
- 技能: md-to-office (pandoc + wkhtmltopdf)
- 修复: 上传文件路径丢失, Agent 搜索浪费, 输出文件跑到 uploads/
- 打包: PyInstaller one-dir, 预打包 pandoc/wkhtmltopdf/chromium
This commit is contained in:
2026-09-13 23:41:04 +08:00
commit 6f402ffcee
638 changed files with 154534 additions and 0 deletions

View File

@@ -0,0 +1,54 @@
from .anthropic_provider import AnthropicProvider
from .bedrock_provider import BedrockProvider
from .base import (
AssistantTurn,
ModelCapabilities,
ProviderClient,
StreamChunk,
ToolCall,
)
from .capabilities import capabilities_for
from .codex_provider import CodexProvider
from .gemini_provider import GeminiProvider
from .openai_provider import OpenAIProvider, resolve_api_key
from .openai_responses import OpenAIResponsesProvider
from .registry import (
ProviderDescriptor,
ProviderField,
build_provider_client,
descriptor_configured,
detect_provider,
get_descriptor,
provider_descriptors,
provider_names,
verify_provider_key,
)
from .router import ProviderRouter
from .vertex_provider import VertexProvider
__all__ = [
"AssistantTurn",
"ModelCapabilities",
"ProviderClient",
"StreamChunk",
"ToolCall",
"AnthropicProvider",
"BedrockProvider",
"CodexProvider",
"GeminiProvider",
"OpenAIProvider",
"OpenAIResponsesProvider",
"VertexProvider",
"resolve_api_key",
"capabilities_for",
"ProviderRouter",
"ProviderDescriptor",
"ProviderField",
"provider_descriptors",
"provider_names",
"get_descriptor",
"build_provider_client",
"descriptor_configured",
"detect_provider",
"verify_provider_key",
]

View File

@@ -0,0 +1,653 @@
"""Anthropic provider — native Claude Messages API.
The runtime's canonical message format is OpenAI-shaped (that is what the engine builds and
persists), so this module is mostly a pair of pure converters: OpenAI-style messages → Anthropic
`messages` + `system`, and OpenAI function schemas → Anthropic `tools`. The Messages API differs
from chat.completions in ways the converters must absorb:
- `system` is a top-level param, not a message role.
- Assistant tool calls are `tool_use` content blocks (input is a dict, not a JSON string).
- Tool results are `tool_result` blocks that must ALL land in the single next user message —
N consecutive `role:"tool"` messages collapse into one user message here.
- `max_tokens` is required.
- Extended thinking (opt-in via the provider profile's `thinking_budget` field): responses
carry `thinking`/`redacted_thinking` blocks that MUST be replayed verbatim (signatures
and all) ahead of the same turn's tool_use blocks when returning tool results — they ride
the canonical assistant message as the `_anthropic` sidecar and are reattached here. The
thinking text also lands on `AssistantTurn.reasoning` for display.
"""
from __future__ import annotations
import json
import re
from typing import Any, Optional
from .base import (
AssistantTurn,
ModelCapabilities,
ProviderClient,
StreamChunk,
TokenUsage,
ToolCall,
)
from .capabilities import capabilities_for
def _usage_from(usage: Any) -> Optional[TokenUsage]:
"""Messages-API usage object → normalized counts (input_tokens excludes cache)."""
if usage is None:
return None
return TokenUsage(
input=int(getattr(usage, "input_tokens", 0) or 0),
output=int(getattr(usage, "output_tokens", 0) or 0),
cache_read=int(getattr(usage, "cache_read_input_tokens", 0) or 0),
cache_write=int(getattr(usage, "cache_creation_input_tokens", 0) or 0),
)
# Required by the Messages API; a ceiling, not a spend target. Sized for file
# generation, not just chat: a coworker writing a self-contained HTML report ships the
# whole file inside one tool call's arguments, and 16k proved too small in the field
# (the call truncates mid-arguments and the write fails). Current Claude models all
# accept ≥32k output.
DEFAULT_MAX_TOKENS = 32000
# Extended thinking is ON by default (owner call 2026-07-23: no user-facing setting —
# most users wouldn't know what a budget is; a per-turn composer control is future work).
# The provider profile's `thinking_budget` remains a hidden override: a number replaces
# the default, 0 disables thinking (where the model allows disabling).
DEFAULT_THINKING_BUDGET = 8192
# API drift (2026): thinking config is MODEL-FAMILY specific.
# - Pre-4.6 models (Haiku 4.5, Sonnet 4.5, Opus 4.5 and older): thinking needs
# {"type": "enabled", "budget_tokens": N}.
# - 4.6+ and the Claude 5 family (Fable/Mythos 5, Opus 4.8/4.7, Sonnet 5, the 4.6 pair):
# budget_tokens is deprecated/REMOVED (hard 400 on 4.7+: '"thinking.type.enabled" is
# not supported for this model') — use {"type": "adaptive"}. Fable 5 thinking is
# always on and can't be disabled. `display: "summarized"` is required to get trace
# text on 4.7+ (default "omitted" streams thinking blocks with EMPTY text).
_BUDGET_THINKING_PREFIXES = (
"claude-haiku-4-5",
"claude-sonnet-4-5",
"claude-opus-4-5",
"claude-opus-4-1",
"claude-opus-4-0",
"claude-sonnet-4-0",
"claude-3",
"claude-2",
)
def _uses_budget_thinking(model: str) -> bool:
return model.startswith(_BUDGET_THINKING_PREFIXES)
# Fable/Mythos 5 run safety classifiers that can decline benign-adjacent requests
# (HTTP 200, stop_reason "refusal", empty or partial content). Recommended posture is
# the server-side fallback: the API re-serves the declined request on Opus 4.8 within
# the same call. Beta header + param, beta messages endpoint.
_FALLBACK_BETA = "server-side-fallback-2026-06-01"
_FALLBACK_MODEL = "claude-opus-4-8"
def _needs_refusal_fallback(model: str) -> bool:
return model.startswith(("claude-fable", "claude-mythos"))
def _raise_on_refusal(stop_reason: Any, raw: Any) -> None:
"""A refusal that survived the fallback chain becomes a normal provider error —
the engine persists it as an error notice with Retry, instead of a silent blank."""
if stop_reason != "refusal":
return
details = getattr(raw, "stop_details", None)
category = getattr(details, "category", None)
suffix = f" (category: {category})" if category else ""
raise RuntimeError(
"Claude's safety filter declined this request"
+ suffix
+ " — try rephrasing, or switch model and press Retry."
)
# Anthropic stop_reason → the engine's OpenAI-shaped finish_reason vocabulary.
_STOP_REASON_MAP = {
"end_turn": "stop",
"tool_use": "tool_calls",
"max_tokens": "length",
"stop_sequence": "stop",
"refusal": "stop",
"pause_turn": "stop",
}
# Settings the Messages API accepts; everything else (frequency_penalty, …) is dropped.
_SETTINGS_WHITELIST = {
"max_tokens",
"temperature",
"top_p",
"top_k",
"stop_sequences",
"metadata",
"thinking",
}
# Sampling knobs the API rejects alongside extended thinking (temperature must stay 1).
_THINKING_INCOMPATIBLE = ("temperature", "top_p", "top_k")
_DATA_URL_RE = re.compile(
r"^data:(image/[a-z0-9.+-]+);base64,(.+)$", re.IGNORECASE | re.DOTALL
)
_PDF_DATA_URL_RE = re.compile(
r"^data:application/pdf;base64,(.+)$", re.IGNORECASE | re.DOTALL
)
def resolve_api_key(secrets: Any = None) -> Optional[str]:
"""Resolve the Anthropic API key: env `ANTHROPIC_API_KEY` first, else the SecretStore
`provider:anthropic` profile (`{api_key}`). Same contract as the OpenAI resolver: the
Tauri-launched sidecar does not inherit the shell env, so Settings-entered keys must work.
"""
import os
key = os.environ.get("ANTHROPIC_API_KEY")
if key:
return key
if secrets is not None:
profile = secrets.get("provider:anthropic") or {}
return profile.get("api_key") or None
return None
def _parse_args(raw: Any) -> dict[str, Any]:
"""Tool-call arguments: dict passthrough, JSON string parse, `{"_raw": …}` fallback."""
if isinstance(raw, dict):
return raw
if not raw:
return {}
try:
parsed = json.loads(raw)
return parsed if isinstance(parsed, dict) else {"_raw": raw}
except (TypeError, json.JSONDecodeError):
return {"_raw": raw}
def _image_block(url: str) -> Optional[dict[str, Any]]:
"""An OpenAI `image_url` part → an Anthropic image block. Attachments are always data URLs
(attachments.py); plain http(s) URLs map to a url source. Anything else → None."""
match = _DATA_URL_RE.match(url or "")
if match:
return {
"type": "image",
"source": {
"type": "base64",
"media_type": match.group(1).lower(),
"data": match.group(2),
},
}
if (url or "").startswith(("http://", "https://")):
return {"type": "image", "source": {"type": "url", "url": url}}
return None
def _document_block(part: dict[str, Any]) -> Optional[dict[str, Any]]:
"""An OpenAI `file` part (PDF data URL, attachments.py) → an Anthropic document block."""
file = part.get("file") or {}
match = _PDF_DATA_URL_RE.match(file.get("file_data") or "")
if not match:
return None
block: dict[str, Any] = {
"type": "document",
"source": {
"type": "base64",
"media_type": "application/pdf",
"data": match.group(1),
},
}
name = file.get("filename")
if name:
block["title"] = str(name)
return block
def _user_blocks(content: Any) -> list[dict[str, Any]]:
"""User content (str or OpenAI parts list) → Anthropic content blocks."""
if isinstance(content, str):
return [{"type": "text", "text": content}] if content else []
blocks: list[dict[str, Any]] = []
for part in content or []:
kind = part.get("type") if isinstance(part, dict) else None
if kind == "text":
text = part.get("text") or ""
if text:
blocks.append({"type": "text", "text": text})
elif kind == "image_url":
url = (part.get("image_url") or {}).get("url") or ""
block = _image_block(url)
blocks.append(
block
if block
else {"type": "text", "text": "[unsupported image attachment]"}
)
elif kind == "file":
block = _document_block(part)
blocks.append(
block
if block
else {"type": "text", "text": "[unsupported file attachment]"}
)
return blocks
def convert_messages(
messages: list[dict[str, Any]],
) -> tuple[Optional[str], list[dict[str, Any]]]:
"""OpenAI-shaped history → (`system`, Anthropic `messages`).
Leading system messages become the `system` param. Consecutive same-role outputs are folded
into one message — this is what collapses a run of `role:"tool"` results (one per parallel
call) into the single user message Anthropic requires, with any steering user text after.
"""
system_parts: list[str] = []
index = 0
while index < len(messages) and messages[index].get("role") == "system":
content = messages[index].get("content")
if isinstance(content, str) and content:
system_parts.append(content)
index += 1
converted: list[dict[str, Any]] = []
for message in messages[index:]:
role = message.get("role")
if role == "system":
# Defensive: a stray mid-thread system message rides as marked user text.
text = message.get("content") or ""
if text:
converted.append(
{
"role": "user",
"content": [
{"type": "text", "text": f"<system>\n{text}\n</system>"}
],
}
)
elif role == "user":
blocks = _user_blocks(message.get("content"))
if blocks:
converted.append({"role": "user", "content": blocks})
elif role == "assistant":
blocks = []
# Replay thinking/redacted_thinking blocks VERBATIM, ahead of the turn's own
# blocks — required whenever the turn's tool calls are being answered.
blocks.extend((message.get("_anthropic") or {}).get("blocks") or [])
text = message.get("content")
if isinstance(text, str) and text:
blocks.append({"type": "text", "text": text})
for call in message.get("tool_calls") or []:
function = call.get("function") or {}
blocks.append(
{
"type": "tool_use",
"id": call.get("id") or "",
"name": function.get("name") or "",
"input": _parse_args(function.get("arguments")),
}
)
if blocks:
converted.append({"role": "assistant", "content": blocks})
elif role == "tool":
converted.append(
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": message.get("tool_call_id") or "",
"content": str(message.get("content") or ""),
}
],
}
)
folded: list[dict[str, Any]] = []
for message in converted:
if folded and folded[-1]["role"] == message["role"]:
folded[-1]["content"].extend(message["content"])
else:
folded.append(message)
if not folded:
raise ValueError("no convertible messages for the Anthropic Messages API")
if folded[0]["role"] != "user":
folded.insert(
0, {"role": "user", "content": [{"type": "text", "text": "(continued)"}]}
)
return ("\n\n".join(system_parts) or None), folded
def convert_tools(tools: Optional[list[dict[str, Any]]]) -> list[dict[str, Any]]:
"""OpenAI function schemas → Anthropic tool definitions. Missing description is omitted;
missing/typeless parameters become an empty object schema (Anthropic requires one).
"""
converted = []
for tool in tools or []:
function = tool.get("function") or {}
entry: dict[str, Any] = {"name": function.get("name") or ""}
if function.get("description"):
entry["description"] = function["description"]
parameters = function.get("parameters")
if not isinstance(parameters, dict) or not parameters.get("type"):
parameters = {"type": "object", "properties": {}}
entry["input_schema"] = parameters
converted.append(entry)
return converted
def _add_cache_breakpoints(kwargs: dict[str, Any]) -> None:
"""Opt the request into prompt caching (5-minute ephemeral, prefix-matched).
Two breakpoints, the standard agent-loop shape:
- last system block — caches tools + system together (tools render first);
- last content block of the final message — caches the whole conversation
prefix, so each request re-reads the previous turns' cache and writes only
the new tail (append-only history keeps the prefix byte-identical).
Outbound-only: the canonical history never carries `cache_control` (the final
message's blocks are freshly built by convert_messages — thinking replays sit
in earlier assistant turns — and the marked block is copied, not mutated).
Prefixes under the model's cacheable minimum silently don't cache; reads bill
~0.1x and show up as `cache_read_input_tokens` (the metering's cache_read).
"""
marker = {"type": "ephemeral"}
system = kwargs.get("system")
if isinstance(system, str) and system:
kwargs["system"] = [{"type": "text", "text": system, "cache_control": marker}]
messages = kwargs.get("messages") or []
if messages:
content = messages[-1].get("content")
if isinstance(content, list) and content:
content[-1] = {**content[-1], "cache_control": marker}
def _reasoning_text(thinking_blocks: list[dict[str, Any]]) -> Optional[str]:
"""Display text for the GUI's disclosure — thinking text only (redacted stays opaque)."""
text = "".join(
b.get("thinking", "") for b in thinking_blocks if b.get("type") == "thinking"
)
return text or None
def _thinking_extras(thinking_blocks: list[dict[str, Any]]) -> dict[str, Any]:
"""Raw blocks → the `_anthropic` sidecar convert_messages replays (empty when none)."""
return {"_anthropic": {"blocks": thinking_blocks}} if thinking_blocks else {}
class AnthropicProvider(ProviderClient):
def __init__(
self,
client: Any = None,
*,
default_model: str = "claude-sonnet-4-6",
api_key: Optional[str] = None,
secrets: Any = None,
thinking_budget: Optional[int] = None,
):
# Mirrors OpenAIProvider: the SDK client is built lazily so engines can be assembled
# before any key exists; the key resolves at call time (explicit → env → SecretStore).
# Tests inject a `client` directly. `thinking_budget` (tokens, from the provider
# profile's optional field) opts every request into extended thinking.
self._client = client
self._api_key = api_key
self._secrets = secrets
self.default_model = default_model
self.thinking_budget = thinking_budget or 0
def _ensure_client(self) -> Any:
if self._client is None:
# Lazy import so the SDK is only required when actually talking to Anthropic.
from anthropic import Anthropic
key = self._api_key or resolve_api_key(self._secrets)
if not key:
raise RuntimeError(
"No Anthropic API key configured. Set ANTHROPIC_API_KEY in the environment, "
"or add your key in Manage → Configure Models."
)
self._client = Anthropic(api_key=key)
return self._client
def _request_kwargs(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]],
settings: dict[str, Any],
) -> dict[str, Any]:
system, converted = convert_messages(messages)
if "stop" in settings and "stop_sequences" not in settings:
stop = settings["stop"]
settings["stop_sequences"] = [stop] if isinstance(stop, str) else list(stop)
filtered = {k: v for k, v in settings.items() if k in _SETTINGS_WHITELIST}
if self.thinking_budget > 0 and "thinking" not in filtered:
if _uses_budget_thinking(model):
filtered["thinking"] = {
"type": "enabled",
"budget_tokens": self.thinking_budget,
}
else:
# 4.6+/Claude 5 family: adaptive only (budget_tokens 400s on 4.7+);
# display opt-in or the trace text arrives empty.
filtered["thinking"] = {"type": "adaptive", "display": "summarized"}
thinking = filtered.get("thinking") or {}
if thinking.get("type") == "enabled":
# Budget must fit under max_tokens.
budget = int(thinking.get("budget_tokens") or 0)
floor = max(DEFAULT_MAX_TOKENS, budget + 4096)
if int(filtered.get("max_tokens") or 0) <= budget:
filtered["max_tokens"] = floor
if thinking.get("type") in ("enabled", "adaptive"):
# Sampling knobs are rejected alongside thinking (and removed outright on 4.7+).
for key in _THINKING_INCOMPATIBLE:
filtered.pop(key, None)
filtered.setdefault("max_tokens", DEFAULT_MAX_TOKENS)
kwargs: dict[str, Any] = {"model": model, "messages": converted, **filtered}
if system:
kwargs["system"] = system
if tools:
kwargs["tools"] = convert_tools(tools)
_add_cache_breakpoints(kwargs)
return kwargs
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
) -> AssistantTurn:
kwargs = self._request_kwargs(
model=model, messages=messages, tools=tools, settings=settings
)
client = self._ensure_client()
# Stream-and-accumulate, not a plain create: the SDK REFUSES non-streaming
# requests whose max_tokens could exceed ~10 minutes (ValueError before any
# network I/O). With DEFAULT_MAX_TOKENS=32000 that killed every consumer of
# the non-streaming path — the auto-approve reviewer errored on ALL rows
# for every Anthropic model (found by the 2026-08-31 eval run; fail-closed,
# so verdicts fell back to asking a human). get_final_message() returns the
# same Message shape create() would.
if _needs_refusal_fallback(model):
with client.beta.messages.stream(
**kwargs,
betas=[_FALLBACK_BETA],
fallbacks=[{"model": _FALLBACK_MODEL}],
) as stream:
response = stream.get_final_message()
else:
with client.messages.stream(**kwargs) as stream:
response = stream.get_final_message()
text_parts: list[str] = []
tool_calls: list[ToolCall] = []
thinking_blocks: list[dict[str, Any]] = []
for block in getattr(response, "content", None) or []:
kind = getattr(block, "type", None)
if kind == "text":
text_parts.append(getattr(block, "text", "") or "")
elif kind == "tool_use":
tool_calls.append(
ToolCall(
id=getattr(block, "id", "") or "",
name=getattr(block, "name", "") or "",
arguments=dict(getattr(block, "input", None) or {}),
)
)
elif kind == "thinking":
thinking_blocks.append(
{
"type": "thinking",
"thinking": getattr(block, "thinking", "") or "",
"signature": getattr(block, "signature", "") or "",
}
)
elif kind == "redacted_thinking":
thinking_blocks.append(
{
"type": "redacted_thinking",
"data": getattr(block, "data", "") or "",
}
)
stop_reason = getattr(response, "stop_reason", None)
_raise_on_refusal(stop_reason, response)
return AssistantTurn(
text="".join(text_parts) or None,
tool_calls=tool_calls,
finish_reason=_STOP_REASON_MAP.get(stop_reason, stop_reason),
raw=response,
reasoning=_reasoning_text(thinking_blocks),
extras=_thinking_extras(thinking_blocks),
usage=_usage_from(getattr(response, "usage", None)),
)
def capabilities(self, model: str) -> ModelCapabilities:
return capabilities_for(model)
def stream(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
kwargs = self._request_kwargs(
model=model, messages=messages, tools=tools, settings=settings
)
kwargs["stream"] = True
client = self._ensure_client()
if _needs_refusal_fallback(model):
events = client.beta.messages.create(
**kwargs,
betas=[_FALLBACK_BETA],
fallbacks=[{"model": _FALLBACK_MODEL}],
)
else:
events = client.messages.create(**kwargs)
text_parts: list[str] = []
tool_accum: dict[int, dict[str, str]] = {}
# Thinking blocks accumulate per stream index and must be replayed verbatim later,
# so both the text and the signature_delta tail are collected (in block order).
thinking_accum: dict[int, dict[str, Any]] = {}
stop_reason = None
usage: Optional[TokenUsage] = None
last_message_delta: Any = None
for event in events:
kind = getattr(event, "type", None)
if kind == "message_start":
# Prompt-side counts (input + cache split) ride the opening event.
usage = (
_usage_from(getattr(getattr(event, "message", None), "usage", None))
or usage
)
elif kind == "content_block_start":
block = getattr(event, "content_block", None)
block_kind = getattr(block, "type", None)
if block_kind == "tool_use":
tool_accum[getattr(event, "index", 0)] = {
"id": getattr(block, "id", "") or "",
"name": getattr(block, "name", "") or "",
"json": "",
}
elif block_kind == "thinking":
thinking_accum[getattr(event, "index", 0)] = {
"type": "thinking",
"thinking": getattr(block, "thinking", "") or "",
"signature": getattr(block, "signature", "") or "",
}
elif block_kind == "redacted_thinking":
# Arrives whole — opaque data, no deltas.
thinking_accum[getattr(event, "index", 0)] = {
"type": "redacted_thinking",
"data": getattr(block, "data", "") or "",
}
elif kind == "content_block_delta":
delta = getattr(event, "delta", None)
delta_kind = getattr(delta, "type", None)
if delta_kind == "text_delta":
text = getattr(delta, "text", "") or ""
if text:
text_parts.append(text)
yield StreamChunk(text_delta=text)
elif delta_kind == "input_json_delta":
acc = tool_accum.get(getattr(event, "index", 0))
if acc is not None:
acc["json"] += getattr(delta, "partial_json", "") or ""
elif delta_kind == "thinking_delta":
acc = thinking_accum.get(getattr(event, "index", 0))
thought = getattr(delta, "thinking", "") or ""
if acc is not None and thought:
acc["thinking"] += thought
yield StreamChunk(reasoning_delta=thought)
elif delta_kind == "signature_delta":
acc = thinking_accum.get(getattr(event, "index", 0))
if acc is not None:
acc["signature"] = (acc.get("signature") or "") + (
getattr(delta, "signature", "") or ""
)
elif kind == "message_delta":
last_message_delta = getattr(event, "delta", None)
reason = getattr(last_message_delta, "stop_reason", None)
if reason:
stop_reason = reason
# Final (cumulative) output-token count rides message_delta.usage.
out = int(
getattr(getattr(event, "usage", None), "output_tokens", 0) or 0
)
if out:
usage = usage or TokenUsage()
usage.output = out
_raise_on_refusal(stop_reason, last_message_delta)
tool_calls = []
for index in sorted(tool_accum):
acc = tool_accum[index]
tool_calls.append(
ToolCall(
id=acc["id"], name=acc["name"], arguments=_parse_args(acc["json"])
)
)
thinking_blocks = [thinking_accum[i] for i in sorted(thinking_accum)]
yield StreamChunk(
turn=AssistantTurn(
text="".join(text_parts) or None,
tool_calls=tool_calls,
finish_reason=_STOP_REASON_MAP.get(stop_reason, stop_reason),
reasoning=_reasoning_text(thinking_blocks),
extras=_thinking_extras(thinking_blocks),
usage=usage,
)
)

136
coworker/providers/base.py Normal file
View File

@@ -0,0 +1,136 @@
"""Provider-agnostic model access layer.
The runtime never imports a provider SDK directly — it talks to a `ProviderClient`.
Implementations: `OpenAIResponsesProvider` (native OpenAI via `/v1/responses`),
`OpenAIProvider` (Chat Completions — the compat world), and the native
Anthropic/Gemini/Bedrock/Vertex providers, all selected by the registry/router.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Optional
@dataclass
class ToolCall:
"""A single tool call requested by the model, with parsed arguments."""
id: str
name: str
arguments: dict[str, Any] = field(default_factory=dict)
@dataclass
class TokenUsage:
"""Normalized token counts for one model round-trip.
`input` counts only fresh (uncached) prompt tokens; cached prompt tokens are
split into `cache_read`/`cache_write`. Providers that don't report a cache
split (Ollama, most compat vendors) leave the cache fields at 0. `output`
includes thinking tokens where the vendor bills them as output (Gemini).
"""
input: int = 0
output: int = 0
cache_read: int = 0
cache_write: int = 0
@property
def context_tokens(self) -> int:
"""Prompt-side total — what actually occupied the context window."""
return self.input + self.cache_read + self.cache_write
def as_dict(self) -> dict[str, int]:
return {
"input": self.input,
"output": self.output,
"cache_read": self.cache_read,
"cache_write": self.cache_write,
}
@dataclass
class AssistantTurn:
"""One assistant response: free text and/or a set of tool calls."""
text: Optional[str] = None
tool_calls: list[ToolCall] = field(default_factory=list)
finish_reason: Optional[str] = None
raw: Any = field(default=None, repr=False, compare=False)
# The model's thinking text (DeepSeek reasoning_content, Gemini thought summaries, …).
# Display-only: persisted on the assistant message as the `reasoning` sidecar and shown
# in the GUI, but stripped before every provider call — never replayed as context.
reasoning: Optional[str] = None
# Provider-private sidecars to persist on the canonical assistant message
# (underscore-prefixed keys, e.g. `_gemini` thought signatures). Contract: the
# owning provider consumes its own key when converting history; every other
# provider must strip or ignore foreign underscore keys before its wire call.
extras: dict[str, Any] = field(default_factory=dict)
# Token counts for this round-trip, normalized across providers. None when the
# backend didn't report usage (some compat servers) — never guessed.
usage: Optional[TokenUsage] = None
@property
def has_tool_calls(self) -> bool:
return bool(self.tool_calls)
@dataclass(frozen=True)
class ModelCapabilities:
"""What a given model/provider can do; used for graceful degradation."""
tools: bool = True
vision: bool = False
# Native PDF ingestion (OpenAI `file` part / Anthropic document / Gemini inline_data).
# Models without it get a local fallback: text extraction or page images (pdf_support.py).
pdf: bool = False
parallel_tool_calls: bool = True
streaming: bool = True
@dataclass
class StreamChunk:
"""One streamed piece: a text and/or reasoning delta, and/or (final) the full turn."""
text_delta: Optional[str] = None
reasoning_delta: Optional[str] = None
turn: Optional[AssistantTurn] = None
class ProviderClient(ABC):
"""Single-shot, provider-agnostic completion interface.
Deliberately blocking (the turn engine wraps it in `asyncio.to_thread`) and
deliberately without a `max_turns` loop — the runtime owns the agent loop.
"""
@abstractmethod
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
) -> AssistantTurn:
"""Return one assistant turn for the given messages/tools."""
@abstractmethod
def capabilities(self, model: str) -> ModelCapabilities:
"""Return capability flags for the given model."""
def stream(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
"""Yield StreamChunks. Default: no token streaming — one final chunk with the
full turn. Providers that support streaming (OpenAIProvider) override this."""
yield StreamChunk(
turn=self.complete(model=model, messages=messages, tools=tools, **settings)
)

View File

@@ -0,0 +1,579 @@
"""AWS Bedrock provider — one entry in Settings, two wire paths by model family.
Routed ids look like `bedrock:<family>/<bedrock model id>`; the router strips `bedrock:`
and this provider splits the family segment:
- `claude/…` → the native `AnthropicProvider` over the SDK's `AnthropicBedrock` client,
so Claude-on-Bedrock gets everything direct Anthropic gets (thinking, refusal handling).
- `other/…` → the Converse API (`bedrock-runtime.converse/converse_stream`), Bedrock's
unified wire format across Llama, Nova, Mistral, Cohere, DeepSeek, …
An id with no family segment falls back to Converse as-is — Converse serves every Bedrock
model (including Claude, minus the native extras), so a raw model id pasted without the
add-model dropdown still works.
Auth is ONE method at a time, selected by the profile's `auth_method` (a segmented choice
in Settings — owner call 2026-07-26, directness over field-precedence rules):
- `api_key` — a **Bedrock API key** (bearer token from the console, the no-CLI path);
rides `AWS_BEARER_TOKEN_BEDROCK`, which boto3 prefers over SigV4 for Bedrock calls.
- `profile` — a named `~/.aws` profile (covers `aws sso login`); blank → the default
credential chain (env vars / ~/.aws / role).
- `iam` — explicit access keys (+ optional STS session token).
Fields from non-selected methods are dropped at construction, so a stale stored value can
never leak into a different auth path (`AnthropicBedrock` raises outright on a mix). A
missing/unknown method falls back to whichever fields are present, api_key first.
boto3 is a lazy import (packaged via the `bedrock` extra) and returns PLAIN DICTS — every
response/stream mapping here is dict-shaped, unlike the attribute objects other SDKs return.
"""
from __future__ import annotations
import base64
import json
import os
import re
from typing import Any, Optional
from .anthropic_provider import AnthropicProvider
from .base import (
AssistantTurn,
ModelCapabilities,
ProviderClient,
StreamChunk,
TokenUsage,
ToolCall,
)
from .capabilities import capabilities_for
def _usage_from(usage: Any) -> Optional[TokenUsage]:
"""Converse `usage` dict → normalized counts (`inputTokens` excludes cache)."""
if not isinstance(usage, dict):
return None
return TokenUsage(
input=int(usage.get("inputTokens") or 0),
output=int(usage.get("outputTokens") or 0),
cache_read=int(usage.get("cacheReadInputTokens") or 0),
cache_write=int(usage.get("cacheWriteInputTokens") or 0),
)
# Converse has no required max token param but per-model defaults vary wildly (Meta's is
# 512 — an agent turn gets truncated mid-tool-call); 4096 fits every family's ceiling.
DEFAULT_MAX_TOKENS = 4096
# Converse stopReason → the engine's OpenAI-shaped finish_reason vocabulary.
_STOP_REASON_MAP = {
"end_turn": "stop",
"tool_use": "tool_calls",
"max_tokens": "length",
"stop_sequence": "stop",
"guardrail_intervened": "stop",
"content_filtered": "stop",
}
_DATA_URL_RE = re.compile(
r"^data:image/([a-z0-9.+-]+);base64,(.+)$", re.IGNORECASE | re.DOTALL
)
_PDF_DATA_URL_RE = re.compile(
r"^data:application/pdf;base64,(.+)$", re.IGNORECASE | re.DOTALL
)
# Bedrock document names: alphanumeric, whitespace, hyphens, parens, brackets only.
_DOC_NAME_RE = re.compile(r"[^A-Za-z0-9\s\-\(\)\[\]]+")
def _session_kwargs(
profile_name: Optional[str],
access_key_id: Optional[str],
secret_access_key: Optional[str],
session_token: Optional[str],
) -> dict[str, Any]:
"""boto3.Session kwargs for the explicit → profile → ambient resolution order."""
if access_key_id and secret_access_key:
kwargs: dict[str, Any] = {
"aws_access_key_id": access_key_id,
"aws_secret_access_key": secret_access_key,
}
if session_token:
kwargs["aws_session_token"] = session_token
return kwargs
if profile_name:
return {"profile_name": profile_name}
return {}
def _parse_args(raw: Any) -> dict[str, Any]:
if isinstance(raw, dict):
return raw
if not raw:
return {}
try:
parsed = json.loads(raw)
return parsed if isinstance(parsed, dict) else {"_raw": raw}
except (TypeError, json.JSONDecodeError):
return {"_raw": raw}
def _user_blocks(content: Any) -> list[dict[str, Any]]:
"""User content (str or OpenAI parts list) → Converse content blocks (bytes, not URLs)."""
if isinstance(content, str):
return [{"text": content}] if content else []
blocks: list[dict[str, Any]] = []
for part in content or []:
kind = part.get("type") if isinstance(part, dict) else None
if kind == "text":
if part.get("text"):
blocks.append({"text": part["text"]})
elif kind == "image_url":
url = (part.get("image_url") or {}).get("url") or ""
match = _DATA_URL_RE.match(url)
if match:
fmt = match.group(1).lower()
blocks.append(
{
"image": {
"format": "jpeg" if fmt == "jpg" else fmt,
"source": {"bytes": base64.b64decode(match.group(2))},
}
}
)
else: # Converse takes bytes only — no URL sources.
blocks.append({"text": "[unsupported image attachment]"})
elif kind == "file":
file = part.get("file") or {}
match = _PDF_DATA_URL_RE.match(file.get("file_data") or "")
if match:
name = _DOC_NAME_RE.sub("-", str(file.get("filename") or "document"))
blocks.append(
{
"document": {
"format": "pdf",
"name": name or "document",
"source": {"bytes": base64.b64decode(match.group(1))},
}
}
)
else:
blocks.append({"text": "[unsupported file attachment]"})
return blocks
def convert_messages(
messages: list[dict[str, Any]],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""OpenAI-shaped history → (Converse `system`, Converse `messages`).
Same shape discipline as the Anthropic converter (it's the same API family): leading
system messages become the top-level param, `role:"tool"` results become toolResult
blocks inside a user message, and consecutive same-role messages fold together so all
of a turn's parallel tool results land in the single next user message.
"""
system_parts: list[str] = []
index = 0
while index < len(messages) and messages[index].get("role") == "system":
content = messages[index].get("content")
if isinstance(content, str) and content:
system_parts.append(content)
index += 1
converted: list[dict[str, Any]] = []
for message in messages[index:]:
role = message.get("role")
if role == "system":
text = message.get("content") or ""
if text:
converted.append(
{"role": "user", "content": [{"text": f"<system>\n{text}\n</system>"}]}
)
elif role == "user":
blocks = _user_blocks(message.get("content"))
if blocks:
converted.append({"role": "user", "content": blocks})
elif role == "assistant":
blocks = []
text = message.get("content")
if isinstance(text, str) and text:
blocks.append({"text": text})
for call in message.get("tool_calls") or []:
function = call.get("function") or {}
blocks.append(
{
"toolUse": {
"toolUseId": call.get("id") or "",
"name": function.get("name") or "",
"input": _parse_args(function.get("arguments")),
}
}
)
if blocks:
converted.append({"role": "assistant", "content": blocks})
elif role == "tool":
converted.append(
{
"role": "user",
"content": [
{
"toolResult": {
"toolUseId": message.get("tool_call_id") or "",
"content": [
{"text": str(message.get("content") or "")}
],
}
}
],
}
)
folded: list[dict[str, Any]] = []
for message in converted:
if folded and folded[-1]["role"] == message["role"]:
folded[-1]["content"].extend(message["content"])
else:
folded.append(message)
if not folded:
raise ValueError("no convertible messages for the Bedrock Converse API")
if folded[0]["role"] != "user":
folded.insert(0, {"role": "user", "content": [{"text": "(continued)"}]})
system = [{"text": "\n\n".join(system_parts)}] if system_parts else []
return system, folded
def convert_tools(tools: Optional[list[dict[str, Any]]]) -> Optional[dict[str, Any]]:
"""OpenAI function schemas → Converse `toolConfig` (None when there are no tools —
Converse rejects an empty tool list)."""
specs = []
for tool in tools or []:
function = tool.get("function") or {}
parameters = function.get("parameters")
if not isinstance(parameters, dict) or not parameters.get("type"):
parameters = {"type": "object", "properties": {}}
spec: dict[str, Any] = {
"name": function.get("name") or "",
"inputSchema": {"json": parameters},
}
if function.get("description"):
spec["description"] = function["description"]
specs.append({"toolSpec": spec})
return {"tools": specs} if specs else None
def _inference_config(settings: dict[str, Any]) -> dict[str, Any]:
"""Whitelisted engine settings → Converse `inferenceConfig` (camelCase)."""
config: dict[str, Any] = {
"maxTokens": int(settings.get("max_tokens") or DEFAULT_MAX_TOKENS)
}
if settings.get("temperature") is not None:
config["temperature"] = settings["temperature"]
if settings.get("top_p") is not None:
config["topP"] = settings["top_p"]
stop = settings.get("stop_sequences") or settings.get("stop")
if stop:
config["stopSequences"] = [stop] if isinstance(stop, str) else list(stop)
return config
class _BedrockConverseClient(ProviderClient):
"""The `other/` family: any Bedrock model over the unified Converse API."""
def __init__(
self,
client: Any = None,
*,
region: Optional[str] = None,
bedrock_api_key: Optional[str] = None,
profile_name: Optional[str] = None,
access_key_id: Optional[str] = None,
secret_access_key: Optional[str] = None,
session_token: Optional[str] = None,
):
self._client = client # tests inject a dict-returning fake
self._region = region
self._bedrock_api_key = bedrock_api_key
self._session_kwargs = _session_kwargs(
profile_name, access_key_id, secret_access_key, session_token
)
def _ensure_client(self) -> Any:
if self._client is None:
try:
import boto3
except ImportError as exc:
raise RuntimeError(
"AWS Bedrock support needs the boto3 package — "
"install with `pip install 'openworker[bedrock]'`."
) from exc
# boto3 has no per-client bearer parameter — it only reads the env var, and
# prefers bearer auth for Bedrock whenever it's set. The sidecar process is
# ours, so publishing the configured key there is the supported path.
if self._bedrock_api_key:
os.environ["AWS_BEARER_TOKEN_BEDROCK"] = self._bedrock_api_key
session = boto3.session.Session(**self._session_kwargs)
self._client = session.client("bedrock-runtime", region_name=self._region)
return self._client
def _request_kwargs(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]],
settings: dict[str, Any],
) -> dict[str, Any]:
system, converted = convert_messages(messages)
kwargs: dict[str, Any] = {
"modelId": model,
"messages": converted,
"inferenceConfig": _inference_config(settings),
}
if system:
kwargs["system"] = system
tool_config = convert_tools(tools)
if tool_config:
kwargs["toolConfig"] = tool_config
return kwargs
@staticmethod
def _call(client: Any, method: str, kwargs: dict[str, Any]) -> Any:
try:
return getattr(client, method)(**kwargs)
except Exception as exc:
# boto3's "Unable to locate credentials" is famously cryptic — name the fix.
if exc.__class__.__name__ == "NoCredentialsError":
raise RuntimeError(
"No AWS credentials found — add keys or a profile in Settings ▸ "
"Models, or configure the AWS CLI (`aws configure` / `aws sso login`)."
) from exc
raise
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
) -> AssistantTurn:
kwargs = self._request_kwargs(
model=model, messages=messages, tools=tools, settings=settings
)
response = self._call(self._ensure_client(), "converse", kwargs)
text_parts: list[str] = []
reasoning_parts: list[str] = []
tool_calls: list[ToolCall] = []
content = ((response.get("output") or {}).get("message") or {}).get(
"content"
) or []
for block in content:
if "text" in block:
text_parts.append(block["text"] or "")
elif "toolUse" in block:
tool = block["toolUse"]
tool_calls.append(
ToolCall(
id=tool.get("toolUseId") or "",
name=tool.get("name") or "",
arguments=_parse_args(tool.get("input")),
)
)
elif "reasoningContent" in block:
text = (block["reasoningContent"].get("reasoningText") or {}).get(
"text"
) or ""
if text:
reasoning_parts.append(text)
stop_reason = response.get("stopReason")
return AssistantTurn(
text="".join(text_parts) or None,
tool_calls=tool_calls,
finish_reason=_STOP_REASON_MAP.get(stop_reason, stop_reason),
raw=response,
reasoning="".join(reasoning_parts) or None,
usage=_usage_from(response.get("usage")),
)
def stream(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
kwargs = self._request_kwargs(
model=model, messages=messages, tools=tools, settings=settings
)
response = self._call(self._ensure_client(), "converse_stream", kwargs)
text_parts: list[str] = []
reasoning_parts: list[str] = []
tool_accum: dict[int, dict[str, str]] = {}
stop_reason = None
usage: Optional[TokenUsage] = None
for event in response.get("stream") or []:
if "contentBlockStart" in event:
start = (event["contentBlockStart"].get("start") or {}).get("toolUse")
if start:
tool_accum[event["contentBlockStart"].get("contentBlockIndex", 0)] = {
"id": start.get("toolUseId") or "",
"name": start.get("name") or "",
"json": "",
}
elif "contentBlockDelta" in event:
block = event["contentBlockDelta"]
delta = block.get("delta") or {}
if delta.get("text"):
text_parts.append(delta["text"])
yield StreamChunk(text_delta=delta["text"])
elif "toolUse" in delta:
acc = tool_accum.get(block.get("contentBlockIndex", 0))
if acc is not None:
acc["json"] += delta["toolUse"].get("input") or ""
elif "reasoningContent" in delta:
thought = delta["reasoningContent"].get("text") or ""
if thought:
reasoning_parts.append(thought)
yield StreamChunk(reasoning_delta=thought)
elif "messageStop" in event:
stop_reason = event["messageStop"].get("stopReason") or stop_reason
elif "metadata" in event:
usage = _usage_from(event["metadata"].get("usage")) or usage
tool_calls = [
ToolCall(
id=tool_accum[i]["id"],
name=tool_accum[i]["name"],
arguments=_parse_args(tool_accum[i]["json"]),
)
for i in sorted(tool_accum)
]
yield StreamChunk(
turn=AssistantTurn(
text="".join(text_parts) or None,
tool_calls=tool_calls,
finish_reason=_STOP_REASON_MAP.get(stop_reason, stop_reason),
reasoning="".join(reasoning_parts) or None,
usage=usage,
)
)
def capabilities(self, model: str) -> ModelCapabilities:
return capabilities_for(f"bedrock:other/{model}")
class BedrockProvider(ProviderClient):
"""Family dispatcher: splits `<family>/<model id>` and delegates to the sub-client."""
def __init__(
self,
*,
region: Optional[str] = None,
auth_method: Optional[str] = None,
bedrock_api_key: Optional[str] = None,
profile_name: Optional[str] = None,
access_key_id: Optional[str] = None,
secret_access_key: Optional[str] = None,
session_token: Optional[str] = None,
claude_client: Optional[ProviderClient] = None,
converse_client: Optional[ProviderClient] = None,
):
# Narrow to the selected auth method here, once — stale values stored under a
# previously-selected method must never reach a different credential path.
if auth_method == "api_key":
profile_name = access_key_id = secret_access_key = session_token = None
elif auth_method == "profile":
bedrock_api_key = access_key_id = secret_access_key = session_token = None
elif auth_method == "iam":
bedrock_api_key = profile_name = None
self._region = region
self._bedrock_api_key = bedrock_api_key
self._profile_name = profile_name
self._access_key_id = access_key_id
self._secret_access_key = secret_access_key
self._session_token = session_token
# Test seams: pre-built sub-providers skip the SDK construction below.
self._clients: dict[str, ProviderClient] = {}
if claude_client is not None:
self._clients["claude"] = claude_client
if converse_client is not None:
self._clients["other"] = converse_client
@staticmethod
def _split(model: str) -> tuple[str, str]:
"""`claude/<id>` → the native path; anything else (including a raw Bedrock id with
no family segment) → Converse, which serves every Bedrock model."""
if "/" in model:
family, rest = model.split("/", 1)
if family in ("claude", "other"):
return family, rest
return "other", model
def _family_client(self, family: str) -> ProviderClient:
client = self._clients.get(family)
if client is None:
if family == "claude":
from anthropic import AnthropicBedrock
# A Bedrock API key (field or ambient env) takes the bearer path and
# EXCLUDES the SigV4 params — AnthropicBedrock raises on a mix.
bearer = self._bedrock_api_key or os.environ.get(
"AWS_BEARER_TOKEN_BEDROCK"
)
if bearer:
sdk = AnthropicBedrock(api_key=bearer, aws_region=self._region)
else:
sdk = AnthropicBedrock(
aws_region=self._region,
aws_profile=self._profile_name,
aws_access_key=self._access_key_id,
aws_secret_key=self._secret_access_key,
aws_session_token=self._session_token,
)
client = AnthropicProvider(client=sdk)
else:
client = _BedrockConverseClient(
region=self._region,
bedrock_api_key=self._bedrock_api_key,
profile_name=self._profile_name,
access_key_id=self._access_key_id,
secret_access_key=self._secret_access_key,
session_token=self._session_token,
)
self._clients[family] = client
return client
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
) -> AssistantTurn:
family, rest = self._split(model)
return self._family_client(family).complete(
model=rest, messages=messages, tools=tools, **settings
)
def stream(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
family, rest = self._split(model)
return self._family_client(family).stream(
model=rest, messages=messages, tools=tools, **settings
)
def capabilities(self, model: str) -> ModelCapabilities:
qualified = model if model.startswith("bedrock:") else f"bedrock:{model}"
return capabilities_for(qualified)

View File

@@ -0,0 +1,147 @@
"""Per-model capability probe.
A heuristic table for now (refined as we probe real providers/endpoints). Accepts
either bare model names (`gpt-5.5`) or provider-qualified ones (`openai:gpt-5.5`).
Custom user-added models can have their capabilities overridden via
`set_custom_capabilities()`, which is populated from preferences.
"""
from __future__ import annotations
import dataclasses
from typing import Optional
from .base import ModelCapabilities
# Per-model custom capability overrides, keyed by full model id (e.g. "openai:smesh-smartops").
# Set by the runtime from user preferences so custom models can declare vision/pdf support.
_custom_caps: dict[str, dict[str, bool]] = {}
def set_custom_capabilities(model: str, caps: dict[str, bool]) -> None:
"""Register or update custom capability flags for a user-added model."""
_custom_caps[model] = caps
def get_custom_capabilities(model: str) -> Optional[dict[str, bool]]:
return _custom_caps.get(model)
def clear_custom_capabilities() -> None:
_custom_caps.clear()
def _apply_custom_overrides(model: str, caps: ModelCapabilities) -> ModelCapabilities:
"""Apply user-configured capability overrides on top of heuristically-detected ones.
ModelCapabilities is a frozen dataclass, so we use dataclasses.replace()
to create a new instance with updated fields rather than mutating in place.
"""
custom = _custom_caps.get(model)
if not custom:
return caps
updates: dict[str, bool] = {}
for key in ("vision", "pdf", "tools", "parallel_tool_calls", "streaming"):
if key in custom:
updates[key] = bool(custom[key])
if not updates:
return caps
return dataclasses.replace(caps, **updates)
def capabilities_for(model: str) -> ModelCapabilities:
# Curated models answer from the matrix (exact full-id match — including reseller ids
# like `together:zai-org/GLM-5.2`, whose names defeat the prefix heuristics below).
# Custom user-added models fall through to the heuristics, at their own risk.
from .matrix import entry_for
entry = entry_for(model)
if entry is not None:
return _apply_custom_overrides(model, entry.caps)
provider = model.split(":", 1)[0].lower() if ":" in model else ""
name = model.split(":", 1)[-1].lower() # strip a provider prefix if present
# Ollama (local) models vary widely and many fake/mishandle parallel tool calls — assume
# tools work (we only point at tool-capable models) but stay conservative otherwise.
# Vision is detected from common model naming conventions (-vl, vision, llava, etc.).
if provider == "ollama":
_vision_patterns = ("-vl", "vision", "llava", "bakllava", "cogvlm", "minicpm-v")
has_vision = any(p in name for p in _vision_patterns)
return _apply_custom_overrides(
model,
ModelCapabilities(
tools=True, vision=has_vision, parallel_tool_calls=False, streaming=True
),
)
# Cloud-account providers (custom-added ids; curated ones answered from the matrix).
# The family segment decides: Claude keeps its native capabilities; everything else
# stays conservative until probed (Converse tool calling works across families, but
# parallel calls and vision vary per model).
if provider in ("bedrock", "vertex"):
if name.startswith(("claude/", "gemini/")):
return _apply_custom_overrides(
model,
ModelCapabilities(
tools=True, vision=True, pdf=True, parallel_tool_calls=True, streaming=True
),
)
return _apply_custom_overrides(
model,
ModelCapabilities(
tools=True, vision=False, parallel_tool_calls=False, streaming=True
),
)
# Claude / Gemini (both native): tools + vision + parallel tool calls + streaming. The
# engine executes parallel calls sequentially and each converter folds the results into
# the single next user message — exactly what both APIs require.
if provider in ("anthropic", "gemini"):
return _apply_custom_overrides(
model,
ModelCapabilities(
tools=True, vision=True, pdf=True, parallel_tool_calls=True, streaming=True
),
)
# Modern OpenAI GPT models: tools + vision + parallel tool calls + streaming.
if name.startswith(("gpt-5", "gpt-4")):
return _apply_custom_overrides(
model,
ModelCapabilities(
tools=True, vision=True, pdf=True, parallel_tool_calls=True, streaming=True
),
)
# OpenAI reasoning models: tools yes, parallel tool calls constrained.
if name.startswith(("o1", "o3", "o4")):
return _apply_custom_overrides(
model,
ModelCapabilities(
tools=True, vision=False, parallel_tool_calls=False, streaming=True
),
)
# OpenAI-compatible vendors (DeepSeek, Z AI/GLM, Kimi, MiniMax, Qwen, xAI/Grok, Mistral):
# tool calling + streaming across their current lineups; vision left off until probed
# per-model (several have vision variants, but the text flagships are what we suggest).
# Custom overrides can flip vision on for user-added fine-tunes like smesh-smartops.
if name.startswith(
("deepseek", "glm", "kimi", "minimax", "qwen", "grok", "mistral", "magistral")
):
return _apply_custom_overrides(
model,
ModelCapabilities(
tools=True, vision=False, parallel_tool_calls=True, streaming=True
),
)
# Conservative default for unknown models.
return _apply_custom_overrides(
model,
ModelCapabilities(
tools=True, vision=False, parallel_tool_calls=False, streaming=True
),
)

View File

@@ -0,0 +1,472 @@
"""Subscription sign-in for the `openai-codex` provider (OAuth 2.0 + PKCE).
Instead of an API key, the user signs in with their ChatGPT plan: a browser flow
against the vendor's auth service using their public subscription client id, with the
loopback redirect that id is registered for (the port is FIXED — any other port fails
the redirect-uri check server-side). Tokens land in the SecretStore profile
`provider:openai-codex` — the same local-only storage every provider profile uses,
never a plaintext config file — mirroring `mcp/oauth.py` (tokens + `tokens_issued_at`).
The pieces:
- `sign_in()` — async, explicit-action only: bind the loopback port, open the
browser, wait for the redirect, exchange the code, persist tokens + account id.
- `CodexTokenStore` — persistence + proactive refresh (JWT `exp`, sync httpx: the
provider is called from engine worker threads, `asyncio.to_thread` like its peers).
- `verify()` — the Test-button probe: one cheap authenticated request that
distinguishes signed-out vs expired vs OK.
The account id rides the token JWTs (the `https://api.openai.com/auth` claim); we
decode without verification — the backend verifies the token, we only route with it.
"""
from __future__ import annotations
import asyncio
import base64
import hashlib
import json
import logging
import secrets as pysecrets
import time
import uuid
from typing import Any, Optional
from urllib.parse import parse_qs, urlencode, urlsplit
logger = logging.getLogger(__name__)
AUTH_ISSUER = "https://auth.openai.com"
AUTHORIZE_URL = AUTH_ISSUER + "/oauth/authorize"
TOKEN_URL = AUTH_ISSUER + "/oauth/token"
# The public subscription client id (ships in the vendor's own tooling — not a secret).
CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
CALLBACK_PORT = 1455
CALLBACK_PATH = "/auth/callback"
# Registered redirect for CLIENT_ID, verbatim — host and port are not ours to choose.
REDIRECT_URI = f"http://localhost:{CALLBACK_PORT}{CALLBACK_PATH}"
SCOPE = "openid profile email offline_access"
ORIGINATOR = "openworker"
CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex"
PROFILE = "provider:openai-codex"
FLOW_TIMEOUT_SECONDS = 300
# Refresh this close to the JWT `exp` instead of sending an about-to-die bearer.
REFRESH_MARGIN_SECONDS = 300
_ACCOUNT_CLAIM = "https://api.openai.com/auth"
# Smallest curated model — the verify probe should cost as close to nothing as possible.
_VERIFY_MODEL = "gpt-5.1-codex-mini"
SIGNED_OUT_ERROR = (
"Not signed in to ChatGPT — connect your account in Settings ▸ Models to use "
"the subscription provider."
)
EXPIRED_ERROR = "ChatGPT session expired — sign in again in Settings ▸ Models."
PLAN_LIMIT_ERROR = (
"ChatGPT plan limit reached — your subscription's rolling usage window (about "
"5 hours) is used up. Wait for it to reset, upgrade the plan, or switch to an "
"API-key provider."
)
PORT_BUSY_ERROR = (
f"Port {CALLBACK_PORT} is already in use — the OpenAI Codex CLI is the usual "
"holder. Quit it and start the sign-in again."
)
class CodexAuthError(RuntimeError):
"""A subscription-auth failure with a user-readable message."""
class CodexSignInRequired(CodexAuthError):
"""No usable tokens — the fix is an explicit sign-in, never a silent browser."""
# -- PKCE / JWT helpers -----------------------------------------------------------
def create_pkce() -> tuple[str, str]:
"""(verifier, S256 challenge) per RFC 7636."""
verifier = pysecrets.token_urlsafe(64)
digest = hashlib.sha256(verifier.encode("ascii")).digest()
challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
return verifier, challenge
def build_authorize_url(state: str, challenge: str) -> str:
params = {
"response_type": "code",
"client_id": CLIENT_ID,
"redirect_uri": REDIRECT_URI,
"scope": SCOPE,
"state": state,
"code_challenge": challenge,
"code_challenge_method": "S256",
# The simplified-flow switch the subscription client id expects, plus the
# client-name tag the backend requires on every call.
"codex_cli_simplified_flow": "true",
"originator": ORIGINATOR,
}
return AUTHORIZE_URL + "?" + urlencode(params)
def _jwt_claims(token: str) -> dict[str, Any]:
"""Decode a JWT payload WITHOUT verification — we only read routing claims
(`exp`, the account object); the backend is the one verifying signatures."""
try:
payload = token.split(".")[1]
payload += "=" * (-len(payload) % 4)
claims = json.loads(base64.urlsafe_b64decode(payload.encode("ascii")))
return claims if isinstance(claims, dict) else {}
except Exception:
return {}
def account_id_from(tokens: dict[str, Any]) -> str:
"""The ChatGPT account id, from the auth claim of the id/access token."""
for key in ("id_token", "access_token"):
auth = _jwt_claims(tokens.get(key) or "").get(_ACCOUNT_CLAIM) or {}
if isinstance(auth, dict):
acct = auth.get("chatgpt_account_id") or auth.get("account_id") or ""
if acct:
return str(acct)
return ""
def backend_headers(account_id: str, session_id: str) -> dict[str, str]:
"""The non-auth headers every backend request must carry (auth is the bearer)."""
return {
"chatgpt-account-id": account_id,
"originator": ORIGINATOR,
"OpenAI-Beta": "responses=experimental",
"session-id": session_id,
}
# -- token persistence + refresh ----------------------------------------------------
def _token_post(data: dict[str, str], timeout: float = 30.0) -> Any:
"""One POST to the token endpoint (module-level so tests stub the wire here)."""
import httpx
return httpx.post(
TOKEN_URL, data=data, headers={"Accept": "application/json"}, timeout=timeout
)
def exchange_code(code: str, verifier: str, timeout: float = 30.0) -> dict[str, Any]:
"""authorization_code + PKCE verifier → the token set. Blocking (httpx sync);
`sign_in` runs it via `asyncio.to_thread`."""
resp = _token_post(
{
"grant_type": "authorization_code",
"code": code,
"redirect_uri": REDIRECT_URI,
"client_id": CLIENT_ID,
"code_verifier": verifier,
},
timeout,
)
if resp.status_code >= 300:
raise CodexAuthError(
f"Sign-in failed — token exchange returned HTTP {resp.status_code}."
)
return resp.json()
class CodexTokenStore:
"""Token set + account metadata in the `provider:openai-codex` SecretStore profile.
`access_token()` is what the provider calls per request: it hands back a live
bearer, refreshing proactively near the JWT `exp` and clearing the profile to a
clean signed-out state when the refresh token is rejected — never a crash loop.
"""
def __init__(self, secrets: Any) -> None:
self._secrets = secrets
def _data(self) -> dict[str, Any]:
if self._secrets is None:
return {}
return self._secrets.get(PROFILE) or {}
def _merge(self, patch: dict[str, Any]) -> None:
self._secrets.put(PROFILE, {**self._data(), **patch})
def signed_in(self) -> bool:
return bool(self._data().get("tokens"))
def account_label(self) -> Optional[str]:
data = self._data()
return data.get("account_email") or data.get("account_id") or None
def save(self, tokens: dict[str, Any]) -> None:
"""Persist a token response, keeping prior values a refresh omitted (the
refresh grant often returns no new refresh/id token)."""
existing = self._data().get("tokens") or {}
merged = {
k: (tokens.get(k) or existing.get(k))
for k in ("access_token", "refresh_token", "id_token")
}
merged = {k: v for k, v in merged.items() if v}
patch: dict[str, Any] = {
"tokens": merged,
"tokens_issued_at": int(time.time()),
}
account_id = account_id_from(merged) or self._data().get("account_id")
if account_id:
patch["account_id"] = account_id
email = _jwt_claims(merged.get("id_token") or "").get("email") or self._data().get(
"account_email"
)
if email:
patch["account_email"] = email
self._merge(patch)
def clear(self) -> bool:
if self._secrets is None:
return False
return bool(self._secrets.delete(PROFILE))
def access_token(self) -> tuple[str, str]:
"""(live access token, account id) — refreshing first when stale/absent."""
data = self._data()
tokens = data.get("tokens") or {}
access = tokens.get("access_token") or ""
if not access and not tokens.get("refresh_token"):
raise CodexSignInRequired(SIGNED_OUT_ERROR)
exp = _jwt_claims(access).get("exp")
stale = not access or (
isinstance(exp, (int, float)) and exp - time.time() < REFRESH_MARGIN_SECONDS
)
if stale:
return self.refresh()
return access, data.get("account_id") or ""
def refresh(self) -> tuple[str, str]:
"""refresh_token grant → fresh (access token, account id). A rejected refresh
token blanks the profile — the provider reads as cleanly signed out."""
refresh = (self._data().get("tokens") or {}).get("refresh_token") or ""
if not refresh:
self.clear()
raise CodexSignInRequired(EXPIRED_ERROR)
try:
resp = _token_post(
{
"grant_type": "refresh_token",
"refresh_token": refresh,
"client_id": CLIENT_ID,
}
)
except Exception as exc:
raise CodexAuthError(
"Couldn't reach the sign-in service to refresh the ChatGPT session "
f"({exc.__class__.__name__})."
) from exc
if 400 <= resp.status_code < 500:
self.clear()
raise CodexSignInRequired(EXPIRED_ERROR)
if resp.status_code >= 300:
raise CodexAuthError(
f"ChatGPT session refresh failed (HTTP {resp.status_code}) — try again."
)
self.save(resp.json())
data = self._data()
return (data.get("tokens") or {}).get("access_token") or "", (
data.get("account_id") or ""
)
# -- interactive sign-in flow -------------------------------------------------------
# The last authorize URL, surfaced over REST so the GUI can offer "reopen sign-in
# page" if the popup was lost (same affordance as mcp/oauth.py).
last_authorize_url: Optional[str] = None
_active_server: Optional[asyncio.AbstractServer] = None
_PAGE = """<!doctype html><meta charset="utf-8"><title>OpenWorker</title>
<body style="font-family: system-ui; margin: 4rem auto; max-width: 28rem; text-align: center;">
<h2>{title}</h2><p>{body}</p></body>"""
def _http_response(status: str, title: str, body: str) -> bytes:
html = _PAGE.format(title=title, body=body).encode("utf-8")
head = (
f"HTTP/1.1 {status}\r\nContent-Type: text/html; charset=utf-8\r\n"
f"Content-Length: {len(html)}\r\nConnection: close\r\n\r\n"
)
return head.encode("ascii") + html
async def _start_callback_server(
expected_state: str,
) -> tuple[asyncio.AbstractServer, "asyncio.Future[str]"]:
"""Bind the fixed loopback port and resolve the future with the auth code when
the redirect (carrying the matching `state`) lands."""
loop = asyncio.get_running_loop()
future: asyncio.Future[str] = loop.create_future()
async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
try:
request_line = await reader.readline()
while True: # drain headers; the redirect is a bare GET
line = await reader.readline()
if line in (b"\r\n", b"\n", b""):
break
parts = request_line.decode("ascii", errors="replace").split()
target = urlsplit(parts[1] if len(parts) > 1 else "/")
if target.path != CALLBACK_PATH:
writer.write(_http_response("404 Not Found", "Not found", ""))
return
query = parse_qs(target.query)
error = (query.get("error") or [""])[0]
code = (query.get("code") or [""])[0]
state = (query.get("state") or [""])[0]
if error:
writer.write(
_http_response(
"400 Bad Request",
"Sign-in failed",
"The service reported an error. Return to OpenWorker and try again.",
)
)
if not future.done():
future.set_exception(
CodexAuthError(f"Sign-in failed — the service returned: {error}")
)
return
# Same loopback gate as mcp/oauth.py: a stray local hit with the wrong
# state must not consume the flow — only the genuine redirect resolves it.
if not code or not pysecrets.compare_digest(state, expected_state):
writer.write(
_http_response(
"400 Bad Request",
"Nothing waiting for this sign-in",
"The sign-in may have timed out. Return to OpenWorker and start it again.",
)
)
return
writer.write(
_http_response(
"200 OK",
"Signed in",
"You can close this tab and return to OpenWorker.",
)
)
if not future.done():
future.set_result(code)
finally:
try:
await writer.drain()
writer.close()
except Exception:
pass
try:
server = await asyncio.start_server(handle, "127.0.0.1", CALLBACK_PORT)
except OSError as exc:
raise CodexAuthError(PORT_BUSY_ERROR) from exc
return server, future
async def sign_in(
secrets: Any,
*,
timeout: float = FLOW_TIMEOUT_SECONDS,
open_browser: bool = True,
) -> dict[str, Any]:
"""Run the full interactive flow: loopback server → browser → code → tokens.
Explicit-action only (a Settings button) — never called from an engine turn, so
unlike mcp/oauth.py it needs no non-interactive refusal path.
"""
global last_authorize_url, _active_server
if _active_server is not None:
# A stale flow lost its browser tab; the new one takes the port.
_active_server.close()
await _active_server.wait_closed()
_active_server = None
verifier, challenge = create_pkce()
state = pysecrets.token_urlsafe(24)
url = build_authorize_url(state, challenge)
last_authorize_url = url
server, code_future = await _start_callback_server(state)
_active_server = server
try:
if open_browser:
import webbrowser
logger.info("codex auth: opening browser for sign-in")
await asyncio.get_running_loop().run_in_executor(None, webbrowser.open, url)
try:
code = await asyncio.wait_for(code_future, timeout)
except asyncio.TimeoutError:
raise CodexAuthError(
"Sign-in timed out — the browser window was not completed in "
f"{int(timeout) // 60} minutes."
)
finally:
server.close()
await server.wait_closed()
if _active_server is server:
_active_server = None
tokens = await asyncio.to_thread(exchange_code, code, verifier)
store = CodexTokenStore(secrets)
store.save(tokens)
if not (store._data().get("tokens") or {}).get("access_token"):
store.clear()
raise CodexAuthError("Sign-in failed — the token response had no access token.")
return {"ok": True, "account": store.account_label()}
# -- verify probe -------------------------------------------------------------------
def verify(secrets: Any, timeout: float = 10.0) -> dict[str, Any]:
"""Test-button probe: one cheap authenticated request against the backend.
Distinguishes signed-out (no/rejected tokens) vs expired (401 with a bearer we
thought was live) vs OK. Never raises; {ok, error?, state?} like the other
provider verifies.
"""
import httpx
store = CodexTokenStore(secrets)
if not store.signed_in():
return {"ok": False, "error": SIGNED_OUT_ERROR, "state": "signed_out"}
try:
token, account = store.access_token()
except CodexSignInRequired as exc:
return {"ok": False, "error": str(exc), "state": "signed_out"}
except CodexAuthError as exc:
return {"ok": False, "error": str(exc)}
try:
resp = httpx.post(
CODEX_BASE_URL + "/responses",
headers={
"Authorization": f"Bearer {token}",
**backend_headers(account, str(uuid.uuid4())),
},
json={
"model": _VERIFY_MODEL,
"input": "Reply with OK.",
"store": False,
"stream": True,
"max_output_tokens": 16,
},
timeout=timeout,
)
except Exception as exc:
return {
"ok": False,
"error": f"Couldn't reach the ChatGPT backend ({exc.__class__.__name__}).",
}
if resp.status_code < 300:
return {"ok": True, "account": store.account_label()}
if resp.status_code in (401, 403):
return {"ok": False, "error": EXPIRED_ERROR, "state": "expired"}
if resp.status_code == 429:
# Auth is fine — the plan window is just used up right now.
return {"ok": True, "account": store.account_label(), "note": PLAN_LIMIT_ERROR}
return {
"ok": False,
"error": f"The ChatGPT backend returned HTTP {resp.status_code}.",
}

View File

@@ -0,0 +1,135 @@
"""`openai-codex` provider — OpenAI models through a ChatGPT subscription.
The backend speaks the same Responses wire as `/v1/responses` (stateless: full
history each turn, `store: false`, encrypted reasoning in the `_openai` sidecar), so
all conversion/parsing is inherited from `OpenAIResponsesProvider` — this subclass
only swaps the credential: a short-lived OAuth bearer from `codex_auth` instead of an
API key, plus the account/originator/session headers the backend requires.
Differences from the API-key path:
- The backend serves streamed responses only, so `complete()` drains `stream()`.
- 401 → one refresh-and-retry (the bearer died mid-flight); a rejected refresh
token surfaces as a typed sign-in-required error, never a crash loop.
- 429 → the plan's rolling usage window, surfaced as a user-readable message.
"""
from __future__ import annotations
import uuid
from typing import Any, Optional
from .base import AssistantTurn
from .codex_auth import (
CODEX_BASE_URL,
PLAN_LIMIT_ERROR,
CodexTokenStore,
backend_headers,
)
from .openai_responses import OpenAIResponsesProvider
def _status_code(exc: Exception) -> Optional[int]:
status = getattr(exc, "status_code", None)
if isinstance(status, int):
return status
status = getattr(getattr(exc, "response", None), "status_code", None)
return status if isinstance(status, int) else None
class CodexProvider(OpenAIResponsesProvider):
def __init__(
self,
client: Any = None,
*,
secrets: Any = None,
default_model: str = "gpt-5.2-codex",
reasoning_summary: bool = True,
):
super().__init__(
client=client,
default_model=default_model,
base_url=CODEX_BASE_URL,
reasoning_summary=reasoning_summary,
)
self._store = CodexTokenStore(secrets)
# One conversation per provider instance in practice (the router caches one
# client per provider); a uuid per instance satisfies the per-conversation
# session header without threading conversation ids through ProviderClient.
self._session_id = str(uuid.uuid4())
self._client_token: Optional[str] = None
self._injected = client is not None
def _ensure_client(self) -> Any:
if self._injected:
return self._client
# The bearer is short-lived: fetch per call (refreshes itself near expiry)
# and rebuild the SDK client whenever the token rotated.
token, account = self._store.access_token()
if self._client is None or token != self._client_token:
from openai import OpenAI
self._client = OpenAI(
api_key=token,
base_url=CODEX_BASE_URL,
default_headers=backend_headers(account, self._session_id),
)
self._client_token = token
return self._client
def _request_kwargs(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]],
settings: dict[str, Any],
) -> dict[str, Any]:
kwargs = super()._request_kwargs(
model=model, messages=messages, tools=tools, settings=settings
)
# This backend 400s ("Unsupported parameter") on standard sampling/cap knobs —
# max_output_tokens and temperature confirmed live, top_p same family — which
# silently killed every autotitle attempt on plan sessions (owner catch
# 2026-08-24). Callers may pass them freely; they just cannot ride to this
# backend.
for unsupported in ("max_output_tokens", "temperature", "top_p"):
kwargs.pop(unsupported, None)
# Unlike stock /v1/responses, this backend honors a reasoning effort knob.
effort = settings.get("reasoning_effort")
if isinstance(effort, str) and effort:
kwargs["reasoning"] = {**kwargs.get("reasoning", {}), "effort": effort}
# The backend rejects requests without instructions; history normally
# carries a system prompt — this is only the bare-call fallback.
kwargs.setdefault("instructions", "You are a helpful assistant.")
return kwargs
def _create(self, client: Any, kwargs: dict[str, Any]) -> Any:
try:
return super()._create(client, kwargs)
except Exception as exc:
status = _status_code(exc)
if status == 401 and not self._injected:
# The bearer died mid-flight: force one refresh and retry once.
# A rejected refresh raises CodexSignInRequired out of the store.
self._store.refresh()
self._client = None
self._client_token = None
return super()._create(self._ensure_client(), kwargs)
if status == 429:
raise RuntimeError(PLAN_LIMIT_ERROR) from exc
raise
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
) -> AssistantTurn:
# The backend only serves streamed responses — aggregate the stream.
turn: Optional[AssistantTurn] = None
for chunk in self.stream(model=model, messages=messages, tools=tools, **settings):
if chunk.turn is not None:
turn = chunk.turn
return turn if turn is not None else AssistantTurn()

View File

@@ -0,0 +1,57 @@
"""Friendly translation of model access + quota failures.
The picker now defaults to brand-new flagships (GPT-5.6 Sol, Claude Fable 5), and not every
account can use them: OpenAI is still rolling GPT-5.6 out per-organization, and both vendors
reject calls once quota/credits run out. Those failures arrive as terse SDK exceptions
wrapping JSON error bodies; this maps the well-known shapes to one actionable sentence.
Anything unrecognized returns None and the caller surfaces the raw error unchanged.
Matching is on the error BODY text (error codes/types), not just HTTP status — a 404 also
means "wrong base_url" and a 429 also means "slow down", and neither of those should be
dressed up as an access problem.
"""
from __future__ import annotations
from typing import Optional
# Error-body markers, verbatim from the vendors' error codes/messages:
# OpenAI: {"error": {"code": "model_not_found", "message": "The model `X` does not exist or
# you do not have access to it."}} (404/403) and {"code": "insufficient_quota"} (429).
# Anthropic: {"type": "not_found_error", "message": "model: X"} (404),
# {"type": "permission_error"} (403), and "credit balance is too low" (400).
_NO_ACCESS = (
"model_not_found",
"does not exist or you do not have access",
"does not have access to model",
"permission_error",
"permission denied",
)
_NO_QUOTA = (
"insufficient_quota",
"exceeded your current quota",
"credit balance is too low",
"billing hard limit",
)
def friendly_model_error(model: str, exc: Exception) -> Optional[str]:
"""One actionable sentence for "your account can't use this model" failures, or None."""
text = str(exc).lower()
no_access = (
f"Your account doesn't have access to {model} — new models can roll out "
"gradually or require a plan upgrade. Pick a different model, or check "
"the provider's console for availability."
)
if any(marker in text for marker in _NO_QUOTA):
return (
f"Your account is out of quota for {model} — add credits or raise the limit "
"in the provider's billing console, or pick a different model."
)
if any(marker in text for marker in _NO_ACCESS):
return no_access
# Anthropic's 404 body is just "model: <id>" under type not_found_error; require both
# halves so unrelated 404s (bad base_url, deleted resource) keep their raw message.
if "not_found_error" in text and f"model: {model.split(':')[-1].lower()}" in text:
return no_access
return None

View File

@@ -0,0 +1,547 @@
"""Gemini provider — native Google GenAI API (`google-genai` SDK).
Like the Anthropic provider, this is mostly a pair of pure converters from our canonical
OpenAI-shaped history to Gemini's `generateContent` format. The differences the converters
must absorb:
- The system prompt is `system_instruction` inside the request config, not a message role.
- Roles are `user`/`model`; tool results ride as `function_response` parts in a user message.
- Function calls carry NO ids — we synthesize `call_<n>` ids for the engine and map results
back by name (an id→name map built from the assistant turns during conversion).
- Tool parameter schemas are an OpenAPI 3.0 subset: unsupported JSON Schema keys
(`additionalProperties`, `$schema`, …) must be stripped or the API rejects the request.
- Gemini 3 thought signatures: response parts carry `thought_signature` (bytes) that MUST
be echoed back on the same parts in later requests — tool loops break without them. They
ride the canonical assistant message as the `_gemini` sidecar (base64 strings; the SDK's
`val_json_bytes="base64"` decodes them on send) and are reattached here. Parts flagged
`thought` are reasoning summaries, never answer text.
"""
from __future__ import annotations
import base64
import json
import re
from dataclasses import dataclass, field as dataclass_field
from typing import Any, Optional
from .base import (
AssistantTurn,
ModelCapabilities,
ProviderClient,
StreamChunk,
TokenUsage,
ToolCall,
)
from .capabilities import capabilities_for
def _usage_from(meta: Any) -> Optional[TokenUsage]:
"""`usage_metadata` → normalized counts. `prompt_token_count` INCLUDES the cached
share; thinking tokens are billed as output, so they fold into `output`."""
if meta is None:
return None
prompt = int(getattr(meta, "prompt_token_count", 0) or 0)
cached = int(getattr(meta, "cached_content_token_count", 0) or 0)
return TokenUsage(
input=max(prompt - cached, 0),
output=int(getattr(meta, "candidates_token_count", 0) or 0)
+ int(getattr(meta, "thoughts_token_count", 0) or 0),
cache_read=cached,
)
# Gemini finishReason → the engine's OpenAI-shaped finish_reason vocabulary. STOP maps to
# "tool_calls" instead when the turn contains function calls (Gemini has no distinct reason).
_FINISH_REASON_MAP = {
"STOP": "stop",
"MAX_TOKENS": "length",
"SAFETY": "stop",
"RECITATION": "stop",
"MALFORMED_FUNCTION_CALL": "stop",
}
# GenerateContentConfig keys we pass through; everything else (frequency_penalty, …) is dropped.
_SETTINGS_WHITELIST = {
"temperature",
"top_p",
"top_k",
"max_output_tokens",
"stop_sequences",
}
# The OpenAPI-subset schema keys Gemini function declarations accept.
_SCHEMA_KEYS = {
"type",
"format",
"description",
"nullable",
"enum",
"items",
"properties",
"required",
"anyOf",
"minimum",
"maximum",
"minItems",
"maxItems",
"minLength",
"maxLength",
"pattern",
"example",
"default",
"title",
}
_DATA_URL_RE = re.compile(
r"^data:(image/[a-z0-9.+-]+);base64,(.+)$", re.IGNORECASE | re.DOTALL
)
_PDF_DATA_URL_RE = re.compile(
r"^data:application/pdf;base64,(.+)$", re.IGNORECASE | re.DOTALL
)
def resolve_api_key(secrets: Any = None) -> Optional[str]:
"""Resolve the Gemini API key: env `GEMINI_API_KEY` (then `GOOGLE_API_KEY`, the SDK's own
convention) first, else the SecretStore `provider:gemini` profile (`{api_key}`)."""
import os
key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
if key:
return key
if secrets is not None:
profile = secrets.get("provider:gemini") or {}
return profile.get("api_key") or None
return None
def _image_part(url: str) -> Optional[dict[str, Any]]:
"""An OpenAI `image_url` part → a Gemini inline_data part. Attachments are always data
URLs (attachments.py). Plain http(s) URLs are not fetchable by the API → None."""
match = _DATA_URL_RE.match(url or "")
if match:
return {
"inline_data": {"mime_type": match.group(1).lower(), "data": match.group(2)}
}
return None
def _pdf_part(part: dict[str, Any]) -> Optional[dict[str, Any]]:
"""An OpenAI `file` part (PDF data URL, attachments.py) → a Gemini inline_data part."""
file = part.get("file") or {}
match = _PDF_DATA_URL_RE.match(file.get("file_data") or "")
if match:
return {"inline_data": {"mime_type": "application/pdf", "data": match.group(1)}}
return None
def _user_parts(content: Any) -> list[dict[str, Any]]:
"""User content (str or OpenAI parts list) → Gemini parts."""
if isinstance(content, str):
return [{"text": content}] if content else []
parts: list[dict[str, Any]] = []
for part in content or []:
kind = part.get("type") if isinstance(part, dict) else None
if kind == "text":
text = part.get("text") or ""
if text:
parts.append({"text": text})
elif kind == "image_url":
url = (part.get("image_url") or {}).get("url") or ""
image = _image_part(url)
parts.append(image if image else {"text": "[unsupported image attachment]"})
elif kind == "file":
pdf = _pdf_part(part)
parts.append(pdf if pdf else {"text": "[unsupported file attachment]"})
return parts
def _parse_args(raw: Any) -> dict[str, Any]:
"""Tool-call arguments: dict passthrough, JSON string parse, `{"_raw": …}` fallback."""
if isinstance(raw, dict):
return raw
if not raw:
return {}
try:
parsed = json.loads(raw)
return parsed if isinstance(parsed, dict) else {"_raw": raw}
except (TypeError, json.JSONDecodeError):
return {"_raw": raw}
def _result_payload(content: Any) -> dict[str, Any]:
"""A tool result string → the JSON object Gemini requires as a function response."""
if isinstance(content, dict):
return content
try:
parsed = json.loads(content)
return parsed if isinstance(parsed, dict) else {"result": parsed}
except (TypeError, json.JSONDecodeError):
return {"result": str(content or "")}
def convert_messages(
messages: list[dict[str, Any]],
) -> tuple[Optional[str], list[dict[str, Any]]]:
"""OpenAI-shaped history → (`system_instruction`, Gemini `contents`).
Function calls have no ids on the wire, so tool results are matched back to their function
NAME via an id→name map built from the assistant turns. Consecutive same-role outputs fold
into one content entry (tool-result runs collapse into a single user message, steering text
merging after — Gemini also dislikes non-alternating roles).
"""
system_parts: list[str] = []
index = 0
while index < len(messages) and messages[index].get("role") == "system":
content = messages[index].get("content")
if isinstance(content, str) and content:
system_parts.append(content)
index += 1
call_names: dict[str, str] = {}
converted: list[dict[str, Any]] = []
for message in messages[index:]:
role = message.get("role")
if role == "system":
# Defensive: a stray mid-thread system message rides as marked user text.
text = message.get("content") or ""
if text:
converted.append(
{
"role": "user",
"parts": [{"text": f"<system>\n{text}\n</system>"}],
}
)
elif role == "user":
parts = _user_parts(message.get("content"))
if parts:
converted.append({"role": "user", "parts": parts})
elif role == "assistant":
sidecar = message.get("_gemini") or {}
call_sigs = sidecar.get("call_sigs") or []
parts = []
text = message.get("content")
if isinstance(text, str) and text:
part: dict[str, Any] = {"text": text}
if sidecar.get("text_sig"):
part["thought_signature"] = sidecar["text_sig"]
parts.append(part)
for i, call in enumerate(message.get("tool_calls") or []):
function = call.get("function") or {}
name = function.get("name") or ""
call_names[call.get("id") or ""] = name
part = {
"function_call": {
"name": name,
"args": _parse_args(function.get("arguments")),
}
}
if i < len(call_sigs) and call_sigs[i]:
part["thought_signature"] = call_sigs[i]
parts.append(part)
if parts:
converted.append({"role": "model", "parts": parts})
elif role == "tool":
call_id = message.get("tool_call_id") or ""
converted.append(
{
"role": "user",
"parts": [
{
"function_response": {
"name": call_names.get(call_id) or call_id,
"response": _result_payload(message.get("content")),
}
}
],
}
)
folded: list[dict[str, Any]] = []
for message in converted:
if folded and folded[-1]["role"] == message["role"]:
folded[-1]["parts"].extend(message["parts"])
else:
folded.append(message)
if not folded:
raise ValueError("no convertible messages for the Gemini API")
if folded[0]["role"] != "user":
folded.insert(0, {"role": "user", "parts": [{"text": "(continued)"}]})
return ("\n\n".join(system_parts) or None), folded
def _sanitize_schema(schema: Any) -> Any:
"""Strip JSON Schema keys Gemini's OpenAPI subset rejects (recursively), and coerce
list-valued `type` (JSON Schema union, e.g. ["string", "number"] — common in vendor
MCP tool schemas) into shapes the API accepts: null joins as `nullable`, a single
remaining type stays `type`, several become `anyOf` (owner-hit 2026-07-23: monday's
compareValue union 400'd every Gemini turn in sessions with MCP tools)."""
if not isinstance(schema, dict):
return schema
cleaned: dict[str, Any] = {}
for key, value in schema.items():
if key not in _SCHEMA_KEYS:
continue
if key == "properties" and isinstance(value, dict):
cleaned[key] = {name: _sanitize_schema(sub) for name, sub in value.items()}
elif key == "items":
cleaned[key] = _sanitize_schema(value)
elif key == "anyOf" and isinstance(value, list):
cleaned[key] = [_sanitize_schema(sub) for sub in value]
elif key == "type" and isinstance(value, list):
types = [t for t in value if t != "null"]
if len(value) != len(types):
cleaned["nullable"] = True
if len(types) == 1:
cleaned["type"] = types[0]
elif types:
cleaned["anyOf"] = [{"type": t} for t in types]
else:
cleaned[key] = value
return cleaned
def convert_tools(tools: Optional[list[dict[str, Any]]]) -> list[dict[str, Any]]:
"""OpenAI function schemas → Gemini tool declarations (one tool, N function_declarations)."""
declarations = []
for tool in tools or []:
function = tool.get("function") or {}
entry: dict[str, Any] = {"name": function.get("name") or ""}
if function.get("description"):
entry["description"] = function["description"]
parameters = function.get("parameters")
if isinstance(parameters, dict) and parameters.get("properties"):
entry["parameters"] = _sanitize_schema(parameters)
# parameter-less functions omit `parameters` entirely (Gemini rejects empty objects)
declarations.append(entry)
return [{"function_declarations": declarations}] if declarations else []
def _sig_str(part: Any) -> Optional[str]:
"""A part's thought signature as a base64 string (jsonl-safe; the SDK's base64 bytes
validation turns it back into the original bytes on send)."""
sig = getattr(part, "thought_signature", None)
if not sig:
return None
if isinstance(sig, (bytes, bytearray)):
return base64.b64encode(bytes(sig)).decode("ascii")
return str(sig)
def _signature_extras(
text_sig: Optional[str], call_sigs: list[Optional[str]]
) -> dict[str, Any]:
"""Captured signatures → the `_gemini` assistant-message sidecar (empty when none)."""
if not text_sig and not any(call_sigs):
return {}
return {"_gemini": {"text_sig": text_sig, "call_sigs": call_sigs}}
@dataclass
class _Parsed:
"""One GenerateContentResponse (or streamed chunk), split into our concerns."""
texts: list[str] = dataclass_field(default_factory=list)
thoughts: list[str] = dataclass_field(default_factory=list) # `thought` summary parts
calls: list[ToolCall] = dataclass_field(default_factory=list)
finish: Optional[str] = None
text_sig: Optional[str] = None
call_sigs: list[Optional[str]] = dataclass_field(default_factory=list)
def _parse_candidate(response: Any) -> _Parsed:
"""Pull answer text, thought summaries, function calls (ids synthesized by the caller),
the finish reason, and thought signatures out of a response or streamed chunk. Parts
flagged `thought` are reasoning — their signature is kept, their text never joins the
answer."""
out = _Parsed()
candidates = getattr(response, "candidates", None) or []
if not candidates:
return out
candidate = candidates[0]
content = getattr(candidate, "content", None)
for part in getattr(content, "parts", None) or []:
sig = _sig_str(part)
function_call = getattr(part, "function_call", None)
if function_call is not None:
out.calls.append(
ToolCall(
id="",
name=getattr(function_call, "name", "") or "",
arguments=dict(getattr(function_call, "args", None) or {}),
)
)
out.call_sigs.append(sig)
continue
if sig:
out.text_sig = sig
text = getattr(part, "text", None)
if getattr(part, "thought", False):
if text:
out.thoughts.append(text)
continue
if text:
out.texts.append(text)
raw_finish = getattr(candidate, "finish_reason", None)
if raw_finish is not None:
out.finish = getattr(raw_finish, "name", None) or str(raw_finish)
return out
def _map_finish(finish: Optional[str], has_calls: bool) -> Optional[str]:
if has_calls:
return "tool_calls"
if finish is None:
return None
return _FINISH_REASON_MAP.get(finish, finish.lower())
class GeminiProvider(ProviderClient):
def __init__(
self,
client: Any = None,
*,
default_model: str = "gemini-2.5-flash",
api_key: Optional[str] = None,
secrets: Any = None,
):
# Mirrors AnthropicProvider: the SDK client is built lazily so engines can be assembled
# before any key exists; the key resolves at call time (explicit → env → SecretStore).
# Tests inject a `client` directly.
self._client = client
self._api_key = api_key
self._secrets = secrets
self.default_model = default_model
def _ensure_client(self) -> Any:
if self._client is None:
# Lazy import so the SDK is only required when actually talking to Gemini.
from google import genai
key = self._api_key or resolve_api_key(self._secrets)
if not key:
raise RuntimeError(
"No Gemini API key configured. Set GEMINI_API_KEY in the environment, "
"or add your key in Manage → Configure Models."
)
self._client = genai.Client(api_key=key)
return self._client
def _request_kwargs(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]],
settings: dict[str, Any],
) -> dict[str, Any]:
system, contents = convert_messages(messages)
if "max_tokens" in settings and "max_output_tokens" not in settings:
settings["max_output_tokens"] = settings["max_tokens"]
if "stop" in settings and "stop_sequences" not in settings:
stop = settings["stop"]
settings["stop_sequences"] = [stop] if isinstance(stop, str) else list(stop)
config: dict[str, Any] = {
k: v for k, v in settings.items() if k in _SETTINGS_WHITELIST
}
# Thinking models (2.5+/3.x — all our curated ids) think by default; ask for the
# thought SUMMARIES too so the GUI can show them. Parse-side keeps them out of
# answer text (`thought` parts → reasoning).
if model.startswith("gemini-"):
config["thinking_config"] = {"include_thoughts": True}
if system:
config["system_instruction"] = system
if tools:
converted = convert_tools(tools)
if converted:
config["tools"] = converted
return {"model": model, "contents": contents, "config": config}
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
) -> AssistantTurn:
kwargs = self._request_kwargs(
model=model, messages=messages, tools=tools, settings=settings
)
response = self._ensure_client().models.generate_content(**kwargs)
parsed = _parse_candidate(response)
tool_calls = [
ToolCall(id=f"call_{i}", name=c.name, arguments=c.arguments)
for i, c in enumerate(parsed.calls)
]
return AssistantTurn(
text="".join(parsed.texts) or None,
tool_calls=tool_calls,
finish_reason=_map_finish(parsed.finish, bool(tool_calls)),
raw=response,
reasoning="".join(parsed.thoughts) or None,
extras=_signature_extras(parsed.text_sig, parsed.call_sigs),
usage=_usage_from(getattr(response, "usage_metadata", None)),
)
def capabilities(self, model: str) -> ModelCapabilities:
return capabilities_for(model)
def stream(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
kwargs = self._request_kwargs(
model=model, messages=messages, tools=tools, settings=settings
)
client = self._ensure_client()
text_parts: list[str] = []
thought_parts: list[str] = []
calls: list[ToolCall] = []
finish = None
text_sig: Optional[str] = None
call_sigs: list[Optional[str]] = []
usage: Optional[TokenUsage] = None
# Unlike Anthropic, function_call parts arrive whole (args are a complete dict per
# part), so there is no JSON accumulation — just collect parts across chunks.
for chunk in client.models.generate_content_stream(**kwargs):
# Counts are cumulative per chunk; the last one seen is the final total.
chunk_usage = _usage_from(getattr(chunk, "usage_metadata", None))
if chunk_usage is not None:
usage = chunk_usage
parsed = _parse_candidate(chunk)
for thought in parsed.thoughts:
thought_parts.append(thought)
yield StreamChunk(reasoning_delta=thought)
for text in parsed.texts:
text_parts.append(text)
yield StreamChunk(text_delta=text)
calls.extend(parsed.calls)
call_sigs.extend(parsed.call_sigs)
if parsed.text_sig:
text_sig = parsed.text_sig
if parsed.finish:
finish = parsed.finish
tool_calls = [
ToolCall(id=f"call_{i}", name=c.name, arguments=c.arguments)
for i, c in enumerate(calls)
]
yield StreamChunk(
turn=AssistantTurn(
text="".join(text_parts) or None,
tool_calls=tool_calls,
finish_reason=_map_finish(finish, bool(tool_calls)),
reasoning="".join(thought_parts) or None,
extras=_signature_extras(text_sig, call_sigs),
usage=usage,
)
)

View File

@@ -0,0 +1,305 @@
"""The curated model matrix — the only models we actively suggest, label, and vouch for.
Keyed by the FULL routed id, exactly as the ProviderRouter receives it — including reseller
"ugly names" like ``together:zai-org/GLM-5.2`` (bare ids route to the OpenAI default). Each
entry carries the UI display label and the model's capabilities, making this the single
source of truth the capability probe and the GUI's pickers read from.
Deliberately SMALL (owner call, 2026-07-04): current-generation, agent-capable (tool-calling)
models only. It is not user-editable — users can still add any custom model string, which
falls back to the conservative heuristics in ``capabilities.py`` at their own risk of
degraded results. Ids verified against vendor/reseller catalogs on 2026-07-04; refresh the
reseller rows when catalogs rotate (they rename on every model generation).
Context windows (``context_window``, tokens) feed the GUI's context-fill meter. Entries
where the vendor spec wasn't re-checked stay ``None`` — the meter simply hides rather than
showing a made-up denominator. Values entered 2026-07-28 from vendor docs; verify alongside
the id refresh.
Resellers: Together + Fireworks + OpenRouter. TODO: add Groq entries here AND its
descriptor in ``registry.py`` once the current provider surface is tested — deliberately
deferred to bound how much needs verifying at once.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
from .base import ModelCapabilities
_AGENTIC = ModelCapabilities(
tools=True, vision=False, parallel_tool_calls=True, streaming=True
)
# The native three (OpenAI, Anthropic, Gemini) all take PDFs directly; every
# OpenAI-compatible vendor and reseller in the matrix does not (their chat APIs have
# no inline file part — checked 2026-07-17), so those fall back via pdf_support.py.
_AGENTIC_VISION = ModelCapabilities(
tools=True, vision=True, pdf=True, parallel_tool_calls=True, streaming=True
)
@dataclass(frozen=True)
class ModelEntry:
label: str # UI display name, e.g. "GLM-5.2 · via Together"
caps: ModelCapabilities = _AGENTIC
# Max context length in tokens (prompt side), for the GUI's context-fill meter.
# None = not verified against the vendor spec yet; the meter hides.
context_window: Optional[int] = None
MATRIX: dict[str, ModelEntry] = {
# -- first-party ------------------------------------------------------------
# GPT-5.6 (2026-07-09): number = generation, Sol/Terra/Luna = capability tiers.
# Bare "gpt-5.6" aliases to Sol server-side; we list the explicit tier ids only.
# Rolling out — accounts without access get a friendly error (providers/errors.py).
"gpt-5.6-sol": ModelEntry("GPT-5.6 Sol · OpenAI", _AGENTIC_VISION, 400_000),
"gpt-5.6-terra": ModelEntry("GPT-5.6 Terra · OpenAI", _AGENTIC_VISION, 400_000),
"gpt-5.6-luna": ModelEntry("GPT-5.6 Luna · OpenAI", _AGENTIC_VISION, 400_000),
"gpt-5.5": ModelEntry("GPT-5.5 · OpenAI", _AGENTIC_VISION, 400_000),
# ChatGPT-subscription catalog (the `openai-codex` OAuth provider). Curated to the
# ids the subscription backend actually serves; vision per the vendor's model docs,
# PDF unverified over this backend → local fallback via pdf_support.py.
# 5.6 tiers (Sol flagship / Terra balanced / Luna fast) serve over the subscription
# backend by plan — Sol is rate-limited on Plus, full on Pro.
"openai-codex:gpt-5.6-sol": ModelEntry(
"GPT-5.6 Sol · ChatGPT plan",
ModelCapabilities(
tools=True, vision=True, parallel_tool_calls=True, streaming=True
),
400_000,
),
"openai-codex:gpt-5.6-terra": ModelEntry(
"GPT-5.6 Terra · ChatGPT plan",
ModelCapabilities(
tools=True, vision=True, parallel_tool_calls=True, streaming=True
),
400_000,
),
"openai-codex:gpt-5.6-luna": ModelEntry(
"GPT-5.6 Luna · ChatGPT plan",
ModelCapabilities(
tools=True, vision=True, parallel_tool_calls=True, streaming=True
),
400_000,
),
"openai-codex:gpt-5.2-codex": ModelEntry(
"GPT-5.2 Codex · ChatGPT plan",
ModelCapabilities(
tools=True, vision=True, parallel_tool_calls=True, streaming=True
),
400_000,
),
"openai-codex:gpt-5.2": ModelEntry(
"GPT-5.2 · ChatGPT plan",
ModelCapabilities(
tools=True, vision=True, parallel_tool_calls=True, streaming=True
),
400_000,
),
"openai-codex:gpt-5.1-codex": ModelEntry(
"GPT-5.1 Codex · ChatGPT plan",
ModelCapabilities(
tools=True, vision=True, parallel_tool_calls=True, streaming=True
),
400_000,
),
"openai-codex:gpt-5.1-codex-mini": ModelEntry(
"GPT-5.1 Codex Mini · ChatGPT plan", _AGENTIC, 400_000
),
# Fable 5 (2026-06-09) is GA; its Mythos 5 sibling is approved-orgs-only, so it
# stays out of a picker meant for the public.
"anthropic:claude-fable-5": ModelEntry(
"Claude Fable 5 · Anthropic", _AGENTIC_VISION, 1_000_000
),
"anthropic:claude-opus-4-8": ModelEntry(
"Claude Opus 4.8 · Anthropic", _AGENTIC_VISION, 200_000
),
"anthropic:claude-sonnet-4-6": ModelEntry(
"Claude Sonnet 4.6 · Anthropic", _AGENTIC_VISION, 200_000
),
"anthropic:claude-haiku-4-5": ModelEntry(
"Claude Haiku 4.5 · Anthropic", _AGENTIC_VISION, 200_000
),
# Gemini 3 (thought signatures required in tool loops — carried via the `_gemini`
# message sidecar, see gemini_provider.py; ids from the vendor catalog 2026-07-22).
"gemini:gemini-3.1-pro-preview": ModelEntry(
"Gemini 3.1 Pro · Google", _AGENTIC_VISION, 1_048_576
),
"gemini:gemini-3.6-flash": ModelEntry(
"Gemini 3.6 Flash · Google", _AGENTIC_VISION, 1_048_576
),
"gemini:gemini-2.5-pro": ModelEntry(
"Gemini 2.5 Pro · Google", _AGENTIC_VISION, 1_048_576
),
"gemini:gemini-2.5-flash": ModelEntry(
"Gemini 2.5 Flash · Google", _AGENTIC_VISION, 1_048_576
),
# Ark Responses API providers (verified 2026-08-14). BytePlus pay-as-you-go and
# Volcengine Agent Plan intentionally use separate provider prefixes because their
# endpoints, credentials, regions, and model catalogs are not interchangeable.
"ark:dola-seed-evolving-latest-version": ModelEntry(
"Dola Seed Evolving · BytePlus Ark", context_window=256_000
),
"ark:dola-seed-2-1-turbo-260628": ModelEntry(
"Dola Seed 2.1 Turbo · BytePlus Ark", context_window=256_000
),
"ark-agent-plan-cn:doubao-seed-evolving": ModelEntry(
"Doubao Seed Evolving · Volcengine Agent Plan", context_window=256_000
),
"ark-agent-plan-cn:doubao-seed-2.1-turbo": ModelEntry(
"Doubao Seed 2.1 Turbo · Volcengine Agent Plan", context_window=256_000
),
# -- direct OpenAI-compatible vendors ----------------------------------------
# Muse Spark (Meta Model API, public preview 2026-07-09): multimodal + tools via
# their OpenAI-compat surface. Vision yes; PDFs unverified over compat — falls
# back via pdf_support.py like the other compat vendors.
"meta:muse-spark-1.1": ModelEntry(
"Muse Spark 1.1 · Meta",
ModelCapabilities(
tools=True, vision=True, parallel_tool_calls=True, streaming=True
),
),
"zai:glm-5.2": ModelEntry("GLM-5.2 · Z AI", _AGENTIC, 128_000),
"deepseek:deepseek-v4-flash": ModelEntry(
"DeepSeek V4 Flash · DeepSeek", _AGENTIC, 128_000
),
"deepseek:deepseek-v4-pro": ModelEntry(
"DeepSeek V4 Pro · DeepSeek", _AGENTIC, 128_000
),
"kimi:kimi-k2.6": ModelEntry("Kimi K2.6 · Moonshot", _AGENTIC, 256_000),
"minimax:MiniMax-M2.5": ModelEntry("MiniMax M2.5 · MiniMax"),
"qwen:qwen3-max": ModelEntry("Qwen3 Max · Alibaba", _AGENTIC, 256_000),
"xai:grok-4.3": ModelEntry("Grok 4.3 · xAI", _AGENTIC, 256_000),
"mistral:mistral-large-latest": ModelEntry(
"Mistral Large · Mistral", _AGENTIC, 128_000
),
# -- resellers (their model namespaces, verbatim) -----------------------------
"together:thinkingmachines/Inkling": ModelEntry("Inkling · via Together"),
"together:zai-org/GLM-5.2": ModelEntry("GLM-5.2 · via Together", _AGENTIC, 128_000),
# Kimi K3 on Together (landed late July 2026): 1M window, native vision; PDFs
# unverified over the compat surface (falls back via pdf_support.py, like Muse Spark).
"together:moonshotai/Kimi-K3": ModelEntry(
"Kimi K3 · via Together",
ModelCapabilities(
tools=True, vision=True, parallel_tool_calls=True, streaming=True
),
1_000_000,
),
"together:moonshotai/Kimi-K2.7-Code": ModelEntry(
"Kimi K2.7 Code · via Together", _AGENTIC, 256_000
),
"together:moonshotai/Kimi-K2.6": ModelEntry(
"Kimi K2.6 · via Together", _AGENTIC, 256_000
),
"together:deepseek-ai/DeepSeek-V4-Pro": ModelEntry(
"DeepSeek V4 Pro · via Together", _AGENTIC, 128_000
),
"together:meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": ModelEntry(
"Llama 4 Maverick · via Together", _AGENTIC, 1_000_000
),
"fireworks:accounts/fireworks/models/glm-5p2": ModelEntry(
"GLM-5.2 · via Fireworks", _AGENTIC, 128_000
),
"fireworks:accounts/fireworks/models/kimi-k2p6": ModelEntry(
"Kimi K2.6 · via Fireworks", _AGENTIC, 256_000
),
"fireworks:accounts/fireworks/models/deepseek-v4-pro": ModelEntry(
"DeepSeek V4 Pro · via Fireworks", _AGENTIC, 128_000
),
"fireworks:accounts/fireworks/models/llama4-maverick-instruct-basic": ModelEntry(
"Llama 4 Maverick · via Fireworks", _AGENTIC, 1_000_000
),
# OpenRouter slugs are lowercase `<lab>/<model>` (checked against their catalog
# 2026-07-25); same labs as above, one key for all of them.
"openrouter:z-ai/glm-5.2": ModelEntry("GLM-5.2 · via OpenRouter", _AGENTIC, 128_000),
"openrouter:moonshotai/kimi-k2.6": ModelEntry(
"Kimi K2.6 · via OpenRouter", _AGENTIC, 256_000
),
"openrouter:deepseek/deepseek-v4-pro": ModelEntry(
"DeepSeek V4 Pro · via OpenRouter", _AGENTIC, 128_000
),
"openrouter:meta-llama/llama-4-maverick": ModelEntry(
"Llama 4 Maverick · via OpenRouter", _AGENTIC, 1_000_000
),
# Stealth/cloaked alpha (catalog-checked 2026-08-24: 1,048,576 ctx, tool calling).
# These are temporary lab previews — expect the slug to vanish when the lab ships
# the real model; keep it until OpenRouter retires it.
"openrouter:stealth/ox-alpha": ModelEntry(
"Ox Alpha · via OpenRouter", _AGENTIC, 1_048_576
),
# -- cloud accounts (models running in the user's own AWS/GCP) ----------------
# Bedrock ids carry a family segment (claude/ → native Anthropic path, other/ →
# Converse) plus AWS's own `-v<n>:<m>` version suffix. Some regions require the
# `us.`/`eu.` cross-region inference-profile prefix — custom add-model accepts those.
"bedrock:claude/anthropic.claude-sonnet-4-6-v1:0": ModelEntry(
"Claude Sonnet 4.6 · AWS Bedrock", _AGENTIC_VISION, 200_000
),
"bedrock:claude/anthropic.claude-haiku-4-5-v1:0": ModelEntry(
"Claude Haiku 4.5 · AWS Bedrock", _AGENTIC_VISION, 200_000
),
"bedrock:other/amazon.nova-2-pro-v1:0": ModelEntry(
"Nova 2 Pro · AWS Bedrock", _AGENTIC, 300_000
),
"bedrock:other/meta.llama4-maverick-17b-instruct-v1:0": ModelEntry(
"Llama 4 Maverick · AWS Bedrock", _AGENTIC, 1_000_000
),
"bedrock:other/mistral.mistral-large-3-v1:0": ModelEntry(
"Mistral Large 3 · AWS Bedrock", _AGENTIC, 128_000
),
# Live-verified on Converse 2026-07-26 (complete/stream/tool round trip); asked for
# two tool calls it emits them one at a time, so parallel stays off.
"bedrock:other/nvidia.nemotron-super-3-120b": ModelEntry(
"Nemotron Super 3 120B · AWS Bedrock",
ModelCapabilities(
tools=True, vision=False, parallel_tool_calls=False, streaming=True
),
),
# Vertex ids carry a family segment too (gemini/ and claude/ → native paths,
# openweight/ → the MaaS OpenAI-compat endpoint, keeping the publisher segment).
"vertex:gemini/gemini-3.1-pro-preview": ModelEntry(
"Gemini 3.1 Pro · Vertex AI", _AGENTIC_VISION, 1_048_576
),
"vertex:gemini/gemini-3.6-flash": ModelEntry(
"Gemini 3.6 Flash · Vertex AI", _AGENTIC_VISION, 1_048_576
),
"vertex:claude/claude-sonnet-4-6": ModelEntry(
"Claude Sonnet 4.6 · Vertex AI", _AGENTIC_VISION, 200_000
),
"vertex:claude/claude-haiku-4-5": ModelEntry(
"Claude Haiku 4.5 · Vertex AI", _AGENTIC_VISION, 200_000
),
"vertex:openweight/meta/llama-4-maverick-17b-128e-instruct-maas": ModelEntry(
"Llama 4 Maverick · Vertex AI", _AGENTIC, 1_000_000
),
"vertex:openweight/qwen/qwen3-coder-480b-a35b-instruct-maas": ModelEntry(
"Qwen3 Coder · Vertex AI", _AGENTIC, 256_000
),
}
def entry_for(model: str) -> ModelEntry | None:
return MATRIX.get(model)
def model_labels() -> dict[str, str]:
"""Full-id → display-label map, shipped to the GUI so every picker shows human names."""
return {mid: e.label for mid, e in MATRIX.items()}
def model_context_windows() -> dict[str, int]:
"""Full-id → context-window map (verified entries only), for the GUI's fill meter."""
return {
mid: e.context_window for mid, e in MATRIX.items() if e.context_window
}
def models_for_provider(provider: str) -> list[str]:
"""BARE model ids (prefix stripped) the matrix curates for a provider — feeds the
Settings pane's suggestions and the composer picker so both stay in lockstep with the
matrix. OpenAI entries are stored without a prefix (bare ids route to the OpenAI
default), so its list is every un-prefixed id."""
if provider == "openai":
return [mid for mid in MATRIX if ":" not in mid]
prefix = provider + ":"
return [mid[len(prefix) :] for mid in MATRIX if mid.startswith(prefix)]

View File

@@ -0,0 +1,621 @@
"""OpenAI Chat Completions provider — the compat workhorse.
Uses the OpenAI Python SDK `chat.completions` API only, which is what the entire
OpenAI-compatible world implements: the compat vendors (DeepSeek, Z AI, Kimi, …),
resellers, Ollama, custom endpoints (Azure OpenAI, vLLM), and the Bedrock/Vertex MaaS
paths. Native OpenAI models (the `openai` provider with no custom endpoint) route to
`openai_responses.OpenAIResponsesProvider` instead — Chat Completions rejects function
tools combined with reasoning on GPT-5.6+, so reasoning + tools needs `/v1/responses`.
"""
from __future__ import annotations
import json
import re
from typing import Any, Optional
from .base import (
AssistantTurn,
ModelCapabilities,
ProviderClient,
StreamChunk,
TokenUsage,
ToolCall,
)
from .capabilities import capabilities_for
def resolve_api_key(secrets: Any = None) -> Optional[str]:
"""Resolve the OpenAI API key: env `OPENAI_API_KEY` first, else the SecretStore
`provider:openai` profile (`{api_key}`). Lets a Tauri-launched sidecar — which does NOT
inherit the shell env — still find a key the user entered in Settings. The value never
enters the model context; it only configures the SDK client.
"""
import os
key = os.environ.get("OPENAI_API_KEY")
if key:
return key
if secrets is not None:
profile = secrets.get("provider:openai") or {}
return profile.get("api_key") or None
return None
# GPT-5.6 (2026-07) defaults reasoning_effort to "medium" server-side, and
# /v1/chat/completions rejects function tools combined with any effort other than
# "none" ("use /v1/responses"). Native OpenAI now routes to the Responses provider,
# but GPT-5.6 can still land here through a custom endpoint (Azure OpenAI serves the
# same wire), so keep pinning effort to none whenever tools ride along on these
# models — and when the API rejects a call with that exact complaint anyway (a future
# generation, an alias we didn't list), retry once at effort none so the user gets a
# working turn instead of a 400.
_EFFORT_ERROR = "function tools with reasoning_effort are not supported"
def _pin_reasoning_effort(kwargs: dict[str, Any]) -> None:
if kwargs.get("tools") and str(kwargs.get("model", "")).startswith("gpt-5.6"):
kwargs.setdefault("reasoning_effort", "none")
def _delta_reasoning(obj: Any) -> Optional[str]:
"""Thinking text off a delta/message: `reasoning_content` (DeepSeek, GLM, Kimi, and
most compat vendors) or `reasoning` (xAI, OpenRouter). Extra response fields survive
the OpenAI SDK's models (extra="allow"), so plain getattr sees them."""
value = getattr(obj, "reasoning_content", None) or getattr(obj, "reasoning", None)
return value if isinstance(value, str) and value else None
def _strip_foreign_sidecars(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Drop provider-private message sidecars (underscore-prefixed keys, e.g. `_gemini`
thought signatures — see providers/base.py): they belong to other providers, and the
OpenAI wire (and its compat servers) rejects unknown message fields."""
return [
(
{k: v for k, v in m.items() if not k.startswith("_")}
if any(k.startswith("_") for k in m)
else m
)
for m in messages
]
_MAX_TOKENS_ERROR = "'max_tokens' is not supported"
# Ceiling, not a spend target — same rationale as the Anthropic provider's default: a
# coworker writing a report ships the whole file inside one tool call's arguments, and
# compat servers left to their OWN defaults cap completions absurdly low (observed
# 2026-08-15: Together defaulted Kimi K3 to ~2k tokens — every ~5KB write truncated).
DEFAULT_MAX_TOKENS = 32000
def _param_fix_retry(kwargs: dict[str, Any], exc: Exception) -> dict[str, Any]:
"""Kwargs for the one retry an unsupported-parameter error earns, or re-raise.
Reasoning-routed OpenAI models reject `max_tokens` outright (they want
`max_completion_tokens`) — but compat servers (Ollama's /v1) know ONLY
`max_tokens`, so the swap must happen on rejection, never up front. Same
contract as the reasoning_effort retry: fix exactly what the server named.
"""
msg = str(exc).lower()
if _EFFORT_ERROR in msg and kwargs.get("reasoning_effort") != "none":
return {**kwargs, "reasoning_effort": "none"}
if _MAX_TOKENS_ERROR in msg and "max_tokens" in kwargs:
fixed = dict(kwargs)
fixed["max_completion_tokens"] = fixed.pop("max_tokens")
return fixed
if "stream_options" in msg and "stream_options" in kwargs:
# Older compat servers don't know the usage opt-in; drop it, lose only metering.
fixed = dict(kwargs)
fixed.pop("stream_options")
return fixed
if ("max_tokens" in msg or "max_new_tokens" in msg) and "max_tokens" in kwargs:
# Our 32k default exceeded this model's completion limit (each server words the
# 400 differently, so no number parsing) — drop the param and retry on the
# server's own default rather than surfacing the 400. Worst case is exactly
# yesterday's behavior; best case the server allows far more once asked.
fixed = dict(kwargs)
fixed.pop("max_tokens")
return fixed
raise exc
def _usage_from(usage: Any) -> Optional[TokenUsage]:
"""chat.completions usage → normalized counts. `prompt_tokens` INCLUDES cached
tokens, so the cached share is subtracted into `cache_read`; no write-side split
exists on this API shape."""
if usage is None:
return None
prompt = int(getattr(usage, "prompt_tokens", 0) or 0)
details = getattr(usage, "prompt_tokens_details", None)
cached = int(getattr(details, "cached_tokens", 0) or 0)
return TokenUsage(
input=max(prompt - cached, 0),
output=int(getattr(usage, "completion_tokens", 0) or 0),
cache_read=cached,
)
class OpenAIProvider(ProviderClient):
def __init__(
self,
client: Any = None,
*,
default_model: str = "gpt-5.6-sol",
api_key: Optional[str] = None,
base_url: Optional[str] = None,
secrets: Any = None,
):
# The SDK client is built lazily on first use, NOT at construction. This lets an engine
# be assembled before any key exists — the desktop app lets you enter the key in Settings
# *after* launch — and the super-agent engine to be built at startup with no key. The key
# is resolved at call time: explicit `api_key` → env `OPENAI_API_KEY` → SecretStore. Tests
# inject a `client` directly, bypassing all of this.
#
# `base_url` points the same OpenAI SDK at any OpenAI-compatible endpoint — used by the
# provider router for Ollama (`http://localhost:11434/v1`, with a placeholder key) and,
# later, other OpenAI-shaped backends. When None, behavior is identical to stock OpenAI.
self._client = client
self._api_key = api_key
self._base_url = base_url
self._secrets = secrets
self.default_model = default_model
def _ensure_client(self) -> Any:
if self._client is None:
# Lazy import so the SDK is only required when actually talking to OpenAI.
from openai import OpenAI
key = self._api_key or resolve_api_key(self._secrets)
if not key:
raise RuntimeError(
"No model API key configured. Set OPENAI_API_KEY in the environment, "
"or add your key in Manage → Settings."
)
kwargs: dict[str, Any] = {"api_key": key}
if self._base_url:
kwargs["base_url"] = self._base_url
self._client = OpenAI(**kwargs)
return self._client
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
) -> AssistantTurn:
kwargs: dict[str, Any] = {
"model": model,
"messages": _strip_foreign_sidecars(messages),
**settings,
}
if tools:
kwargs["tools"] = tools
kwargs.setdefault("max_tokens", DEFAULT_MAX_TOKENS)
_pin_reasoning_effort(kwargs)
client = self._ensure_client()
# Up to three param-fix retries: effort, the max_tokens rename, and the
# max_tokens over-limit drop can ALL need fixing on one call.
for _ in range(3):
try:
response = client.chat.completions.create(**kwargs)
break
except Exception as exc:
kwargs = _param_fix_retry(kwargs, exc)
else:
response = client.chat.completions.create(**kwargs)
choice = response.choices[0]
message = choice.message
text = getattr(message, "content", None)
tool_calls = _parse_tool_calls(getattr(message, "tool_calls", None))
text, tool_calls = _maybe_salvage_tool_calls(text, tool_calls, tools=tools)
return AssistantTurn(
text=text,
tool_calls=tool_calls,
finish_reason=getattr(choice, "finish_reason", None),
raw=response,
reasoning=_delta_reasoning(message),
usage=_usage_from(getattr(response, "usage", None)),
)
def capabilities(self, model: str) -> ModelCapabilities:
return capabilities_for(model)
def stream(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
kwargs: dict[str, Any] = {
"model": model,
"messages": _strip_foreign_sidecars(messages),
"stream": True,
# Usage on the final chunk (empty `choices`). Compat servers that reject
# the option get a one-shot retry without it (_param_fix_retry).
"stream_options": {"include_usage": True},
**settings,
}
if tools:
kwargs["tools"] = tools
kwargs.setdefault("max_tokens", DEFAULT_MAX_TOKENS)
_pin_reasoning_effort(kwargs)
client = self._ensure_client()
text_parts: list[str] = []
reasoning_parts: list[str] = []
tool_accum: dict[int, dict[str, str]] = {}
finish_reason = None
usage: Optional[TokenUsage] = None
# Up to three param-fix retries: effort, the max_tokens rename, and the
# max_tokens over-limit drop can ALL need fixing on one call.
for _ in range(3):
try:
chunks = client.chat.completions.create(**kwargs)
break
except Exception as exc:
kwargs = _param_fix_retry(kwargs, exc)
else:
chunks = client.chat.completions.create(**kwargs)
for chunk in chunks:
chunk_usage = _usage_from(getattr(chunk, "usage", None))
if chunk_usage is not None:
usage = chunk_usage
choices = getattr(chunk, "choices", None)
if not choices:
continue
choice = choices[0]
delta = getattr(choice, "delta", None)
if delta is not None:
reasoning = _delta_reasoning(delta)
if reasoning:
reasoning_parts.append(reasoning)
yield StreamChunk(reasoning_delta=reasoning)
content = getattr(delta, "content", None)
if content:
text_parts.append(content)
yield StreamChunk(text_delta=content)
for tc in getattr(delta, "tool_calls", None) or []:
acc = tool_accum.setdefault(
getattr(tc, "index", 0), {"id": "", "name": "", "args": ""}
)
if getattr(tc, "id", None):
acc["id"] = tc.id
fn = getattr(tc, "function", None)
if fn is not None:
if getattr(fn, "name", None):
acc["name"] = fn.name
if getattr(fn, "arguments", None):
acc["args"] += fn.arguments
if getattr(choice, "finish_reason", None):
finish_reason = choice.finish_reason
tool_calls = []
for index in sorted(tool_accum):
acc = tool_accum[index]
try:
arguments = json.loads(acc["args"]) if acc["args"] else {}
except (TypeError, json.JSONDecodeError):
arguments = {"_raw": acc["args"]}
tool_calls.append(
ToolCall(id=acc["id"], name=acc["name"], arguments=arguments)
)
text, tool_calls = _maybe_salvage_tool_calls(
"".join(text_parts) or None, tool_calls, tools=tools
)
yield StreamChunk(
turn=AssistantTurn(
text=text,
tool_calls=tool_calls,
finish_reason=finish_reason,
reasoning="".join(reasoning_parts) or None,
usage=usage,
)
)
def _parse_tool_calls(raw_tool_calls: Any) -> list[ToolCall]:
calls: list[ToolCall] = []
for tc in raw_tool_calls or []:
function = tc.function
raw_args = getattr(function, "arguments", None)
try:
arguments = json.loads(raw_args) if raw_args else {}
except (TypeError, json.JSONDecodeError):
# Surface unparseable arguments rather than dropping the call; the engine
# can return a tool-error so the model corrects itself.
arguments = {"_raw": raw_args}
calls.append(
ToolCall(id=getattr(tc, "id", ""), name=function.name, arguments=arguments)
)
return calls
# Some OpenAI-compatible backends — notably Ollama for several local models (qwen, etc.) —
# fail to populate the structured `tool_calls` field and instead emit the call as TEXT, in
# wildly varied shapes: a `<tool_call>{…}</tool_call>` block, a bare `{"name","arguments"}` object
# (often mixed in with prose), or a `toolname {args}` / `toolname [args]` shorthand. Our agent
# loop needs structured calls, so we recover them — using the requested tool SCHEMAS to recognize
# tool-name forms and to filter out anything whose name isn't a real tool (no false positives).
# Gated on: tools were requested AND no structured calls came back. Never fires for OpenAI.
_TOOLCALL_OPEN = re.compile(r"<tool_call>\s*", re.IGNORECASE)
# Qwen/Hermes native tool-call template — NOT JSON. The model writes the call as nested XML:
# <function=write_file><parameter=path>hello.txt</parameter><parameter=content>hi</parameter></function>
# (usually wrapped in <tool_call>…</tool_call>). qwen3-coder emits exactly this, so we parse the
# function/parameter tags directly. Values are taken verbatim (stripped); only no-whitespace JSON
# tokens (numbers, bools, objects/arrays) are coerced, so free-text content stays a string.
_FUNCTION_BLOCK = re.compile(
r"<function\s*=\s*(?P<name>[^>\s]+)\s*>(?P<body>.*?)</function\s*>",
re.IGNORECASE | re.DOTALL,
)
_PARAM_BLOCK = re.compile(
r"<parameter\s*=\s*(?P<key>[^>\s]+)\s*>(?P<val>.*?)</parameter\s*>",
re.IGNORECASE | re.DOTALL,
)
# A `<function=NAME>` that never closes — the model ran out of tokens (or drifted) partway
# through writing the call. Anchored to end-of-text so it only matches a genuinely unfinished
# tail, never a well-formed block earlier in the message. Small local models hit this often on
# a large tool schema, and the turn used to end silently on the leftover text.
_FUNCTION_OPEN_TRUNCATED = re.compile(
r"<function\s*=\s*(?P<name>[^>\s]+)\s*>(?P<body>(?:(?!</function\s*>).)*)$",
re.IGNORECASE | re.DOTALL,
)
# Markers that mean "this text IS a tool call the endpoint failed to parse", used to tell a
# real answer from a leaked one. Fenced code is stripped first: a model *explaining* tool-call
# syntax in a ``` block is answering, not calling.
_LEAKED_TOOL_SYNTAX = (
"<tool_call>",
"</tool_call>",
"<function=",
"</function>",
"<parameter=",
"</parameter>",
"<function_calls>",
"<invoke ",
)
_FENCED = re.compile(r"```.*?```|~~~.*?~~~|`[^`\n]*`", re.DOTALL)
def looks_like_unparsed_tool_call(
text: Optional[str], tools: Optional[list[dict[str, Any]]] = None
) -> bool:
"""True when assistant text still carries tool-call markup that salvage couldn't turn into
a call — i.e. the model tried to call a tool and the syntax was mangled or cut off.
Only meaningful when tools were actually offered, and only over OpenAI-compatible endpoints
that parse tool calls out of the model's raw output (LM Studio, Ollama, vLLM). The caller
uses it to end the turn as a retriable error instead of presenting the fragment as an answer.
"""
if not tools or not text:
return False
return any(m in _FENCED.sub("", text).lower() for m in _LEAKED_TOOL_SYNTAX)
def _coerce_param(raw: str) -> Any:
"""Keep free-text verbatim (the common case: file content), but recover real JSON values when
the whole token is unambiguous JSON (no embedded whitespace) — e.g. `3`, `true`, `{"a":1}`.
"""
s = raw.strip()
if s and not any(c.isspace() for c in s):
v = _loads(s)
if isinstance(v, (dict, list, int, float, bool)):
return v
return s
def _maybe_salvage_tool_calls(
text: Optional[str],
tool_calls: list[ToolCall],
*,
tools: Optional[list[dict[str, Any]]],
) -> tuple[Optional[str], list[ToolCall]]:
"""If the model returned tool calls as text, convert them. Returns (text, tool_calls):
on success the salvaged calls replace `tool_calls` and `text` is cleared."""
if tool_calls or not tools or not text:
return text, tool_calls
salvaged = _salvage_tool_calls_from_text(text, tools)
if salvaged:
return None, salvaged
return text, tool_calls
def _tool_index(
tools: Optional[list[dict[str, Any]]],
) -> tuple[Optional[set[str]], dict[str, Optional[str]]]:
"""(known tool names, {name: sole-parameter-name}) from OpenAI tool schemas. The sole-param
map lets us map a bare `toolname [args]` to `{param: args}` when a tool has one parameter.
"""
if not tools:
return None, {}
names: set[str] = set()
single: dict[str, Optional[str]] = {}
for t in tools:
fn = (t or {}).get("function") or {}
name = fn.get("name")
if not isinstance(name, str) or not name:
continue
names.add(name)
params = fn.get("parameters") or {}
props = params.get("properties") or {}
if len(props) == 1:
single[name] = next(iter(props))
else:
required = params.get("required") or []
single[name] = required[0] if len(required) == 1 else None
return names, single
def _loads(s: str) -> Any:
try:
return json.loads(s)
except (TypeError, json.JSONDecodeError):
return None
def _extract_balanced(text: str, start: int) -> Optional[str]:
"""Return the balanced `{…}`/`[…]` substring beginning at `text[start]` (string-aware), or
None if it doesn't close — so nested braces/brackets are handled correctly."""
open_ch = text[start]
close_ch = "]" if open_ch == "[" else "}"
depth = 0
in_str = False
esc = False
for i in range(start, len(text)):
ch = text[i]
if in_str:
if esc:
esc = False
elif ch == "\\":
esc = True
elif ch == '"':
in_str = False
elif ch == '"':
in_str = True
elif ch == open_ch:
depth += 1
elif ch == close_ch:
depth -= 1
if depth == 0:
return text[start : i + 1]
return None
def _iter_top_objects(text: str):
"""Yield balanced `{…}` substrings at brace-depth 0 (array brackets ignored), so embedded
JSON objects are found even amid surrounding prose."""
i = 0
while i < len(text):
if text[i] == "{":
sub = _extract_balanced(text, i)
if sub:
yield sub
i += len(sub)
continue
i += 1
def _call_from_dict(d: Any, names: Optional[set[str]]) -> Optional[ToolCall]:
"""Build a ToolCall from a `{"name","arguments"}` dict, or None if it isn't one / the name
isn't a known tool."""
if not isinstance(d, dict):
return None
name = d.get("name")
if not isinstance(name, str) or not name:
return None
if names is not None and name not in names:
return None
args = d.get("arguments", d.get("parameters"))
if args is None:
args = {}
if isinstance(args, str):
args = _loads(args)
if not isinstance(args, dict):
args = {"_raw": d.get("arguments")}
if not isinstance(args, dict):
args = {"_raw": args}
return ToolCall(id="", name=name, arguments=args)
def _renumber(calls: list[ToolCall]) -> list[ToolCall]:
return [
ToolCall(id=f"call_salvaged_{i}", name=c.name, arguments=c.arguments)
for i, c in enumerate(calls)
]
def _salvage_tool_calls_from_text(
content: str, tools: Optional[list[dict[str, Any]]] = None
) -> list[ToolCall]:
"""Best-effort recovery of tool calls embedded in assistant text. Tries, in order:
1. `<tool_call>…</tool_call>` blocks (anywhere, balanced); 2. embedded `{"name","arguments"}`
objects (even mixed with prose); 3. `toolname {args}` / `toolname [args]` for known tools.
Returns [] (treat as plain text) when nothing tool-shaped is found."""
text = (content or "").strip()
if not text:
return []
names, single = _tool_index(tools)
# 1) <tool_call> … </tool_call> blocks.
calls: list[ToolCall] = []
for m in _TOOLCALL_OPEN.finditer(text):
j = m.end()
if j < len(text) and text[j] in "{[":
sub = _extract_balanced(text, j)
parsed = _loads(sub) if sub else None
for d in parsed if isinstance(parsed, list) else [parsed]:
c = _call_from_dict(d, names)
if c:
calls.append(c)
if calls:
return _renumber(calls)
# 1b) Qwen/Hermes XML calls: <function=NAME><parameter=KEY>VAL</parameter>…</function>.
for fm in _FUNCTION_BLOCK.finditer(text):
name = fm.group("name").strip()
if names is not None and name not in names:
continue
args = {
pm.group("key").strip(): _coerce_param(pm.group("val"))
for pm in _PARAM_BLOCK.finditer(fm.group("body"))
}
calls.append(ToolCall(id="", name=name, arguments=args))
if calls:
return _renumber(calls)
# 1c) A TRUNCATED XML call: `<function=NAME>` with no closing tag, because the model ran
# out of tokens mid-call. Take the name plus every parameter that DID close; a trailing
# unterminated `<parameter=…>` is dropped rather than guessed, so a half-written path or
# file body can never reach a tool. If that leaves a required argument missing the call
# fails validation and the model gets a corrective tool error — which is the agent loop
# working, and strictly better than the turn ending on the leftover fragment.
tm = _FUNCTION_OPEN_TRUNCATED.search(text)
if tm:
name = tm.group("name").strip()
if names is None or name in names:
args = {
pm.group("key").strip(): _coerce_param(pm.group("val"))
for pm in _PARAM_BLOCK.finditer(tm.group("body"))
}
return _renumber([ToolCall(id="", name=name, arguments=args)])
# 2) Embedded {"name": …, "arguments": …} objects, even surrounded by prose.
for sub in _iter_top_objects(text):
d = _loads(sub)
if isinstance(d, dict) and "name" in d:
c = _call_from_dict(d, names)
if c:
calls.append(c)
if calls:
return _renumber(calls)
# 3) `toolname {args}` / `toolname [args]` shorthand — only for tools we actually offered.
if names:
for name in names:
for m in re.finditer(re.escape(name) + r"\s*[:=]?\s*", text):
j = m.end()
if j >= len(text) or text[j] not in "{[":
continue
sub = _extract_balanced(text, j)
parsed = _loads(sub) if sub else None
if parsed is None:
continue
if isinstance(parsed, dict):
args = parsed
else:
param = single.get(name)
if not param:
continue
args = {param: parsed}
calls.append(ToolCall(id="", name=name, arguments=args))
break # one salvaged call per tool name
return _renumber(calls)

View File

@@ -0,0 +1,462 @@
"""OpenAI Responses provider — native and compatible models via `/responses`.
Chat Completions rejects function tools combined with any `reasoning_effort` other than
`none` on GPT-5.6+ ("use /v1/responses"), which had reasoning pinned OFF for native OpenAI
models (see `openai_provider._pin_reasoning_effort`). This provider is the Responses path:
reasoning + tools at real effort levels, streamed reasoning summaries (→ the same
`reasoning_delta` / `AssistantTurn.reasoning` plumbing the GUI already renders), and
chain-of-thought continuity across tool round-trips via `store: false` +
`include: ["reasoning.encrypted_content"]` — nothing retained server-side.
Routing: the `openai` provider entry with NO custom base_url builds this class. Most custom
endpoints (Azure, vLLM, and the existing compat vendors) keep the Chat Completions
`OpenAIProvider`; vendors that explicitly implement the Responses wire can opt into this
class with their own base URL (registry.py).
Like the other native providers, this is mostly a pair of pure converters from the
canonical OpenAI-chat-shaped history to Responses `input` items. What the converters
must absorb:
- The system prompt is the `instructions` request field, not a message role.
- Assistant tool calls are top-level `function_call` items; tool results are
`function_call_output` items paired by `call_id` (ids only need to pair up, so foreign
`toolu_…` ids from a mid-conversation provider switch are fine).
- Tool schemas are FLAT (`{"type": "function", "name", …}` — no nested `function` key).
- Reasoning continuity: the raw output items (reasoning item with `encrypted_content`,
`function_call` items with their ids) ride the canonical assistant message as the
`_openai` sidecar (see providers/base.py). Present → replayed verbatim for exact CoT
continuity; absent (history from another provider) → items are synthesized from the
canonical fields. Reasoning items WITHOUT `encrypted_content` never enter the sidecar:
with `store: false` the server can't resolve them and would reject the replay.
"""
from __future__ import annotations
import json
import re
from typing import Any, Optional
from .base import (
AssistantTurn,
ModelCapabilities,
ProviderClient,
StreamChunk,
TokenUsage,
ToolCall,
)
from .capabilities import capabilities_for
from .openai_provider import resolve_api_key
# Request params passed through from model settings; everything else (frequency_penalty,
# reasoning_effort — no effort knob in v1, the server default rides) is dropped.
_SETTINGS_WHITELIST = {
"temperature",
"top_p",
"max_output_tokens",
"tool_choice",
"parallel_tool_calls",
}
# "Unsupported parameter: 'temperature' is not supported with this model." — reasoning
# models reject sampling params; non-reasoning models reject `reasoning`/`include`. The
# server names exactly one offender per error, so each retry drops exactly that.
_UNSUPPORTED_PARAM = re.compile(r"unsupported (?:parameter|value)s?:?\s*'([^']+)'")
def _param_fix_retry(kwargs: dict[str, Any], exc: Exception) -> dict[str, Any]:
"""Kwargs for the one retry an unsupported-parameter error earns, or re-raise.
Same contract as the Chat Completions retries: fix exactly what the server named.
A dotted name (`reasoning.summary`) drops its top-level param.
"""
match = _UNSUPPORTED_PARAM.search(str(exc).lower())
if match:
param = match.group(1).split(".", 1)[0].split("[", 1)[0]
if param in kwargs and param not in ("model", "input"):
fixed = dict(kwargs)
del fixed[param]
return fixed
raise exc
def _user_content(content: Any) -> Any:
"""User content (str or OpenAI chat parts) → Responses content (str or input parts)."""
if isinstance(content, str):
return content
parts: list[dict[str, Any]] = []
for part in content or []:
kind = part.get("type") if isinstance(part, dict) else None
if kind == "text":
parts.append({"type": "input_text", "text": part.get("text") or ""})
elif kind == "image_url":
url = (part.get("image_url") or {}).get("url") or ""
parts.append({"type": "input_image", "image_url": url})
elif kind == "file":
file = part.get("file") or {}
entry: dict[str, Any] = {"type": "input_file"}
if file.get("filename"):
entry["filename"] = file["filename"]
if file.get("file_data"):
entry["file_data"] = file["file_data"]
parts.append(entry)
return parts
def _synthesized_items(message: dict[str, Any]) -> list[dict[str, Any]]:
"""An assistant message WITHOUT a usable `_openai` sidecar (history produced by another
provider before a switch) → items rebuilt from the canonical fields."""
items: list[dict[str, Any]] = []
text = message.get("content")
if isinstance(text, str) and text:
items.append({"role": "assistant", "content": text})
for call in message.get("tool_calls") or []:
function = call.get("function") or {}
arguments = function.get("arguments")
if not isinstance(arguments, str):
arguments = json.dumps(arguments or {})
items.append(
{
"type": "function_call",
"call_id": call.get("id") or "",
"name": function.get("name") or "",
"arguments": arguments,
}
)
return items
def convert_messages(
messages: list[dict[str, Any]],
) -> tuple[Optional[str], list[dict[str, Any]]]:
"""Canonical OpenAI-chat history → (`instructions`, Responses `input` items).
Leading system messages join into `instructions`; a stray mid-thread system message
rides as a system message item. Assistant messages replay their `_openai` sidecar
verbatim when present (exact CoT continuity), else synthesize from canonical fields.
"""
system_parts: list[str] = []
index = 0
while index < len(messages) and messages[index].get("role") == "system":
content = messages[index].get("content")
if isinstance(content, str) and content:
system_parts.append(content)
index += 1
items: list[dict[str, Any]] = []
for message in messages[index:]:
role = message.get("role")
if role == "system":
text = message.get("content") or ""
if text:
items.append({"role": "system", "content": text})
elif role == "user":
content = _user_content(message.get("content"))
if content:
items.append({"role": "user", "content": content})
elif role == "assistant":
sidecar = message.get("_openai") or {}
replay = sidecar.get("items") or []
if replay:
items.extend(replay)
else:
items.extend(_synthesized_items(message))
elif role == "tool":
content = message.get("content")
items.append(
{
"type": "function_call_output",
"call_id": message.get("tool_call_id") or "",
"output": content if isinstance(content, str) else str(content or ""),
}
)
return ("\n\n".join(system_parts) or None), items
def convert_tools(tools: Optional[list[dict[str, Any]]]) -> list[dict[str, Any]]:
"""OpenAI chat function schemas → Responses FLAT tool entries (no nested `function`)."""
converted: list[dict[str, Any]] = []
for tool in tools or []:
function = (tool or {}).get("function") or {}
name = function.get("name")
if not name:
continue
entry: dict[str, Any] = {"type": "function", "name": name}
if function.get("description"):
entry["description"] = function["description"]
if function.get("parameters") is not None:
entry["parameters"] = function["parameters"]
converted.append(entry)
return converted
def _dump(value: Any) -> Any:
"""An output item (SDK model, dict, or test namespace) → plain jsonl-safe data."""
if isinstance(value, dict):
return {k: _dump(v) for k, v in value.items() if v is not None}
if isinstance(value, (list, tuple)):
return [_dump(v) for v in value]
dump = getattr(value, "model_dump", None)
if callable(dump):
return dump(exclude_none=True)
if hasattr(value, "__dict__"): # SimpleNamespace fakes in tests
return {k: _dump(v) for k, v in vars(value).items() if v is not None}
return value
def _parse_arguments(raw: Any) -> dict[str, Any]:
if isinstance(raw, dict):
return raw
if not raw:
return {}
try:
parsed = json.loads(raw)
return parsed if isinstance(parsed, dict) else {"_raw": raw}
except (TypeError, json.JSONDecodeError):
# Surface unparseable arguments rather than dropping the call; the engine
# can return a tool-error so the model corrects itself.
return {"_raw": raw}
def _sidecar_extras(items: list[dict[str, Any]]) -> dict[str, Any]:
"""Output items → the `_openai` sidecar, or {} when replay would add nothing.
Reasoning items without `encrypted_content` are dropped: under `store: false` the
server can't resolve them by id and rejects the replay. The sidecar is only worth
persisting when something beyond plain answer text needs continuity.
"""
kept = [
item
for item in items
if item.get("type") != "reasoning" or item.get("encrypted_content")
]
if any(item.get("type") in ("reasoning", "function_call") for item in kept):
return {"_openai": {"items": kept}}
return {}
def _usage_from(usage: Any) -> Optional[TokenUsage]:
"""Responses-API usage → normalized counts (OPE-101). `input_tokens` INCLUDES the
cached share, so fresh input = input_tokens cached_tokens (the same convention as
the Chat Completions and Anthropic adapters); `output_tokens` already includes
reasoning tokens (billed as output). Defensive reads throughout — compat/older
servers may omit `input_tokens_details`."""
if usage is None:
return None
prompt = int(getattr(usage, "input_tokens", 0) or 0)
details = getattr(usage, "input_tokens_details", None)
cached = int(getattr(details, "cached_tokens", 0) or 0)
return TokenUsage(
input=max(prompt - cached, 0),
output=int(getattr(usage, "output_tokens", 0) or 0),
cache_read=cached,
)
def _parse_response(response: Any) -> AssistantTurn:
"""One Responses result → an AssistantTurn (+ `_openai` extras)."""
items = [_dump(item) for item in getattr(response, "output", None) or []]
texts: list[str] = []
summaries: list[str] = []
tool_calls: list[ToolCall] = []
for item in items:
kind = item.get("type")
if kind == "message" or (kind is None and "content" in item):
content = item.get("content")
if isinstance(content, str):
texts.append(content)
else:
for part in content or []:
if part.get("type") == "output_text" and part.get("text"):
texts.append(part["text"])
elif kind == "reasoning":
for part in item.get("summary") or []:
text = part.get("text") if isinstance(part, dict) else part
if text:
summaries.append(text)
elif kind == "function_call":
tool_calls.append(
ToolCall(
id=item.get("call_id") or item.get("id") or "",
name=item.get("name") or "",
arguments=_parse_arguments(item.get("arguments")),
)
)
incomplete = _dump(getattr(response, "incomplete_details", None)) or {}
if tool_calls:
finish = "tool_calls"
elif incomplete.get("reason") == "max_output_tokens":
finish = "length"
else:
finish = "stop"
return AssistantTurn(
text="".join(texts) or None,
tool_calls=tool_calls,
finish_reason=finish,
raw=response,
reasoning="".join(summaries) or None,
extras=_sidecar_extras(items),
usage=_usage_from(getattr(response, "usage", None)),
)
class OpenAIResponsesProvider(ProviderClient):
def __init__(
self,
client: Any = None,
*,
default_model: str = "gpt-5.6-sol",
api_key: Optional[str] = None,
secrets: Any = None,
base_url: Optional[str] = None,
reasoning_summary: bool = True,
):
# Same deferred-client contract as OpenAIProvider: built lazily so an engine can be
# assembled before any key exists; key resolves at call time (explicit → env →
# SecretStore). Tests inject a `client` directly. `base_url` is opt-in: stock OpenAI
# leaves it unset, while Responses-compatible vendors can supply their own endpoint.
self._client = client
self._api_key = api_key
self._secrets = secrets
self._base_url = (base_url or "").strip().rstrip("/") or None
if not isinstance(reasoning_summary, bool):
raise TypeError("reasoning_summary must be a bool")
self._reasoning_summary = reasoning_summary
self.default_model = default_model
def _ensure_client(self) -> Any:
if self._client is None:
# Lazy import so the SDK is only required when actually talking to OpenAI.
from openai import OpenAI
key = self._api_key or resolve_api_key(self._secrets)
if not key:
raise RuntimeError(
"No model API key configured. Set OPENAI_API_KEY in the environment, "
"or add your key in Manage → Settings."
)
kwargs = {"api_key": key}
if self._base_url:
kwargs["base_url"] = self._base_url
self._client = OpenAI(**kwargs)
return self._client
def _request_kwargs(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]],
settings: dict[str, Any],
) -> dict[str, Any]:
instructions, items = convert_messages(messages)
if "max_tokens" in settings and "max_output_tokens" not in settings:
settings = {**settings, "max_output_tokens": settings["max_tokens"]}
kwargs: dict[str, Any] = {
"model": model,
"input": items,
# Stateless: nothing retained server-side; the encrypted reasoning rides the
# `_openai` sidecar instead, and summaries feed the GUI's thinking display.
"store": False,
"include": ["reasoning.encrypted_content"],
**{k: v for k, v in settings.items() if k in _SETTINGS_WHITELIST},
}
if self._reasoning_summary:
kwargs["reasoning"] = {"summary": "auto"}
if instructions:
kwargs["instructions"] = instructions
if tools:
converted = convert_tools(tools)
if converted:
kwargs["tools"] = converted
return kwargs
def _create(self, client: Any, kwargs: dict[str, Any]) -> Any:
# Up to three param-fix retries: sampling params, `reasoning`, and `include` can
# each need dropping depending on the model (reasoning vs not).
for _ in range(3):
try:
return client.responses.create(**kwargs)
except Exception as exc:
kwargs = _param_fix_retry(kwargs, exc)
return client.responses.create(**kwargs)
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
) -> AssistantTurn:
kwargs = self._request_kwargs(
model=model, messages=messages, tools=tools, settings=settings
)
response = self._create(self._ensure_client(), kwargs)
return _parse_response(response)
def capabilities(self, model: str) -> ModelCapabilities:
return capabilities_for(model)
def stream(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
kwargs = self._request_kwargs(
model=model, messages=messages, tools=tools, settings=settings
)
kwargs["stream"] = True
events = self._create(self._ensure_client(), kwargs)
text_parts: list[str] = []
reasoning_parts: list[str] = []
done_items: list[Any] = []
final: Optional[Any] = None
for event in events:
kind = getattr(event, "type", None)
if kind == "response.output_text.delta":
delta = getattr(event, "delta", None)
if delta:
text_parts.append(delta)
yield StreamChunk(text_delta=delta)
elif kind == "response.reasoning_summary_text.delta":
delta = getattr(event, "delta", None)
if delta:
reasoning_parts.append(delta)
yield StreamChunk(reasoning_delta=delta)
elif kind == "response.output_item.done":
item = getattr(event, "item", None)
if item is not None:
done_items.append(item)
elif kind in ("response.completed", "response.incomplete", "response.failed"):
final = getattr(event, "response", None)
if final is not None:
# The terminal event carries the full response — parse it whole so tool
# calls, finish reason, and the `_openai` sidecar come from one place.
# Some Responses-compatible backends (the subscription backend) leave the
# terminal response's `output` EMPTY — the items only ever stream — so
# graft the streamed output_item.done items back on before parsing, or a
# turn's text and tool calls silently vanish.
if not (getattr(final, "output", None) or []) and done_items:
try:
final.output = done_items
except Exception:
pass
turn = _parse_response(final)
if turn.text is None and not turn.tool_calls and text_parts:
turn.text = "".join(text_parts)
yield StreamChunk(turn=turn)
else:
yield StreamChunk(
turn=AssistantTurn(
text="".join(text_parts) or None,
reasoning="".join(reasoning_parts) or None,
)
)

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,118 @@
"""ProviderRouter — one `ProviderClient` that dispatches by the `provider:` prefix of a model
string to a per-provider client, built lazily from its SecretStore profile and cached.
This is the single provider the `SessionManager` hands to every engine, so `complete()/stream()`
(which already receive the full model string per-call) route themselves: `ollama:llama3.3` →
the Ollama client (Ollama's OpenAI-compatible `/v1`), bare `gpt-5.5` → the default (OpenAI). The
prefix is stripped before delegating, since the underlying SDKs want the bare model name.
Config changes (a new key, a new Ollama URL) call `invalidate()` to drop cached clients, so
existing engines pick up the change without a rebuild.
"""
from __future__ import annotations
import threading
from typing import Any, Optional
from .base import ProviderClient
from .capabilities import capabilities_for
from .registry import build_provider_client, get_descriptor
class ProviderRouter(ProviderClient):
def __init__(
self,
secrets: Any = None,
*,
default_provider: str = "openai",
on_use: Any = None,
) -> None:
self._secrets = secrets
self._default = default_provider
self._clients: dict[str, ProviderClient] = {}
self._lock = threading.Lock()
# Optional callable(provider_name) fired when a completion is dispatched — drives the
# Settings pane's "Last used" line. Best-effort: its failures never break a model call.
self._on_use = on_use
def _note_use(self, model: str) -> None:
if self._on_use is None:
return
try:
self._on_use(self._provider_name(model))
except Exception:
pass
# -- routing ----------------------------------------------------------------
def _provider_name(self, model: str) -> str:
"""The provider for a model: the `prefix` of `prefix:rest` if it's a known provider,
else the default. (A colon that isn't a known provider — unlikely — falls through.)
"""
if ":" in model:
prefix = model.split(":", 1)[0]
if get_descriptor(prefix) is not None:
return prefix
return self._default
def _client_for(self, model: str) -> ProviderClient:
name = self._provider_name(model)
with self._lock:
client = self._clients.get(name)
if client is None:
profile = {}
if self._secrets is not None:
profile = self._secrets.get(f"provider:{name}") or {}
client = build_provider_client(name, profile, self._secrets)
self._clients[name] = client
return client
@staticmethod
def _bare(model: str) -> str:
"""Strip a KNOWN provider prefix; the underlying SDK wants the bare model name. A model
whose first segment isn't a provider (e.g. `qwen2.5-coder:32b` — a version tag, not a
prefix) is returned unchanged, so the colon isn't mistaken for a provider separator.
"""
if ":" in model:
prefix, rest = model.split(":", 1)
if get_descriptor(prefix) is not None:
return rest
return model
def invalidate(self, name: Optional[str] = None) -> None:
"""Drop cached client(s) so the next call rebuilds with fresh config."""
with self._lock:
if name is None:
self._clients.clear()
else:
self._clients.pop(name, None)
# -- ProviderClient ---------------------------------------------------------
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
self._note_use(model)
return self._client_for(model).complete(
model=self._bare(model), messages=messages, tools=tools, **settings
)
def stream(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
self._note_use(model)
return self._client_for(model).stream(
model=self._bare(model), messages=messages, tools=tools, **settings
)
def capabilities(self, model: str):
return capabilities_for(model)

View File

@@ -0,0 +1,238 @@
"""Google Vertex AI provider — one entry in Settings, three wire paths by model family.
Routed ids look like `vertex:<family>/<model id>`; the router strips `vertex:` and this
provider splits the family segment, reusing an existing provider class per family:
- `gemini/…` → the native `GeminiProvider` over `genai.Client(vertexai=True)`.
- `claude/…` → the native `AnthropicProvider` over the SDK's `AnthropicVertex` client.
- `openweight/…` → `OpenAIProvider` against Vertex's OpenAI-compatible MaaS endpoint
(Llama, Qwen, DeepSeek, …; ids keep their publisher segment, e.g. `openweight/meta/…`).
An id with no recognized family segment is best-effort routed by name (gemini* → Gemini,
claude* → Claude, anything else → MaaS) so a raw id pasted without the add-model dropdown
still works.
Auth is ONE method at a time, selected by the profile's `auth_method` (a segmented choice
in Settings, mirroring Bedrock — owner call 2026-07-26):
- `adc` — Application Default Credentials (`gcloud auth application-default
login`), Google's own recommended path. Nothing stored.
- `service_account` — an explicit service-account JSON (pasted content or a file path).
- `api_key` — a Vertex API key (express mode). GEMINI FAMILY ONLY: the genai SDK
takes it (and it excludes project/location — mutually exclusive there), but Claude
(AnthropicVertex) and the MaaS endpoint require OAuth credentials, so those families
raise a clear error directing the user to the other methods.
The MaaS path authenticates with a google-auth bearer token that expires ~hourly — this
wrapper refreshes it and rebuilds the OpenAI sub-client as needed; the two native SDK
clients take the credentials object and refresh internally. Fields from non-selected
methods are dropped at construction; a missing/unknown method falls back to whichever
fields are present (service account, else ADC).
"""
from __future__ import annotations
import json
from typing import Any, Optional
from .anthropic_provider import AnthropicProvider
from .base import AssistantTurn, ModelCapabilities, ProviderClient
from .capabilities import capabilities_for
from .gemini_provider import GeminiProvider
from .openai_provider import OpenAIProvider
_SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
_FAMILIES = ("gemini", "claude", "openweight")
def _regional_host(location: Optional[str]) -> str:
"""Vertex REST host for a location — `global` (newer Gemini models) has no region
prefix (checked live 2026-07-26)."""
if not location or location == "global":
return "aiplatform.googleapis.com"
return f"{location}-aiplatform.googleapis.com"
def load_credentials(service_account_json: Optional[str]) -> Any:
"""Explicit service-account JSON (content or path) → Credentials; blank → None (the
SDKs and the token path then fall back to Application Default Credentials)."""
raw = (service_account_json or "").strip()
if not raw:
return None
from google.oauth2 import service_account
if raw.startswith("{"):
info = json.loads(raw)
return service_account.Credentials.from_service_account_info(
info, scopes=_SCOPES
)
return service_account.Credentials.from_service_account_file(raw, scopes=_SCOPES)
class VertexProvider(ProviderClient):
"""Family dispatcher: splits `<family>/<model id>` and delegates to the sub-client."""
def __init__(
self,
*,
project: Optional[str] = None,
location: Optional[str] = None,
auth_method: Optional[str] = None,
service_account_json: Optional[str] = None,
api_key: Optional[str] = None,
credentials: Any = None,
gemini_client: Optional[ProviderClient] = None,
claude_client: Optional[ProviderClient] = None,
openweight_client: Optional[ProviderClient] = None,
):
# Narrow to the selected auth method here, once — stale values stored under a
# previously-selected method must never reach a different credential path.
if auth_method == "adc":
service_account_json = api_key = None
elif auth_method == "service_account":
api_key = None
elif auth_method == "api_key":
service_account_json = None
self._project = project
self._location = location
self._api_key = api_key
self._service_account_json = service_account_json
self._credentials = credentials # test seam; normally resolved lazily
# Test seams: pre-built sub-providers skip the SDK construction below.
self._clients: dict[str, ProviderClient] = {}
if gemini_client is not None:
self._clients["gemini"] = gemini_client
if claude_client is not None:
self._clients["claude"] = claude_client
if openweight_client is not None:
self._clients["openweight"] = openweight_client
self._openweight_injected = openweight_client is not None
@staticmethod
def _split(model: str) -> tuple[str, str]:
if "/" in model:
family, rest = model.split("/", 1)
if family in _FAMILIES:
return family, rest
# Raw id without a family segment: route by name, best effort.
if model.startswith("gemini"):
return "gemini", model
if model.startswith("claude"):
return "claude", model
return "openweight", model
# -- credentials -------------------------------------------------------------
def _explicit_credentials(self) -> Any:
"""The service-account credentials, or None to let each SDK use ADC."""
if self._credentials is None:
self._credentials = load_credentials(self._service_account_json)
return self._credentials
def _bearer_credentials(self) -> Any:
"""Credentials for the MaaS bearer token: explicit service account, else ADC."""
creds = self._explicit_credentials()
if creds is None:
import google.auth
try:
creds, _ = google.auth.default(scopes=_SCOPES)
except Exception as exc:
raise RuntimeError(
"No Google Cloud credentials found — paste a service-account JSON "
"in Settings ▸ Models, or run `gcloud auth application-default login`."
) from exc
self._credentials = creds
return creds
# -- family sub-clients --------------------------------------------------------
def _family_client(self, family: str) -> ProviderClient:
if self._api_key and family != "gemini":
raise RuntimeError(
"Vertex API keys cover Gemini models only — switch the Vertex provider "
"to Google Cloud login or a service account for Claude and open-weight "
"models (Settings ▸ Models)."
)
if family == "openweight":
return self._openweight_client()
client = self._clients.get(family)
if client is None:
if family == "gemini":
from google import genai
if self._api_key:
# Express mode: the key excludes project/location (SDK enforces
# mutual exclusivity — the key already identifies the project).
sdk = genai.Client(vertexai=True, api_key=self._api_key)
else:
sdk = genai.Client(
vertexai=True,
project=self._project,
location=self._location,
credentials=self._explicit_credentials(),
)
client = GeminiProvider(client=sdk)
else:
from anthropic import AnthropicVertex
client = AnthropicProvider(
client=AnthropicVertex(
project_id=self._project,
region=self._location,
credentials=self._explicit_credentials(),
)
)
self._clients[family] = client
return client
def _openweight_client(self) -> ProviderClient:
"""OpenAIProvider over the Vertex MaaS endpoint, rebuilt whenever the bearer
token has to be refreshed (google-auth tokens expire ~hourly)."""
if self._openweight_injected:
return self._clients["openweight"]
creds = self._bearer_credentials()
if not getattr(creds, "valid", False):
from google.auth.transport.requests import Request
creds.refresh(Request())
self._clients.pop("openweight", None) # stale token — rebuild below
client = self._clients.get("openweight")
if client is None:
base = (
f"https://{_regional_host(self._location)}/v1/projects/"
f"{self._project}/locations/{self._location}/endpoints/openapi"
)
client = OpenAIProvider(api_key=creds.token, base_url=base)
self._clients["openweight"] = client
return client
# -- ProviderClient -------------------------------------------------------------
def complete(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
) -> AssistantTurn:
family, rest = self._split(model)
return self._family_client(family).complete(
model=rest, messages=messages, tools=tools, **settings
)
def stream(
self,
*,
model: str,
messages: list[dict[str, Any]],
tools: Optional[list[dict[str, Any]]] = None,
**settings: Any,
):
family, rest = self._split(model)
return self._family_client(family).stream(
model=rest, messages=messages, tools=tools, **settings
)
def capabilities(self, model: str) -> ModelCapabilities:
qualified = model if model.startswith("vertex:") else f"vertex:{model}"
return capabilities_for(qualified)