Merge pull request #115 from nankingjing/contrib/add-openai-compatible-backend
feat(model): add generic OpenAI-compatible LLM backend
This commit is contained in:
@@ -2,6 +2,59 @@
|
|||||||
|
|
||||||
SkillOpt supports multiple LLM backends. This guide shows how to add your own.
|
SkillOpt supports multiple LLM backends. This guide shows how to add your own.
|
||||||
|
|
||||||
|
## Built-in: the generic OpenAI-compatible backend
|
||||||
|
|
||||||
|
Before writing a new backend, check whether your provider already speaks the
|
||||||
|
OpenAI Chat Completions protocol. Most do — in which case you can use the
|
||||||
|
built-in **`openai_compatible`** backend
|
||||||
|
(`skillopt/model/openai_compatible_backend.py`) with no code changes.
|
||||||
|
|
||||||
|
A single `base_url` + `api_key` pair lets you point SkillOpt at, for example:
|
||||||
|
|
||||||
|
| Provider | `base_url` | Example model |
|
||||||
|
|---|---|---|
|
||||||
|
| DeepSeek | `https://api.deepseek.com/v1` | `deepseek-chat` |
|
||||||
|
| Groq | `https://api.groq.com/openai/v1` | `llama-3.3-70b-versatile` |
|
||||||
|
| Together AI | `https://api.together.xyz/v1` | `meta-llama/Llama-3.3-70B-Instruct-Turbo` |
|
||||||
|
| Ollama (local) | `http://localhost:11434/v1` | `qwen2.5:7b` |
|
||||||
|
| vLLM / SGLang / TGI | `http://localhost:8000/v1` | your served model |
|
||||||
|
| LiteLLM proxy | `http://localhost:4000` | any proxied model |
|
||||||
|
| OpenRouter / Fireworks / xAI / … | provider base URL | provider model id |
|
||||||
|
|
||||||
|
Select it as the optimizer and/or target backend:
|
||||||
|
|
||||||
|
```python
|
||||||
|
import skillopt.model as model
|
||||||
|
|
||||||
|
# Shorthand: use it for both optimizer and target.
|
||||||
|
model.set_backend("openai_compatible")
|
||||||
|
|
||||||
|
# Point it at a provider (shared, or per-role with optimizer_*/target_*).
|
||||||
|
model.configure_openai_compatible(
|
||||||
|
base_url="https://api.deepseek.com/v1",
|
||||||
|
api_key="sk-...",
|
||||||
|
model="deepseek-chat",
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
Or configure it entirely through environment variables (role-specific
|
||||||
|
`OPTIMIZER_*` / `TARGET_*` variants override the shared ones):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export TARGET_BACKEND=openai_compatible
|
||||||
|
export OPENAI_COMPATIBLE_BASE_URL="https://api.groq.com/openai/v1"
|
||||||
|
export OPENAI_COMPATIBLE_API_KEY="gsk_..."
|
||||||
|
export OPENAI_COMPATIBLE_MODEL="llama-3.3-70b-versatile"
|
||||||
|
# Optional: OPENAI_COMPATIBLE_TEMPERATURE, _MAX_TOKENS, _TIMEOUT_SECONDS
|
||||||
|
```
|
||||||
|
|
||||||
|
The backend uses the official `openai` SDK, records token usage through the
|
||||||
|
shared tracker, supports tool/function calling via
|
||||||
|
`chat_target_messages(..., tools=...)`, and exposes
|
||||||
|
`count_tokens()` (tiktoken with a character-based fallback for non-OpenAI
|
||||||
|
models). Only write a brand-new backend if your provider is *not*
|
||||||
|
OpenAI-compatible.
|
||||||
|
|
||||||
## Backend Architecture
|
## Backend Architecture
|
||||||
|
|
||||||
```
|
```
|
||||||
|
|||||||
+101
-4
@@ -7,6 +7,7 @@ from typing import Any
|
|||||||
from skillopt.model import azure_openai as _openai
|
from skillopt.model import azure_openai as _openai
|
||||||
from skillopt.model import claude_backend as _claude
|
from skillopt.model import claude_backend as _claude
|
||||||
from skillopt.model import minimax_backend as _minimax
|
from skillopt.model import minimax_backend as _minimax
|
||||||
|
from skillopt.model import openai_compatible_backend as _openai_compat
|
||||||
from skillopt.model import qwen_backend as _qwen
|
from skillopt.model import qwen_backend as _qwen
|
||||||
from skillopt.model.backend_config import ( # noqa: F401
|
from skillopt.model.backend_config import ( # noqa: F401
|
||||||
configure_claude_code_exec,
|
configure_claude_code_exec,
|
||||||
@@ -55,6 +56,10 @@ def set_backend(name: str | None) -> str:
|
|||||||
set_optimizer_backend("openai_chat")
|
set_optimizer_backend("openai_chat")
|
||||||
set_target_backend("minimax_chat")
|
set_target_backend("minimax_chat")
|
||||||
return "minimax_chat"
|
return "minimax_chat"
|
||||||
|
if normalized in {"openai_compatible", "openai_compatible_chat", "openai-compatible", "compat"}:
|
||||||
|
set_optimizer_backend("openai_compatible")
|
||||||
|
set_target_backend("openai_compatible")
|
||||||
|
return "openai_compatible"
|
||||||
raise ValueError(f"Unsupported legacy backend: {name!r}")
|
raise ValueError(f"Unsupported legacy backend: {name!r}")
|
||||||
|
|
||||||
|
|
||||||
@@ -74,6 +79,8 @@ def get_backend_name() -> str:
|
|||||||
return "qwen_chat"
|
return "qwen_chat"
|
||||||
if optimizer == "openai_chat" and target == "minimax_chat":
|
if optimizer == "openai_chat" and target == "minimax_chat":
|
||||||
return "minimax_chat"
|
return "minimax_chat"
|
||||||
|
if optimizer == "openai_compatible" and target == "openai_compatible":
|
||||||
|
return "openai_compatible"
|
||||||
return f"{optimizer}+{target}"
|
return f"{optimizer}+{target}"
|
||||||
|
|
||||||
|
|
||||||
@@ -115,6 +122,16 @@ def chat_optimizer(
|
|||||||
reasoning_effort=reasoning_effort,
|
reasoning_effort=reasoning_effort,
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
)
|
)
|
||||||
|
if get_optimizer_backend() == "openai_compatible":
|
||||||
|
return _openai_compat.chat_optimizer(
|
||||||
|
system=system,
|
||||||
|
user=user,
|
||||||
|
max_completion_tokens=max_completion_tokens,
|
||||||
|
retries=retries,
|
||||||
|
stage=stage,
|
||||||
|
reasoning_effort=reasoning_effort,
|
||||||
|
timeout=timeout,
|
||||||
|
)
|
||||||
return _openai.chat_optimizer(
|
return _openai.chat_optimizer(
|
||||||
system=system,
|
system=system,
|
||||||
user=user,
|
user=user,
|
||||||
@@ -163,10 +180,20 @@ def chat_target(
|
|||||||
stage=stage,
|
stage=stage,
|
||||||
reasoning_effort=reasoning_effort,
|
reasoning_effort=reasoning_effort,
|
||||||
)
|
)
|
||||||
|
if get_target_backend() == "openai_compatible":
|
||||||
|
return _openai_compat.chat_target(
|
||||||
|
system=system,
|
||||||
|
user=user,
|
||||||
|
max_completion_tokens=max_completion_tokens,
|
||||||
|
retries=retries,
|
||||||
|
stage=stage,
|
||||||
|
reasoning_effort=reasoning_effort,
|
||||||
|
timeout=timeout,
|
||||||
|
)
|
||||||
if not is_target_chat_backend():
|
if not is_target_chat_backend():
|
||||||
raise NotImplementedError(
|
raise NotImplementedError(
|
||||||
"chat_target is only supported with target_backend=openai_chat, claude_chat, qwen_chat, or minimax_chat. "
|
"chat_target is only supported with target_backend=openai_chat, claude_chat, qwen_chat, minimax_chat, "
|
||||||
"Exec backends are handled in environment-specific rollout code."
|
"or openai_compatible. Exec backends are handled in environment-specific rollout code."
|
||||||
)
|
)
|
||||||
return _openai.chat_target(
|
return _openai.chat_target(
|
||||||
system=system,
|
system=system,
|
||||||
@@ -226,6 +253,18 @@ def chat_optimizer_messages(
|
|||||||
return_message=return_message,
|
return_message=return_message,
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
)
|
)
|
||||||
|
if get_optimizer_backend() == "openai_compatible":
|
||||||
|
return _openai_compat.chat_optimizer_messages(
|
||||||
|
messages=messages,
|
||||||
|
max_completion_tokens=max_completion_tokens,
|
||||||
|
retries=retries,
|
||||||
|
stage=stage,
|
||||||
|
reasoning_effort=reasoning_effort,
|
||||||
|
tools=tools,
|
||||||
|
tool_choice=tool_choice,
|
||||||
|
return_message=return_message,
|
||||||
|
timeout=timeout,
|
||||||
|
)
|
||||||
return _openai.chat_optimizer_messages(
|
return _openai.chat_optimizer_messages(
|
||||||
messages=messages,
|
messages=messages,
|
||||||
max_completion_tokens=max_completion_tokens,
|
max_completion_tokens=max_completion_tokens,
|
||||||
@@ -285,10 +324,22 @@ def chat_target_messages(
|
|||||||
tool_choice=tool_choice,
|
tool_choice=tool_choice,
|
||||||
return_message=return_message,
|
return_message=return_message,
|
||||||
)
|
)
|
||||||
|
if get_target_backend() == "openai_compatible":
|
||||||
|
return _openai_compat.chat_target_messages(
|
||||||
|
messages=messages,
|
||||||
|
max_completion_tokens=max_completion_tokens,
|
||||||
|
retries=retries,
|
||||||
|
stage=stage,
|
||||||
|
reasoning_effort=reasoning_effort,
|
||||||
|
tools=tools,
|
||||||
|
tool_choice=tool_choice,
|
||||||
|
return_message=return_message,
|
||||||
|
timeout=timeout,
|
||||||
|
)
|
||||||
if not is_target_chat_backend():
|
if not is_target_chat_backend():
|
||||||
raise NotImplementedError(
|
raise NotImplementedError(
|
||||||
"chat_target_messages is only supported with target_backend=openai_chat, claude_chat, qwen_chat, or minimax_chat. "
|
"chat_target_messages is only supported with target_backend=openai_chat, claude_chat, qwen_chat, "
|
||||||
"Exec backends are handled in environment-specific rollout code."
|
"minimax_chat, or openai_compatible. Exec backends are handled in environment-specific rollout code."
|
||||||
)
|
)
|
||||||
return _openai.chat_target_messages(
|
return _openai.chat_target_messages(
|
||||||
messages=messages,
|
messages=messages,
|
||||||
@@ -387,6 +438,17 @@ def get_token_summary() -> dict:
|
|||||||
summary[stage]["prompt_tokens"] += values["prompt_tokens"]
|
summary[stage]["prompt_tokens"] += values["prompt_tokens"]
|
||||||
summary[stage]["completion_tokens"] += values["completion_tokens"]
|
summary[stage]["completion_tokens"] += values["completion_tokens"]
|
||||||
summary[stage]["total_tokens"] += values["total_tokens"]
|
summary[stage]["total_tokens"] += values["total_tokens"]
|
||||||
|
openai_compat_summary = _openai_compat.get_token_summary()
|
||||||
|
for stage, values in openai_compat_summary.items():
|
||||||
|
if stage == "_total":
|
||||||
|
continue
|
||||||
|
if stage not in summary:
|
||||||
|
summary[stage] = values
|
||||||
|
continue
|
||||||
|
summary[stage]["calls"] += values["calls"]
|
||||||
|
summary[stage]["prompt_tokens"] += values["prompt_tokens"]
|
||||||
|
summary[stage]["completion_tokens"] += values["completion_tokens"]
|
||||||
|
summary[stage]["total_tokens"] += values["total_tokens"]
|
||||||
total = {
|
total = {
|
||||||
"calls": 0,
|
"calls": 0,
|
||||||
"prompt_tokens": 0,
|
"prompt_tokens": 0,
|
||||||
@@ -409,6 +471,7 @@ def reset_token_tracker() -> None:
|
|||||||
_claude.reset_token_tracker()
|
_claude.reset_token_tracker()
|
||||||
_qwen.reset_token_tracker()
|
_qwen.reset_token_tracker()
|
||||||
_minimax.reset_token_tracker()
|
_minimax.reset_token_tracker()
|
||||||
|
_openai_compat.reset_token_tracker()
|
||||||
|
|
||||||
|
|
||||||
def configure_azure_openai(
|
def configure_azure_openai(
|
||||||
@@ -522,11 +585,43 @@ def configure_minimax_chat(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def configure_openai_compatible(
|
||||||
|
*,
|
||||||
|
base_url: str | None = None,
|
||||||
|
api_key: str | None = None,
|
||||||
|
model: str | None = None,
|
||||||
|
temperature: float | str | None = None,
|
||||||
|
timeout_seconds: float | str | None = None,
|
||||||
|
max_tokens: int | str | None = None,
|
||||||
|
optimizer_base_url: str | None = None,
|
||||||
|
optimizer_api_key: str | None = None,
|
||||||
|
optimizer_model: str | None = None,
|
||||||
|
target_base_url: str | None = None,
|
||||||
|
target_api_key: str | None = None,
|
||||||
|
target_model: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
_openai_compat.configure_openai_compatible(
|
||||||
|
base_url=base_url,
|
||||||
|
api_key=api_key,
|
||||||
|
model=model,
|
||||||
|
temperature=temperature,
|
||||||
|
timeout_seconds=timeout_seconds,
|
||||||
|
max_tokens=max_tokens,
|
||||||
|
optimizer_base_url=optimizer_base_url,
|
||||||
|
optimizer_api_key=optimizer_api_key,
|
||||||
|
optimizer_model=optimizer_model,
|
||||||
|
target_base_url=target_base_url,
|
||||||
|
target_api_key=target_api_key,
|
||||||
|
target_model=target_model,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def set_reasoning_effort(effort: str | None) -> None:
|
def set_reasoning_effort(effort: str | None) -> None:
|
||||||
_openai.set_reasoning_effort(effort)
|
_openai.set_reasoning_effort(effort)
|
||||||
_claude.set_reasoning_effort(effort)
|
_claude.set_reasoning_effort(effort)
|
||||||
_qwen.set_reasoning_effort(effort)
|
_qwen.set_reasoning_effort(effort)
|
||||||
_minimax.set_reasoning_effort(effort)
|
_minimax.set_reasoning_effort(effort)
|
||||||
|
_openai_compat.set_reasoning_effort(effort)
|
||||||
|
|
||||||
|
|
||||||
def set_target_deployment(deployment: str) -> None:
|
def set_target_deployment(deployment: str) -> None:
|
||||||
@@ -534,9 +629,11 @@ def set_target_deployment(deployment: str) -> None:
|
|||||||
_claude.set_target_deployment(deployment)
|
_claude.set_target_deployment(deployment)
|
||||||
_qwen.set_target_deployment(deployment)
|
_qwen.set_target_deployment(deployment)
|
||||||
_minimax.set_target_deployment(deployment)
|
_minimax.set_target_deployment(deployment)
|
||||||
|
_openai_compat.set_target_deployment(deployment)
|
||||||
|
|
||||||
|
|
||||||
def set_optimizer_deployment(deployment: str) -> None:
|
def set_optimizer_deployment(deployment: str) -> None:
|
||||||
_openai.set_optimizer_deployment(deployment)
|
_openai.set_optimizer_deployment(deployment)
|
||||||
_claude.set_optimizer_deployment(deployment)
|
_claude.set_optimizer_deployment(deployment)
|
||||||
_qwen.set_optimizer_deployment(deployment)
|
_qwen.set_optimizer_deployment(deployment)
|
||||||
|
_openai_compat.set_optimizer_deployment(deployment)
|
||||||
|
|||||||
@@ -49,10 +49,10 @@ CLAUDE_CODE_EXEC_MAX_THINKING_TOKENS = max(
|
|||||||
def set_optimizer_backend(backend: str) -> None:
|
def set_optimizer_backend(backend: str) -> None:
|
||||||
global OPTIMIZER_BACKEND
|
global OPTIMIZER_BACKEND
|
||||||
OPTIMIZER_BACKEND = normalize_backend_name(backend or "openai_chat")
|
OPTIMIZER_BACKEND = normalize_backend_name(backend or "openai_chat")
|
||||||
if OPTIMIZER_BACKEND not in {"openai_chat", "claude_chat", "qwen_chat", "minimax_chat"}:
|
if OPTIMIZER_BACKEND not in {"openai_chat", "claude_chat", "qwen_chat", "minimax_chat", "openai_compatible"}:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Unsupported optimizer backend: {OPTIMIZER_BACKEND!r}. "
|
f"Unsupported optimizer backend: {OPTIMIZER_BACKEND!r}. "
|
||||||
"Supported values are 'openai_chat', 'claude_chat', 'qwen_chat', and 'minimax_chat'."
|
"Supported values are 'openai_chat', 'claude_chat', 'qwen_chat', 'minimax_chat', and 'openai_compatible'."
|
||||||
)
|
)
|
||||||
os.environ["OPTIMIZER_BACKEND"] = OPTIMIZER_BACKEND
|
os.environ["OPTIMIZER_BACKEND"] = OPTIMIZER_BACKEND
|
||||||
|
|
||||||
@@ -64,10 +64,10 @@ def get_optimizer_backend() -> str:
|
|||||||
def set_target_backend(backend: str) -> None:
|
def set_target_backend(backend: str) -> None:
|
||||||
global TARGET_BACKEND
|
global TARGET_BACKEND
|
||||||
TARGET_BACKEND = normalize_backend_name(backend or "openai_chat")
|
TARGET_BACKEND = normalize_backend_name(backend or "openai_chat")
|
||||||
if TARGET_BACKEND not in {"openai_chat", "claude_chat", "qwen_chat", "minimax_chat", "codex_exec", "claude_code_exec"}:
|
if TARGET_BACKEND not in {"openai_chat", "claude_chat", "qwen_chat", "minimax_chat", "openai_compatible", "codex_exec", "claude_code_exec"}:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Unsupported target backend: {TARGET_BACKEND!r}. "
|
f"Unsupported target backend: {TARGET_BACKEND!r}. "
|
||||||
"Supported values are 'openai_chat', 'claude_chat', 'qwen_chat', 'minimax_chat', 'codex_exec', and 'claude_code_exec'."
|
"Supported values are 'openai_chat', 'claude_chat', 'qwen_chat', 'minimax_chat', 'openai_compatible', 'codex_exec', and 'claude_code_exec'."
|
||||||
)
|
)
|
||||||
os.environ["TARGET_BACKEND"] = TARGET_BACKEND
|
os.environ["TARGET_BACKEND"] = TARGET_BACKEND
|
||||||
|
|
||||||
@@ -81,11 +81,11 @@ def is_target_exec_backend() -> bool:
|
|||||||
|
|
||||||
|
|
||||||
def is_optimizer_chat_backend() -> bool:
|
def is_optimizer_chat_backend() -> bool:
|
||||||
return OPTIMIZER_BACKEND in {"openai_chat", "claude_chat", "qwen_chat", "minimax_chat"}
|
return OPTIMIZER_BACKEND in {"openai_chat", "claude_chat", "qwen_chat", "minimax_chat", "openai_compatible"}
|
||||||
|
|
||||||
|
|
||||||
def is_target_chat_backend() -> bool:
|
def is_target_chat_backend() -> bool:
|
||||||
return TARGET_BACKEND in {"openai_chat", "claude_chat", "qwen_chat", "minimax_chat"}
|
return TARGET_BACKEND in {"openai_chat", "claude_chat", "qwen_chat", "minimax_chat", "openai_compatible"}
|
||||||
|
|
||||||
|
|
||||||
def configure_codex_exec(
|
def configure_codex_exec(
|
||||||
|
|||||||
@@ -26,6 +26,7 @@ _BACKEND_DEFAULT_MODELS = {
|
|||||||
"claude_code_exec": "claude-sonnet-4-6",
|
"claude_code_exec": "claude-sonnet-4-6",
|
||||||
"qwen_chat": "Qwen/Qwen3.5-4B",
|
"qwen_chat": "Qwen/Qwen3.5-4B",
|
||||||
"minimax_chat": "MiniMax-M2.7",
|
"minimax_chat": "MiniMax-M2.7",
|
||||||
|
"openai_compatible": "gpt-4o-mini",
|
||||||
}
|
}
|
||||||
|
|
||||||
_BACKEND_ALIASES = {
|
_BACKEND_ALIASES = {
|
||||||
@@ -44,6 +45,10 @@ _BACKEND_ALIASES = {
|
|||||||
"qwen_chat": "qwen_chat",
|
"qwen_chat": "qwen_chat",
|
||||||
"minimax": "minimax_chat",
|
"minimax": "minimax_chat",
|
||||||
"minimax_chat": "minimax_chat",
|
"minimax_chat": "minimax_chat",
|
||||||
|
"openai_compatible": "openai_compatible",
|
||||||
|
"openai_compatible_chat": "openai_compatible",
|
||||||
|
"openai-compatible": "openai_compatible",
|
||||||
|
"compat": "openai_compatible",
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,447 @@
|
|||||||
|
"""Generic OpenAI-compatible chat backend for optimizer and target paths.
|
||||||
|
|
||||||
|
This backend talks to *any* service that exposes an OpenAI-compatible
|
||||||
|
``/chat/completions`` endpoint through the official ``openai`` SDK. A single
|
||||||
|
implementation therefore covers a large family of providers, for example:
|
||||||
|
|
||||||
|
* DeepSeek (``https://api.deepseek.com``)
|
||||||
|
* Groq (``https://api.groq.com/openai/v1``)
|
||||||
|
* Together AI (``https://api.together.xyz/v1``)
|
||||||
|
* Mistral / Fireworks / OpenRouter / Perplexity / xAI Grok
|
||||||
|
* Ollama (``http://localhost:11434/v1``)
|
||||||
|
* vLLM / SGLang / TGI self-hosted servers
|
||||||
|
* LiteLLM proxy (``http://localhost:4000``)
|
||||||
|
* Azure OpenAI and OpenAI itself
|
||||||
|
|
||||||
|
Unlike the Azure backend it never assumes Azure-specific auth or the Responses
|
||||||
|
API — it only needs a ``base_url`` and an ``api_key`` (some local servers accept
|
||||||
|
any key, so the key is optional and falls back to a harmless placeholder).
|
||||||
|
|
||||||
|
The module mirrors the callable surface of the other chat backends
|
||||||
|
(:mod:`skillopt.model.qwen_backend`, :mod:`skillopt.model.minimax_backend`) so
|
||||||
|
it can be selected as the optimizer and/or target backend and routed through
|
||||||
|
:mod:`skillopt.model`.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from openai import OpenAI
|
||||||
|
|
||||||
|
from skillopt.model.common import (
|
||||||
|
TokenTracker,
|
||||||
|
compat_message_from_chat_message,
|
||||||
|
default_model_for_backend,
|
||||||
|
usage_from_openai_usage,
|
||||||
|
)
|
||||||
|
|
||||||
|
BACKEND_NAME = "openai_compatible"
|
||||||
|
|
||||||
|
# A neutral, widely-available default. Real deployments should set the model
|
||||||
|
# explicitly (e.g. "deepseek-chat", "llama-3.3-70b-versatile", "qwen2.5:7b").
|
||||||
|
_DEFAULT_BASE_URL = "https://api.openai.com/v1"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class OpenAICompatibleConfig:
|
||||||
|
base_url: str
|
||||||
|
api_key: str
|
||||||
|
deployment: str
|
||||||
|
timeout_seconds: float
|
||||||
|
max_tokens: int
|
||||||
|
temperature: float | None
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_optional_float(value: Any) -> float | None:
|
||||||
|
if value is None:
|
||||||
|
return None
|
||||||
|
raw = str(value).strip()
|
||||||
|
return float(raw) if raw else None
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_int(value: Any, default: int) -> int:
|
||||||
|
if value is None:
|
||||||
|
return default
|
||||||
|
raw = str(value).strip()
|
||||||
|
return int(raw) if raw else default
|
||||||
|
|
||||||
|
|
||||||
|
def _role_env(role: str, key: str, default: str) -> str:
|
||||||
|
"""Resolve a config value, preferring role-specific over shared env vars."""
|
||||||
|
role_key = f"{role.upper()}_OPENAI_COMPATIBLE_{key}"
|
||||||
|
generic_key = f"OPENAI_COMPATIBLE_{key}"
|
||||||
|
return os.environ.get(role_key) or os.environ.get(generic_key) or default
|
||||||
|
|
||||||
|
|
||||||
|
def _initial_config(role: str) -> OpenAICompatibleConfig:
|
||||||
|
role_upper = role.upper()
|
||||||
|
deployment_env = "OPTIMIZER_DEPLOYMENT" if role == "optimizer" else "TARGET_DEPLOYMENT"
|
||||||
|
return OpenAICompatibleConfig(
|
||||||
|
base_url=_role_env(role, "BASE_URL", _DEFAULT_BASE_URL),
|
||||||
|
api_key=_role_env(role, "API_KEY", ""),
|
||||||
|
deployment=(
|
||||||
|
os.environ.get(f"{role_upper}_OPENAI_COMPATIBLE_MODEL")
|
||||||
|
or os.environ.get("OPENAI_COMPATIBLE_MODEL")
|
||||||
|
or os.environ.get(deployment_env)
|
||||||
|
or default_model_for_backend(BACKEND_NAME)
|
||||||
|
),
|
||||||
|
timeout_seconds=float(_role_env(role, "TIMEOUT_SECONDS", "300") or 300),
|
||||||
|
max_tokens=_parse_int(_role_env(role, "MAX_TOKENS", "8000"), 8000),
|
||||||
|
temperature=_parse_optional_float(_role_env(role, "TEMPERATURE", "")),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
OPTIMIZER_CONFIG = _initial_config("optimizer")
|
||||||
|
TARGET_CONFIG = _initial_config("target")
|
||||||
|
|
||||||
|
_config_lock = threading.Lock()
|
||||||
|
_client_lock = threading.Lock()
|
||||||
|
tracker = TokenTracker()
|
||||||
|
|
||||||
|
_optimizer_client: OpenAI | None = None
|
||||||
|
_target_client: OpenAI | None = None
|
||||||
|
|
||||||
|
|
||||||
|
def _config_for(role: str) -> OpenAICompatibleConfig:
|
||||||
|
return OPTIMIZER_CONFIG if role == "optimizer" else TARGET_CONFIG
|
||||||
|
|
||||||
|
|
||||||
|
def _build_client(config: OpenAICompatibleConfig) -> OpenAI:
|
||||||
|
return OpenAI(
|
||||||
|
base_url=config.base_url.rstrip("/") or _DEFAULT_BASE_URL,
|
||||||
|
# Some OpenAI-compatible servers (Ollama, vLLM, local proxies) do not
|
||||||
|
# require an API key. The SDK still expects a non-empty string, so fall
|
||||||
|
# back to a harmless placeholder when none is configured.
|
||||||
|
api_key=config.api_key or "dummy",
|
||||||
|
timeout=config.timeout_seconds,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _get_client(role: str) -> OpenAI:
|
||||||
|
global _optimizer_client, _target_client
|
||||||
|
with _client_lock:
|
||||||
|
if role == "optimizer":
|
||||||
|
if _optimizer_client is None:
|
||||||
|
_optimizer_client = _build_client(OPTIMIZER_CONFIG)
|
||||||
|
return _optimizer_client
|
||||||
|
if _target_client is None:
|
||||||
|
_target_client = _build_client(TARGET_CONFIG)
|
||||||
|
return _target_client
|
||||||
|
|
||||||
|
|
||||||
|
def _reset_clients() -> None:
|
||||||
|
global _optimizer_client, _target_client
|
||||||
|
with _client_lock:
|
||||||
|
_optimizer_client = None
|
||||||
|
_target_client = None
|
||||||
|
|
||||||
|
|
||||||
|
def count_tokens(text: str, model: str | None = None) -> int:
|
||||||
|
"""Best-effort token count for a string.
|
||||||
|
|
||||||
|
Uses ``tiktoken`` when available (per-model encoding, falling back to the
|
||||||
|
``cl100k_base`` encoding). If ``tiktoken`` is not installed or fails — which
|
||||||
|
is common for non-OpenAI models served through compatible APIs — it falls
|
||||||
|
back to a character-based estimate of roughly four characters per token.
|
||||||
|
"""
|
||||||
|
if not text:
|
||||||
|
return 0
|
||||||
|
try:
|
||||||
|
import tiktoken
|
||||||
|
|
||||||
|
try:
|
||||||
|
encoding = tiktoken.encoding_for_model(model or "gpt-4o")
|
||||||
|
except Exception: # noqa: BLE001 - unknown/non-OpenAI model name
|
||||||
|
encoding = tiktoken.get_encoding("cl100k_base")
|
||||||
|
return len(encoding.encode(text))
|
||||||
|
except Exception: # noqa: BLE001 - tiktoken missing or encoding failure
|
||||||
|
# Rough heuristic: ~4 characters per token for English-like text.
|
||||||
|
return max(1, (len(text) + 3) // 4)
|
||||||
|
|
||||||
|
|
||||||
|
def _chat_messages_impl(
|
||||||
|
messages: list[dict[str, Any]],
|
||||||
|
max_completion_tokens: int,
|
||||||
|
retries: int,
|
||||||
|
stage: str,
|
||||||
|
*,
|
||||||
|
role: str,
|
||||||
|
tools: list[dict[str, Any]] | None = None,
|
||||||
|
tool_choice: str | dict[str, Any] | None = None,
|
||||||
|
return_message: bool = False,
|
||||||
|
deployment: str | None = None,
|
||||||
|
timeout: float | None = None,
|
||||||
|
) -> tuple[Any, dict[str, int]]:
|
||||||
|
config = _config_for(role)
|
||||||
|
client = _get_client(role)
|
||||||
|
kwargs: dict[str, Any] = {
|
||||||
|
"model": deployment or config.deployment,
|
||||||
|
"messages": messages,
|
||||||
|
# ``max_tokens`` (rather than ``max_completion_tokens``) is the field
|
||||||
|
# understood by the broadest set of OpenAI-compatible providers.
|
||||||
|
"max_tokens": min(max_completion_tokens, config.max_tokens),
|
||||||
|
}
|
||||||
|
if config.temperature is not None:
|
||||||
|
kwargs["temperature"] = config.temperature
|
||||||
|
if tools:
|
||||||
|
kwargs["tools"] = tools
|
||||||
|
if tool_choice is not None:
|
||||||
|
kwargs["tool_choice"] = tool_choice
|
||||||
|
if timeout is not None:
|
||||||
|
kwargs["timeout"] = timeout
|
||||||
|
|
||||||
|
last_err: Exception | None = None
|
||||||
|
for attempt in range(retries):
|
||||||
|
try:
|
||||||
|
resp = client.chat.completions.create(**kwargs)
|
||||||
|
choices = getattr(resp, "choices", None) or []
|
||||||
|
if not choices:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"OpenAI-compatible API returned no choices: {resp!r}"
|
||||||
|
)
|
||||||
|
message = choices[0].message
|
||||||
|
text = message.content or ""
|
||||||
|
usage_info = usage_from_openai_usage(getattr(resp, "usage", None))
|
||||||
|
tracker.record(
|
||||||
|
stage,
|
||||||
|
usage_info["prompt_tokens"],
|
||||||
|
usage_info["completion_tokens"],
|
||||||
|
)
|
||||||
|
if return_message:
|
||||||
|
return compat_message_from_chat_message(message), usage_info
|
||||||
|
return text, usage_info
|
||||||
|
except Exception as e: # noqa: BLE001
|
||||||
|
last_err = e
|
||||||
|
time.sleep(min(2 ** attempt, 30))
|
||||||
|
raise RuntimeError(
|
||||||
|
f"OpenAI-compatible chat call failed after {retries} retries: {last_err}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Public API (mirrors the other chat backends) ─────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def chat_optimizer(
|
||||||
|
system: str,
|
||||||
|
user: str,
|
||||||
|
max_completion_tokens: int = 16384,
|
||||||
|
retries: int = 5,
|
||||||
|
stage: str = "optimizer",
|
||||||
|
reasoning_effort: str | None = None,
|
||||||
|
timeout: float | None = None,
|
||||||
|
) -> tuple[str, dict[str, int]]:
|
||||||
|
del reasoning_effort # not forwarded — kept for a uniform signature
|
||||||
|
messages = [
|
||||||
|
{"role": "system", "content": system},
|
||||||
|
{"role": "user", "content": user},
|
||||||
|
]
|
||||||
|
return _chat_messages_impl(
|
||||||
|
messages,
|
||||||
|
max_completion_tokens,
|
||||||
|
retries,
|
||||||
|
stage,
|
||||||
|
role="optimizer",
|
||||||
|
timeout=timeout,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def chat_target(
|
||||||
|
system: str,
|
||||||
|
user: str,
|
||||||
|
max_completion_tokens: int = 16384,
|
||||||
|
retries: int = 5,
|
||||||
|
stage: str = "target",
|
||||||
|
reasoning_effort: str | None = None,
|
||||||
|
timeout: float | None = None,
|
||||||
|
) -> tuple[str, dict[str, int]]:
|
||||||
|
del reasoning_effort
|
||||||
|
messages = [
|
||||||
|
{"role": "system", "content": system},
|
||||||
|
{"role": "user", "content": user},
|
||||||
|
]
|
||||||
|
return _chat_messages_impl(
|
||||||
|
messages,
|
||||||
|
max_completion_tokens,
|
||||||
|
retries,
|
||||||
|
stage,
|
||||||
|
role="target",
|
||||||
|
timeout=timeout,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def chat_optimizer_messages(
|
||||||
|
messages: list[dict[str, Any]],
|
||||||
|
max_completion_tokens: int = 16384,
|
||||||
|
retries: int = 5,
|
||||||
|
stage: str = "optimizer",
|
||||||
|
reasoning_effort: str | None = None,
|
||||||
|
*,
|
||||||
|
tools: list[dict[str, Any]] | None = None,
|
||||||
|
tool_choice: str | dict[str, Any] | None = None,
|
||||||
|
return_message: bool = False,
|
||||||
|
timeout: float | None = None,
|
||||||
|
) -> tuple[Any, dict[str, int]]:
|
||||||
|
del reasoning_effort
|
||||||
|
return _chat_messages_impl(
|
||||||
|
messages,
|
||||||
|
max_completion_tokens,
|
||||||
|
retries,
|
||||||
|
stage,
|
||||||
|
role="optimizer",
|
||||||
|
tools=tools,
|
||||||
|
tool_choice=tool_choice,
|
||||||
|
return_message=return_message,
|
||||||
|
timeout=timeout,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def chat_target_messages(
|
||||||
|
messages: list[dict[str, Any]],
|
||||||
|
max_completion_tokens: int = 16384,
|
||||||
|
retries: int = 5,
|
||||||
|
stage: str = "target",
|
||||||
|
reasoning_effort: str | None = None,
|
||||||
|
*,
|
||||||
|
tools: list[dict[str, Any]] | None = None,
|
||||||
|
tool_choice: str | dict[str, Any] | None = None,
|
||||||
|
return_message: bool = False,
|
||||||
|
timeout: float | None = None,
|
||||||
|
) -> tuple[Any, dict[str, int]]:
|
||||||
|
del reasoning_effort
|
||||||
|
return _chat_messages_impl(
|
||||||
|
messages,
|
||||||
|
max_completion_tokens,
|
||||||
|
retries,
|
||||||
|
stage,
|
||||||
|
role="target",
|
||||||
|
tools=tools,
|
||||||
|
tool_choice=tool_choice,
|
||||||
|
return_message=return_message,
|
||||||
|
timeout=timeout,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Configuration / lifecycle ────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def _update_config(
|
||||||
|
config: OpenAICompatibleConfig,
|
||||||
|
role: str,
|
||||||
|
*,
|
||||||
|
base_url: str | None = None,
|
||||||
|
api_key: str | None = None,
|
||||||
|
deployment: str | None = None,
|
||||||
|
temperature: float | str | None = None,
|
||||||
|
timeout_seconds: float | str | None = None,
|
||||||
|
max_tokens: int | str | None = None,
|
||||||
|
) -> None:
|
||||||
|
env_prefix = role.upper()
|
||||||
|
if base_url is not None:
|
||||||
|
config.base_url = str(base_url).strip() or config.base_url
|
||||||
|
os.environ[f"{env_prefix}_OPENAI_COMPATIBLE_BASE_URL"] = config.base_url
|
||||||
|
if api_key is not None:
|
||||||
|
config.api_key = str(api_key).strip()
|
||||||
|
os.environ[f"{env_prefix}_OPENAI_COMPATIBLE_API_KEY"] = config.api_key
|
||||||
|
if deployment is not None:
|
||||||
|
config.deployment = str(deployment).strip() or config.deployment
|
||||||
|
os.environ[f"{env_prefix}_OPENAI_COMPATIBLE_MODEL"] = config.deployment
|
||||||
|
if temperature is not None:
|
||||||
|
raw = str(temperature).strip()
|
||||||
|
config.temperature = float(raw) if raw else None
|
||||||
|
os.environ[f"{env_prefix}_OPENAI_COMPATIBLE_TEMPERATURE"] = raw
|
||||||
|
if timeout_seconds is not None:
|
||||||
|
config.timeout_seconds = float(timeout_seconds)
|
||||||
|
os.environ[f"{env_prefix}_OPENAI_COMPATIBLE_TIMEOUT_SECONDS"] = str(timeout_seconds)
|
||||||
|
if max_tokens is not None:
|
||||||
|
config.max_tokens = int(max_tokens)
|
||||||
|
os.environ[f"{env_prefix}_OPENAI_COMPATIBLE_MAX_TOKENS"] = str(max_tokens)
|
||||||
|
|
||||||
|
|
||||||
|
def configure_openai_compatible(
|
||||||
|
*,
|
||||||
|
base_url: str | None = None,
|
||||||
|
api_key: str | None = None,
|
||||||
|
model: str | None = None,
|
||||||
|
temperature: float | str | None = None,
|
||||||
|
timeout_seconds: float | str | None = None,
|
||||||
|
max_tokens: int | str | None = None,
|
||||||
|
optimizer_base_url: str | None = None,
|
||||||
|
optimizer_api_key: str | None = None,
|
||||||
|
optimizer_model: str | None = None,
|
||||||
|
target_base_url: str | None = None,
|
||||||
|
target_api_key: str | None = None,
|
||||||
|
target_model: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Configure the generic OpenAI-compatible backend at runtime.
|
||||||
|
|
||||||
|
Shared values apply to both the optimizer and target roles; the
|
||||||
|
``optimizer_*`` / ``target_*`` variants override them per role.
|
||||||
|
"""
|
||||||
|
with _config_lock:
|
||||||
|
if base_url is not None:
|
||||||
|
os.environ["OPENAI_COMPATIBLE_BASE_URL"] = str(base_url).strip()
|
||||||
|
if api_key is not None:
|
||||||
|
os.environ["OPENAI_COMPATIBLE_API_KEY"] = str(api_key).strip()
|
||||||
|
if model is not None:
|
||||||
|
os.environ["OPENAI_COMPATIBLE_MODEL"] = str(model).strip()
|
||||||
|
if temperature is not None:
|
||||||
|
os.environ["OPENAI_COMPATIBLE_TEMPERATURE"] = str(temperature).strip()
|
||||||
|
if timeout_seconds is not None:
|
||||||
|
os.environ["OPENAI_COMPATIBLE_TIMEOUT_SECONDS"] = str(timeout_seconds)
|
||||||
|
if max_tokens is not None:
|
||||||
|
os.environ["OPENAI_COMPATIBLE_MAX_TOKENS"] = str(max_tokens)
|
||||||
|
_update_config(
|
||||||
|
OPTIMIZER_CONFIG,
|
||||||
|
"optimizer",
|
||||||
|
base_url=optimizer_base_url if optimizer_base_url is not None else base_url,
|
||||||
|
api_key=optimizer_api_key if optimizer_api_key is not None else api_key,
|
||||||
|
deployment=optimizer_model if optimizer_model is not None else model,
|
||||||
|
temperature=temperature,
|
||||||
|
timeout_seconds=timeout_seconds,
|
||||||
|
max_tokens=max_tokens,
|
||||||
|
)
|
||||||
|
_update_config(
|
||||||
|
TARGET_CONFIG,
|
||||||
|
"target",
|
||||||
|
base_url=target_base_url if target_base_url is not None else base_url,
|
||||||
|
api_key=target_api_key if target_api_key is not None else api_key,
|
||||||
|
deployment=target_model if target_model is not None else model,
|
||||||
|
temperature=temperature,
|
||||||
|
timeout_seconds=timeout_seconds,
|
||||||
|
max_tokens=max_tokens,
|
||||||
|
)
|
||||||
|
_reset_clients()
|
||||||
|
|
||||||
|
|
||||||
|
def get_max_tokens() -> int:
|
||||||
|
return TARGET_CONFIG.max_tokens
|
||||||
|
|
||||||
|
|
||||||
|
def get_token_summary() -> dict[str, dict[str, int]]:
|
||||||
|
return tracker.summary()
|
||||||
|
|
||||||
|
|
||||||
|
def reset_token_tracker() -> None:
|
||||||
|
tracker.reset()
|
||||||
|
|
||||||
|
|
||||||
|
def set_reasoning_effort(effort: str | None) -> None:
|
||||||
|
# Reasoning effort is provider-specific and not universally supported by
|
||||||
|
# OpenAI-compatible endpoints, so it is intentionally a no-op here.
|
||||||
|
del effort
|
||||||
|
|
||||||
|
|
||||||
|
def set_target_deployment(deployment: str) -> None:
|
||||||
|
TARGET_CONFIG.deployment = deployment or default_model_for_backend(BACKEND_NAME)
|
||||||
|
os.environ["TARGET_DEPLOYMENT"] = TARGET_CONFIG.deployment
|
||||||
|
_reset_clients()
|
||||||
|
|
||||||
|
|
||||||
|
def set_optimizer_deployment(deployment: str) -> None:
|
||||||
|
OPTIMIZER_CONFIG.deployment = deployment or default_model_for_backend(BACKEND_NAME)
|
||||||
|
os.environ["OPTIMIZER_DEPLOYMENT"] = OPTIMIZER_CONFIG.deployment
|
||||||
|
_reset_clients()
|
||||||
@@ -0,0 +1,146 @@
|
|||||||
|
"""Tests for the generic OpenAI-compatible model backend."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
import skillopt.model as model
|
||||||
|
from skillopt.model import backend_config
|
||||||
|
from skillopt.model import openai_compatible_backend as backend
|
||||||
|
|
||||||
|
|
||||||
|
class _CompletionRecorder:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.calls: list[dict[str, Any]] = []
|
||||||
|
|
||||||
|
def create(self, **kwargs: Any) -> Any:
|
||||||
|
self.calls.append(kwargs)
|
||||||
|
message = SimpleNamespace(content="ok", tool_calls=[])
|
||||||
|
usage = SimpleNamespace(prompt_tokens=2, completion_tokens=3, total_tokens=5)
|
||||||
|
return SimpleNamespace(choices=[SimpleNamespace(message=message)], usage=usage)
|
||||||
|
|
||||||
|
|
||||||
|
class _Client:
|
||||||
|
def __init__(self, recorder: _CompletionRecorder) -> None:
|
||||||
|
self.chat = SimpleNamespace(completions=recorder)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(autouse=True)
|
||||||
|
def isolate_backend_state(monkeypatch: pytest.MonkeyPatch):
|
||||||
|
optimizer_backend = backend_config.get_optimizer_backend()
|
||||||
|
target_backend = backend_config.get_target_backend()
|
||||||
|
optimizer_config = vars(backend.OPTIMIZER_CONFIG).copy()
|
||||||
|
target_config = vars(backend.TARGET_CONFIG).copy()
|
||||||
|
backend.reset_token_tracker()
|
||||||
|
yield
|
||||||
|
backend.reset_token_tracker()
|
||||||
|
vars(backend.OPTIMIZER_CONFIG).update(optimizer_config)
|
||||||
|
vars(backend.TARGET_CONFIG).update(target_config)
|
||||||
|
backend_config.set_optimizer_backend(optimizer_backend)
|
||||||
|
backend_config.set_target_backend(target_backend)
|
||||||
|
backend._reset_clients()
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_preserves_role_specific_values() -> None:
|
||||||
|
model.configure_openai_compatible(
|
||||||
|
base_url="https://shared.example/v1",
|
||||||
|
api_key="shared-key",
|
||||||
|
model="shared-model",
|
||||||
|
optimizer_base_url="https://optimizer.example/v1",
|
||||||
|
optimizer_api_key="optimizer-key",
|
||||||
|
optimizer_model="optimizer-model",
|
||||||
|
target_base_url="https://target.example/v1",
|
||||||
|
target_api_key="target-key",
|
||||||
|
target_model="target-model",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert backend.OPTIMIZER_CONFIG.base_url == "https://optimizer.example/v1"
|
||||||
|
assert backend.OPTIMIZER_CONFIG.api_key == "optimizer-key"
|
||||||
|
assert backend.OPTIMIZER_CONFIG.deployment == "optimizer-model"
|
||||||
|
assert backend.TARGET_CONFIG.base_url == "https://target.example/v1"
|
||||||
|
assert backend.TARGET_CONFIG.api_key == "target-key"
|
||||||
|
assert backend.TARGET_CONFIG.deployment == "target-model"
|
||||||
|
|
||||||
|
|
||||||
|
def test_optimizer_and_target_route_to_their_own_clients(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
optimizer_calls = _CompletionRecorder()
|
||||||
|
target_calls = _CompletionRecorder()
|
||||||
|
monkeypatch.setattr(
|
||||||
|
backend,
|
||||||
|
"_get_client",
|
||||||
|
lambda role: _Client(optimizer_calls if role == "optimizer" else target_calls),
|
||||||
|
)
|
||||||
|
model.set_optimizer_backend("openai_compatible")
|
||||||
|
model.set_target_backend("openai_compatible")
|
||||||
|
backend.OPTIMIZER_CONFIG.deployment = "optimizer-model"
|
||||||
|
backend.TARGET_CONFIG.deployment = "target-model"
|
||||||
|
|
||||||
|
model.chat_optimizer("system", "user", retries=1)
|
||||||
|
model.chat_target_messages([{"role": "user", "content": "question"}], retries=1)
|
||||||
|
|
||||||
|
assert optimizer_calls.calls[0]["model"] == "optimizer-model"
|
||||||
|
assert target_calls.calls[0]["model"] == "target-model"
|
||||||
|
|
||||||
|
|
||||||
|
def test_client_creation_is_lazy(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
builds: list[str] = []
|
||||||
|
|
||||||
|
def build(config: backend.OpenAICompatibleConfig) -> _Client:
|
||||||
|
builds.append(config.deployment)
|
||||||
|
return _Client(_CompletionRecorder())
|
||||||
|
|
||||||
|
backend._reset_clients()
|
||||||
|
monkeypatch.setattr(backend, "_build_client", build)
|
||||||
|
assert builds == []
|
||||||
|
|
||||||
|
backend._get_client("optimizer")
|
||||||
|
backend._get_client("optimizer")
|
||||||
|
|
||||||
|
assert builds == [backend.OPTIMIZER_CONFIG.deployment]
|
||||||
|
|
||||||
|
|
||||||
|
def test_combined_token_summary_counts_each_backend_once(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
def make(stage: str) -> dict[str, dict[str, int]]:
|
||||||
|
return {
|
||||||
|
stage: {
|
||||||
|
"calls": 1,
|
||||||
|
"prompt_tokens": 2,
|
||||||
|
"completion_tokens": 3,
|
||||||
|
"total_tokens": 5,
|
||||||
|
},
|
||||||
|
"_total": {
|
||||||
|
"calls": 1,
|
||||||
|
"prompt_tokens": 2,
|
||||||
|
"completion_tokens": 3,
|
||||||
|
"total_tokens": 5,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
monkeypatch.setattr(model._openai, "get_token_summary", lambda: make("azure"))
|
||||||
|
monkeypatch.setattr(model._claude, "get_token_summary", lambda: make("claude"))
|
||||||
|
monkeypatch.setattr(model._qwen, "get_token_summary", lambda: make("qwen"))
|
||||||
|
monkeypatch.setattr(model._minimax, "get_token_summary", lambda: make("minimax"))
|
||||||
|
monkeypatch.setattr(model._openai_compat, "get_token_summary", lambda: make("openai_compatible"))
|
||||||
|
|
||||||
|
combined = model.get_token_summary()
|
||||||
|
|
||||||
|
assert set(combined) - {"_total"} == {"azure", "claude", "qwen", "minimax", "openai_compatible"}
|
||||||
|
expected_stage_total = {
|
||||||
|
"calls": 1,
|
||||||
|
"prompt_tokens": 2,
|
||||||
|
"completion_tokens": 3,
|
||||||
|
"total_tokens": 5,
|
||||||
|
}
|
||||||
|
for stage in {"azure", "claude", "qwen", "minimax", "openai_compatible"}:
|
||||||
|
assert combined[stage] == expected_stage_total
|
||||||
|
assert combined["_total"] == {
|
||||||
|
"calls": 5,
|
||||||
|
"prompt_tokens": 10,
|
||||||
|
"completion_tokens": 15,
|
||||||
|
"total_tokens": 25,
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user