implementation_engineer
97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
"""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,
|
|
)
|