agent-platform/agent_platform/llm.py

"""LiteLLM-Wrapper mit Token-Tracking."""
import litellm
from typing import List, Dict, Optional
from config import settings


if settings.LLM_API_KEY:
    litellm.api_key = settings.LLM_API_KEY
if settings.LLM_API_BASE:
    litellm.api_base = settings.LLM_API_BASE


class LLMResponse:
    def __init__(self, content: str, tool_calls: Optional[List[Dict]] = None, usage: Optional[Dict] = None):
        self.content = content
        self.tool_calls = tool_calls or []
        self.usage = usage or {}

    @property
    def input_tokens(self) -> int:
        return self.usage.get("prompt_tokens", 0)

    @property
    def output_tokens(self) -> int:
        return self.usage.get("completion_tokens", 0)

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens

    @property
    def finish_reason(self) -> str:
        return self.usage.get("finish_reason", "stop")


async def chat(
    messages: List[Dict[str, str]],
    tools: Optional[List[Dict]] = None,
    model: Optional[str] = None,
    temperature: float = 0.7,
    max_tokens: int = 2000,
) -> LLMResponse:
    """LLM-Call via LiteLLM.

    messages: Liste von {role, content} Dicts.
    tools: Optional, MCP-Tool-Definitionen für Function-Calling.
    """
    model = model or settings.LLM_MODEL

    kwargs = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }

    if tools:
        kwargs["tools"] = tools

    response = await litellm.acompletion(**kwargs)

    message = response.choices[0].message

    tool_calls = []
    if hasattr(message, "tool_calls") and message.tool_calls:
        for tc in message.tool_calls:
            tool_calls.append({
                "id": tc.id,
                "name": tc.function.name,
                "arguments": tc.function.arguments,
            })

    usage = {}
    if hasattr(response, "usage") and response.usage:
        # LiteLLM exposes cached prompt tokens via prompt_tokens_details.cached_tokens
        # (see litellm.types.utils.Usage). Providers like Anthropic/OpenAI surface this
        # when prompt caching is active; surfacing it here flows through LLMResponse.usage
        # to any audit log that serialises the usage dict.
        cached_tokens = 0
        details = getattr(response.usage, "prompt_tokens_details", None)
        if details is not None:
            cached_tokens = getattr(details, "cached_tokens", None) or 0

        usage = {
            "prompt_tokens": response.usage.prompt_tokens or 0,
            "completion_tokens": response.usage.completion_tokens or 0,
            "total_tokens": response.usage.total_tokens or 0,
            "cached_tokens": cached_tokens,
            "finish_reason": response.choices[0].finish_reason or "stop",
        }

    return LLMResponse(
        content=message.content or "",
        tool_calls=tool_calls,
        usage=usage,
    )