feat(model): add OpenAI-compatible generic backend for any OpenAI-compatible API
This commit is contained in:
黄云龙
@@ -0,0 +1,448 @@
|
||||
"""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 (
|
||||
CompatAssistantMessage,
|
||||
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()
|
||||
Reference in New Issue
Block a user