databao/configs/llm.py

import os
from pathlib import Path
from typing import Any, Literal

from langchain.chat_models import init_chat_model
from langchain_core.language_models import BaseChatModel
from pydantic import BaseModel, ConfigDict, Field
from typing_extensions import Self

_OPENAI_PREFIXES = ["gpt", "o1", "o3", "o4"]
_ANTHROPIC_PREFIXES = ["claude", "anthropic"]
_OPENAI_REASONING_INFIXES = ["o1", "o3", "o4", "gpt-5", "openai/gpt-oss"]


# TODO: add a config folder for LLM configs, make it initializable from hydra configs
class LLMConfig(BaseModel):
    """Base class with all fields and computed logic for LLM configurations."""

    # Fields declared in parent - can be overridden in children with different defaults
    name: str
    """The model name can be of the form 'provider:name' or 'name'."""

    temperature: float = 0.0
    max_tokens: int = 8192
    """Maximum number of tokens to generate."""
    reasoning_effort: str = "medium"
    """Reasoning effort is used for OpenAI reasoning models only. 
    Warning: reasoning can use a lot of tokens! OpenAI recommends at least 25000 tokens"""
    cache_system_prompt: bool = True
    """Cache system prompt with prompt caching. Only used for Anthropic models."""
    # TODO multi-turn prompt caching

    max_tokens_before_cleaning: int = 10000
    """Number of tokens to start history cleaning. Each Executor has it's own cleaning strategy."""

    timeout: int | None | Literal["auto"] = "auto"
    """Timeout in seconds for LLM calls. If None, use the LLM provider's defaults. 
    If 'auto', use a default timeout (60s) that increases for reasoning models."""

    api_base_url: str | None = None
    """Base URL for an OpenAI-compatible API like 'http://localhost:8080/v1'. Mostly used for running local models."""
    use_responses_api: bool = True
    """Use the [responses API](https://platform.openai.com/docs/guides/migrate-to-responses) for OpenAI models. 
    If False, use the old Chat Completions API (useful for local models that don't support the new responses API)."""

    ollama_pull_model: bool = True
    """Pull the model from Ollama if it's not already downloaded. Only applicable for the 'ollama' model provider."""

    model_kwargs: dict[str, Any] = Field(default_factory=dict)
    """Additional kwargs for the model constructor."""

    model_config = ConfigDict(frozen=True, extra="forbid")

    def _resolve_timeout(self) -> float | None:
        if self.timeout == "auto":
            return 360 if _is_reasoning_model(self.name) else 60
        else:
            return self.timeout

    def new_chat_model(self) -> BaseChatModel:
        """Create a chat model from this config using init_chat_model for provider detection."""
        provider, name = _parse_model_provider(self.name)
        if provider == "openai" or self.api_base_url is not None:
            from langchain_openai import ChatOpenAI

            # Use the verbatim name if using an OAI server
            model_name = self.name if self.api_base_url is not None else name

            is_reasoning = _is_reasoning_model(model_name)
            extra_kwargs: dict[str, Any] = {}
            if self.use_responses_api:
                extra_kwargs.update(
                    # Without "summary", no reasoning traces will be returned by the API
                    reasoning={"effort": self.reasoning_effort, "summary": "auto"} if is_reasoning else None,
                    temperature=self.temperature,
                    # TODO output_version="responses/v1"
                )
            else:
                extra_kwargs.update(
                    reasoning_effort=self.reasoning_effort if is_reasoning else None,
                    # The old API errors out if you provide a temperature for reasoning models
                    temperature=self.temperature if not is_reasoning else None,
                )

            # Set a default API key for local models if the user didn't provide one
            if (
                self.api_base_url is not None
                and "api_key" not in self.model_kwargs
                and "OPENAI_API_KEY" not in os.environ
            ):
                extra_kwargs["api_key"] = "local-api-key"

            return ChatOpenAI(
                model=model_name,
                timeout=self._resolve_timeout(),
                max_tokens=self.max_tokens,
                base_url=self.api_base_url,
                use_responses_api=self.use_responses_api,
                **extra_kwargs,
                **self.model_kwargs,
            )
        elif provider == "anthropic":
            from langchain_anthropic import ChatAnthropic

            return ChatAnthropic(
                model_name=name,
                timeout=self._resolve_timeout(),
                temperature=self.temperature,
                max_tokens_to_sample=self.max_tokens,
                **self.model_kwargs,
            )

        if provider == "ollama" and self.ollama_pull_model:
            import ollama

            # Download with ollama. If the model already exists it will not be re-downloaded.
            ollama.pull(name)

        return init_chat_model(
            self.name,
            configurable_fields=None,  # Ensures we match the BaseChatModel overload
            temperature=self.temperature,
            max_tokens=self.max_tokens,
            timeout=self._resolve_timeout(),
            **self.model_kwargs,
        )

    @classmethod
    def from_yaml(cls, path: str | Path) -> Self:
        """Load an LLM config from a YAML file."""

        import yaml

        path = Path(path)
        if path.exists():
            content = path.read_text()
        else:
            raise ValueError(f"LLM config file {path} not found.")

        model_dict = yaml.safe_load(content)
        return cls.model_validate(model_dict)


def _is_reasoning_model(model_name: str) -> bool:
    """Check if a model is a reasoning model based on its name."""
    return any(prefix in model_name for prefix in _OPENAI_REASONING_INFIXES)


def _is_openai_model(model_name: str) -> bool:
    """Check if a model is an OpenAI model based on its name."""
    return any(model_name.startswith(prefix) for prefix in _OPENAI_PREFIXES)


def _is_anthropic_model(model_name: str) -> bool:
    """Check if a model is an Anthropic model based on its name."""
    return any(model_name.startswith(prefix) for prefix in _ANTHROPIC_PREFIXES)


def _parse_model_provider(model: str) -> tuple[str, str]:
    """Parse the provider and model name from a string of the form 'provider:name' or 'name'."""
    provider, sep, name = model.partition(":")
    if len(sep) == 0 and len(name) == 0:
        if _is_openai_model(model):
            return "openai", model
        elif _is_anthropic_model(model):
            return "anthropic", model
        else:
            return "", model
    return provider, name


class LLMConfigDirectory:
    """Namespace for preconfigured LLM configurations."""

    @classmethod
    def list_all(cls) -> list[LLMConfig]:
        return [config for name, config in vars(cls).items() if name.isupper()]

    DEFAULT = LLMConfig(name="gpt-4o-mini")

    # https://huggingface.co/Qwen/Qwen3-8B-GGUF#best-practices
    QWEN3_8B_OAI = LLMConfig(
        name="qwen/qwen3-8b",
        api_base_url="http://localhost:8080/v1",
        max_tokens=32768,
        temperature=0.6,
        use_responses_api=False,
        timeout=600,
    )

    # https://huggingface.co/Qwen/Qwen3-8B-GGUF#best-practices
    QWEN3_8B_OLLAMA = LLMConfig(
        name="ollama:qwen3:8b",
        max_tokens=32768,
        temperature=0.6,
        timeout=600,
        # Refer to https://python.langchain.com/api_reference/ollama/chat_models/langchain_ollama.chat_models.ChatOllama.html
        model_kwargs={
            "reasoning": True,
            "num_ctx": 40960,  # Override the global context size: https://docs.ollama.com/context-length
            "num_predict": 32768,
            "validate_model_on_init": True,
        },
    )