feat(model): add OpenAI-compatible generic backend for any OpenAI-compatible API

skillopt/model/openai_compatible_backend.py

@@ -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()