DataClaw/dataclaw-api/app/core/llm_provider.py

import os
from typing import Optional, Dict

from nanobot.providers.azure_openai_provider import AzureOpenAIProvider
from nanobot.providers.openai_codex_provider import OpenAICodexProvider
from nanobot.providers.registry import find_by_name

from app.core.patched_openai_compat_provider import PatchedOpenAICompatProvider


def normalize_provider_name(provider: Optional[str]) -> Optional[str]:
    if not provider:
        return None
    normalized = provider.strip().lower()
    alias_map = {
        "azure": "azure_openai",
        "local": "vllm",
    }
    return alias_map.get(normalized, normalized)


def _running_in_docker() -> bool:
    # Best-effort, cross-platform detection.
    if os.environ.get("DATACLAW_RUNNING_IN_DOCKER", "").strip().lower() in ("1", "true", "yes", "y"):
        return True
    return os.path.exists("/.dockerenv")


def _rewrite_localhost_api_base(api_base: Optional[str]) -> Optional[str]:
    """
    When running inside Docker, `localhost` points to the container itself.
    For host-local LLMs (Ollama/vLLM), users often configure `http://localhost:...`,
    which breaks in containers. We rewrite it to `host.docker.internal`.
    """
    if not api_base:
        return api_base
    base = api_base.strip()
    if base.startswith("http://localhost") or base.startswith("https://localhost"):
        return base.replace("://localhost", "://host.docker.internal", 1)
    if base.startswith("http://127.0.0.1") or base.startswith("https://127.0.0.1"):
        return base.replace("://127.0.0.1", "://host.docker.internal", 1)
    return api_base


def build_llm_provider(
    *,
    model: str,
    provider: Optional[str] = None,
    api_key: Optional[str] = None,
    api_base: Optional[str] = None,
    extra_headers: Optional[Dict[str, str]] = None,
):
    provider_name = normalize_provider_name(provider)
    spec = find_by_name(provider_name) if provider_name else None
    backend = spec.backend if spec else "openai_compat"
    if _running_in_docker():
        api_base = _rewrite_localhost_api_base(api_base)

    if backend == "openai_codex" or model.startswith("openai-codex/"):
        return OpenAICodexProvider(default_model=model)

    if backend == "azure_openai":
        if not api_key or not api_base:
            raise ValueError("Azure OpenAI requires api_key and api_base.")
        return AzureOpenAIProvider(
            api_key=api_key,
            api_base=api_base,
            default_model=model,
        )

    if backend == "anthropic":
        from nanobot.providers.anthropic_provider import AnthropicProvider

        return AnthropicProvider(
            api_key=api_key,
            api_base=api_base,
            default_model=model,
            extra_headers=extra_headers,
        )

    return PatchedOpenAICompatProvider(
        api_key=api_key,
        api_base=api_base,
        default_model=model,
        extra_headers=extra_headers,
        spec=spec,
    )