first-aid-rag-assistant / src / config.py
config.py
Raw
"""Central configuration: model name, LLM client factory, and cost pricing."""

from litellm import Router

# ── Model ────────────────────────────────────────────────────────────────────
MODEL_NAME = "ministral-3b-2512"

_MODEL_LIST = [
    {
        "model_name": MODEL_NAME,
        "litellm_params": {
            "model": "mistral/ministral-3b-2512",
        },
    }
]

# ── Pricing ($ per million tokens) ──────────────────────────────────────────
PRICE_PER_MILLION_INPUT = 0.1
PRICE_PER_MILLION_OUTPUT = 0.1


def make_llm_client() -> Router:
    """Return a new LiteLLM Router configured for the project model."""
    return Router(model_list=_MODEL_LIST)


def calculate_cost(model: str, usage) -> float:
    """Return the dollar cost for a single LLM call.

    Handles missing usage gracefully (returns 0).
    """
    if usage is None:
        return 0.0
    try:
        prompt_tokens = usage.prompt_tokens
        completion_tokens = usage.completion_tokens
    except AttributeError:
        return 0.0

    if MODEL_NAME in model:
        return (
            prompt_tokens * PRICE_PER_MILLION_INPUT
            + completion_tokens * PRICE_PER_MILLION_OUTPUT
        ) / 1_000_000

    return 0.0