feat: OpenMesh 基础平台与 MD/PDF 转换技能
- 后端: 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:
54
coworker/providers/__init__.py
Normal file
54
coworker/providers/__init__.py
Normal 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",
|
||||
]
|
||||
653
coworker/providers/anthropic_provider.py
Normal file
653
coworker/providers/anthropic_provider.py
Normal 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
136
coworker/providers/base.py
Normal 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)
|
||||
)
|
||||
579
coworker/providers/bedrock_provider.py
Normal file
579
coworker/providers/bedrock_provider.py
Normal 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)
|
||||
147
coworker/providers/capabilities.py
Normal file
147
coworker/providers/capabilities.py
Normal 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
|
||||
),
|
||||
)
|
||||
472
coworker/providers/codex_auth.py
Normal file
472
coworker/providers/codex_auth.py
Normal 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}.",
|
||||
}
|
||||
135
coworker/providers/codex_provider.py
Normal file
135
coworker/providers/codex_provider.py
Normal 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()
|
||||
57
coworker/providers/errors.py
Normal file
57
coworker/providers/errors.py
Normal 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
|
||||
547
coworker/providers/gemini_provider.py
Normal file
547
coworker/providers/gemini_provider.py
Normal 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,
|
||||
)
|
||||
)
|
||||
305
coworker/providers/matrix.py
Normal file
305
coworker/providers/matrix.py
Normal 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)]
|
||||
621
coworker/providers/openai_provider.py
Normal file
621
coworker/providers/openai_provider.py
Normal 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)
|
||||
462
coworker/providers/openai_responses.py
Normal file
462
coworker/providers/openai_responses.py
Normal 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,
|
||||
)
|
||||
)
|
||||
1020
coworker/providers/registry.py
Normal file
1020
coworker/providers/registry.py
Normal file
File diff suppressed because it is too large
Load Diff
118
coworker/providers/router.py
Normal file
118
coworker/providers/router.py
Normal 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)
|
||||
238
coworker/providers/vertex_provider.py
Normal file
238
coworker/providers/vertex_provider.py
Normal 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)
|
||||
Reference in New Issue
Block a user