115 lines
2.9 KiB
Python
115 lines
2.9 KiB
Python
"""LLM Protocol - 数据类与抽象基类"""
|
||
|
||
from abc import ABC, abstractmethod
|
||
from dataclasses import dataclass, field
|
||
|
||
|
||
@dataclass
|
||
class TokenUsage:
|
||
"""Token 使用量"""
|
||
|
||
prompt_tokens: int = 0
|
||
completion_tokens: int = 0
|
||
|
||
@property
|
||
def total_tokens(self) -> int:
|
||
return self.prompt_tokens + self.completion_tokens
|
||
|
||
|
||
@dataclass
|
||
class ToolCall:
|
||
"""工具调用"""
|
||
|
||
id: str
|
||
name: str
|
||
arguments: dict[str, object]
|
||
|
||
|
||
@dataclass
|
||
class LLMRequest:
|
||
"""LLM 请求"""
|
||
|
||
messages: list[dict[str, str]]
|
||
model: str
|
||
tools: list[dict[str, object]] | None = None
|
||
tool_choice: str = "auto"
|
||
temperature: float = 0.7
|
||
max_tokens: int = 2000
|
||
timeout: float | None = None
|
||
|
||
def __init__(
|
||
self,
|
||
messages: list[dict[str, str]],
|
||
model: str,
|
||
tools: list[dict[str, object]] | None = None,
|
||
tool_choice: str = "auto",
|
||
temperature: float = 0.7,
|
||
max_tokens: int = 2000,
|
||
timeout: float | None = None,
|
||
cache: dict[str, object] | None = None,
|
||
**kwargs: object,
|
||
):
|
||
self.messages = messages
|
||
self.model = model
|
||
self.tools = tools
|
||
self.tool_choice = tool_choice
|
||
self.temperature = temperature
|
||
self.max_tokens = max_tokens
|
||
self.timeout = timeout
|
||
self._extra = kwargs
|
||
# U17 — LiteLLM cache 参数（cache_key 或 no-cache），透传到 litellm.acompletion
|
||
self._cache: dict[str, object] | None = cache
|
||
|
||
|
||
@dataclass
|
||
class StreamChunk:
|
||
"""LLM 流式响应块"""
|
||
|
||
content: str # Delta content
|
||
model: str
|
||
tool_calls: list[ToolCall] = field(
|
||
default_factory=list
|
||
) # Accumulated tool calls (only in final chunk)
|
||
usage: TokenUsage | None = None # Only in final chunk
|
||
is_final: bool = False # True for the last chunk
|
||
|
||
|
||
@dataclass
|
||
class LLMResponse:
|
||
"""LLM 响应"""
|
||
|
||
content: str
|
||
model: str
|
||
usage: TokenUsage
|
||
tool_calls: list[ToolCall] = field(default_factory=list)
|
||
latency_ms: float = 0.0
|
||
cache_hit: bool = False # U17 — 缓存命中标记，用于 usage tracking（cost=0）
|
||
|
||
@property
|
||
def has_tool_calls(self) -> bool:
|
||
return len(self.tool_calls) > 0
|
||
|
||
|
||
class LLMProvider(ABC):
|
||
"""LLM Provider 抽象基类"""
|
||
|
||
@abstractmethod
|
||
async def chat(self, request: LLMRequest) -> LLMResponse:
|
||
"""发送 chat 请求并返回响应"""
|
||
...
|
||
|
||
async def chat_stream(self, request: LLMRequest):
|
||
"""Stream chat response. Override in subclasses that support streaming.
|
||
|
||
Yields StreamChunk objects. Default implementation falls back to
|
||
non-streaming chat and yields a single chunk.
|
||
"""
|
||
response = await self.chat(request)
|
||
yield StreamChunk(
|
||
content=response.content,
|
||
model=response.model,
|
||
tool_calls=response.tool_calls,
|
||
usage=response.usage,
|
||
is_final=True,
|
||
)
|