Agentic-OS/backend/app/services/llm_usage.py

"""Track LLM calls per task: provider, model, tokens, and estimated cost."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from langchain_core.messages import BaseMessage
from langchain_openai import ChatOpenAI

# USD per 1M tokens (input, output). Local/Ollama is always zero.
MODEL_PRICING: dict[str, tuple[float, float]] = {
    "deepseek-chat": (0.14, 0.28),
    "deepseek/deepseek-chat": (0.14, 0.28),
    "qwen2.5:7b-instruct": (0.0, 0.0),
    "qwen2.5-coder:14b": (0.0, 0.0),
    # OpenRouter common fallbacks (approximate)
    "anthropic/claude-3.5-sonnet": (3.0, 15.0),
    "anthropic/claude-3.7-sonnet": (3.0, 15.0),
    "openai/gpt-4o-mini": (0.15, 0.60),
    "meta-llama/llama-3.1-70b-instruct": (0.52, 0.75),
}


@dataclass
class ModelInfo:
    step: str
    provider: str  # local | api
    backend: str  # ollama | gemini | deepseek | openrouter
    model: str
    display: str


@dataclass
class LlmUsageTracker:
    """Collects LLM call records for one agent run."""

    calls: list[dict[str, Any]] = field(default_factory=list)

    async def invoke(
        self,
        step: str,
        model: ChatOpenAI,
        info: ModelInfo,
        messages: list[BaseMessage],
    ) -> Any:
        response = await model.ainvoke(messages)
        usage = extract_usage(response, info)
        usage["step"] = step
        self.calls.append(usage)
        return response

    def summary(self) -> dict[str, Any]:
        run_cost = round(sum(c.get("cost_usd", 0.0) for c in self.calls), 6)
        return {
            "llm_usage": self.calls,
            "run_cost_usd": run_cost,
            "total_input_tokens": sum(c.get("input_tokens", 0) for c in self.calls),
            "total_output_tokens": sum(c.get("output_tokens", 0) for c in self.calls),
        }


def extract_usage(response: Any, info: ModelInfo) -> dict[str, Any]:
    meta = getattr(response, "response_metadata", None) or {}
    token_usage = meta.get("token_usage") or {}
    if not token_usage:
        usage_meta = getattr(response, "usage_metadata", None)
        if usage_meta:
            token_usage = {
                "prompt_tokens": usage_meta.get("input_tokens", 0),
                "completion_tokens": usage_meta.get("output_tokens", 0),
                "total_tokens": usage_meta.get("total_tokens", 0),
            }

    input_tokens = int(token_usage.get("prompt_tokens") or token_usage.get("input_tokens") or 0)
    output_tokens = int(
        token_usage.get("completion_tokens") or token_usage.get("output_tokens") or 0
    )
    model_name = meta.get("model_name") or info.model
    cost = estimate_cost(model_name, info.backend, input_tokens, output_tokens)

    return {
        "provider": info.provider,
        "backend": info.backend,
        "model": model_name,
        "display": info.display,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": input_tokens + output_tokens,
        "cost_usd": cost,
    }


def estimate_cost(model: str, backend: str, input_tokens: int, output_tokens: int) -> float:
    if backend == "ollama":
        return 0.0
    from app.services.model_catalog import CATALOG, entry_for_model

    key = model.removeprefix("openrouter/")
    for e in CATALOG:
        if e.backend == backend and (e.model == model or e.model == key):
            in_price, out_price = e.input_per_m, e.output_per_m
            return round((input_tokens * in_price + output_tokens * out_price) / 1_000_000, 6)
    entry = entry_for_model(backend, model, "economy")
    in_price, out_price = entry.input_per_m, entry.output_per_m
    return round((input_tokens * in_price + output_tokens * out_price) / 1_000_000, 6)


def merge_run_history(
    prior_runs: list[dict],
    run_number: int,
    run_report: dict,
    llm_summary: dict,
) -> dict:
    """Build cumulative report with per-run and total cost."""
    run_entry = {
        "run_number": run_number,
        "follow_up": run_report.get("follow_up"),
        "summary": run_report.get("summary"),
        "resolved": run_report.get("resolved"),
        "devices_checked": run_report.get("devices_checked"),
        "llm_usage": llm_summary.get("llm_usage", []),
        "run_cost_usd": llm_summary.get("run_cost_usd", 0.0),
        "generated_at": run_report.get("generated_at"),
    }
    all_runs = prior_runs + [run_entry]
    total_cost = round(sum(r.get("run_cost_usd", 0.0) for r in all_runs), 6)

    merged = {**run_report}
    merged["run_number"] = run_number
    merged["runs"] = all_runs
    merged["llm_usage"] = llm_summary.get("llm_usage", [])
    merged["run_cost_usd"] = llm_summary.get("run_cost_usd", 0.0)
    merged["total_cost_usd"] = total_cost
    merged["total_input_tokens"] = sum(
        t.get("input_tokens", 0) for r in all_runs for t in r.get("llm_usage", [])
    )
    merged["total_output_tokens"] = sum(
        t.get("output_tokens", 0) for r in all_runs for t in r.get("llm_usage", [])
    )
    return merged