chat_backend/llm_be/chat_backend/ollama_config.py

"""Shared Ollama client helpers — always use settings.OLLAMA_BASE_URL.

Role-scoped models (#62): THINKING / FAST / UTILITY / EMBED each resolve
independently, with ``OLLAMA_MODEL`` kept as a fallback so existing deploys
keep working until they set the role-specific vars.
"""

from __future__ import annotations

from django.conf import settings

ROLE_THINKING = "thinking"
ROLE_FAST = "fast"
ROLE_UTILITY = "utility"
ROLE_EMBED = "embed"
# Agentic task execution (#63): orchestrator plans/synthesises; sub-agents run
# independent plan steps concurrently on a smaller/cheaper model.
ROLE_ORCHESTRATOR = "orchestrator"
ROLE_SUBAGENT = "subagent"

_VALID_ROLES = {
    ROLE_THINKING,
    ROLE_FAST,
    ROLE_UTILITY,
    ROLE_EMBED,
    ROLE_ORCHESTRATOR,
    ROLE_SUBAGENT,
}


def ollama_base_url() -> str:
    return getattr(settings, "OLLAMA_BASE_URL", "http://127.0.0.1:11434")


def ollama_model(default: str | None = None) -> str:
    """Legacy single-model accessor. Prefer :func:`ollama_model_for_role`."""
    if default:
        return default
    return getattr(settings, "OLLAMA_MODEL", "gpt-oss:20b")


def ollama_model_for_role(role: str) -> str:
    """Resolve the model name for a generation role.

    Lookup order: role-specific setting → ``OLLAMA_MODEL`` fallback →
    hard-coded role default (never falls back from embed → chat model).
    """
    role = (role or ROLE_THINKING).lower()
    if role not in _VALID_ROLES:
        raise ValueError(f"Unknown Ollama role: {role!r}")

    role_setting = {
        ROLE_THINKING: "OLLAMA_MODEL_THINKING",
        ROLE_FAST: "OLLAMA_MODEL_FAST",
        ROLE_UTILITY: "OLLAMA_MODEL_UTILITY",
        ROLE_EMBED: "OLLAMA_EMBED_MODEL",
        ROLE_ORCHESTRATOR: "OLLAMA_MODEL_ORCHESTRATOR",
        ROLE_SUBAGENT: "OLLAMA_MODEL_SUBAGENT",
    }[role]
    role_default = {
        ROLE_THINKING: "gpt-oss:20b",
        ROLE_FAST: "gemma4:latest",
        ROLE_UTILITY: "llama3.2",
        ROLE_EMBED: "nomic-embed-text",
        ROLE_ORCHESTRATOR: "gpt-oss:20b",
        ROLE_SUBAGENT: "llama3.2",
    }[role]

    configured = getattr(settings, role_setting, None)
    if configured:
        return configured

    # Embeddings must never silently fall back to a chat model (#62).
    if role == ROLE_EMBED:
        return role_default

    legacy = getattr(settings, "OLLAMA_MODEL", None)
    if legacy:
        return legacy
    return role_default


def ollama_embed_model() -> str:
    return ollama_model_for_role(ROLE_EMBED)


def ollama_num_ctx_for_role(role: str) -> int:
    role = (role or ROLE_THINKING).lower()
    if role == ROLE_FAST:
        return int(getattr(settings, "OLLAMA_NUM_CTX_FAST", 8192) or 8192)
    if role == ROLE_UTILITY:
        return int(getattr(settings, "OLLAMA_NUM_CTX_UTILITY", 4096) or 4096)
    if role == ROLE_SUBAGENT:
        return int(getattr(settings, "OLLAMA_NUM_CTX_SUBAGENT", 8192) or 8192)
    if role == ROLE_ORCHESTRATOR:
        return int(getattr(settings, "OLLAMA_NUM_CTX_ORCHESTRATOR", 16384) or 16384)
    return int(getattr(settings, "OLLAMA_NUM_CTX_THINKING", 16384) or 16384)


def resolve_chat_role(model_name: str | None) -> str:
    """Map FE mode selector (FAST / THINKING / …) to an Ollama role."""
    if (model_name or "").upper() == "FAST":
        return ROLE_FAST
    return ROLE_THINKING


def ollama_llm_kwargs(role: str = ROLE_THINKING, **extra):
    """Keyword args for langchain_ollama.OllamaLLM / ChatOllama."""
    kwargs = {
        "base_url": ollama_base_url(),
        "model": ollama_model_for_role(role),
        "num_ctx": ollama_num_ctx_for_role(role),
    }
    kwargs.update(extra)
    return kwargs


def ollama_embeddings_kwargs(**extra):
    kwargs = {"base_url": ollama_base_url(), "model": ollama_embed_model()}
    kwargs.update(extra)
    return kwargs