File size: 1,303 Bytes
dc893fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | from enum import Enum
from typing import Any
from pydantic import BaseModel
class LLMProvider(str, Enum):
"""LLM provider types."""
ANTHROPIC = "anthropic"
OPENAI = "openai"
class FunctionCall(BaseModel):
"""Function call details."""
name: str
arguments: dict[str, Any] # Function arguments as dict
class ToolCall(BaseModel):
"""Tool call structure."""
id: str
type: str # "function"
function: FunctionCall
class Message(BaseModel):
"""Chat message."""
role: str # "system", "user", "assistant", "tool"
content: str | list[dict[str, Any]] # Can be string or list of content blocks
thinking: str | None = None # Extended thinking content for assistant messages
tool_calls: list[ToolCall] | None = None
tool_call_id: str | None = None
name: str | None = None # For tool role
class TokenUsage(BaseModel):
"""Token usage statistics from LLM API response."""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
class LLMResponse(BaseModel):
"""LLM response."""
content: str
thinking: str | None = None # Extended thinking blocks
tool_calls: list[ToolCall] | None = None
finish_reason: str
usage: TokenUsage | None = None # Token usage from API response
|