268 lines
9.9 KiB
Python
268 lines
9.9 KiB
Python
"""LLM Gateway - 统一 LLM 调用入口"""
|
|
|
|
import logging
|
|
import time
|
|
|
|
from agentkit.core.exceptions import LLMProviderError, ModelNotFoundError
|
|
from agentkit.llm.config import LLMConfig
|
|
from agentkit.llm.protocol import LLMProvider, LLMRequest, LLMResponse, StreamChunk, TokenUsage
|
|
from agentkit.llm.providers.tracker import UsageSummary, UsageTracker
|
|
from agentkit.telemetry.tracing import get_tracer, _OTEL_AVAILABLE
|
|
from agentkit.telemetry.metrics import llm_token_histogram
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class LLMGateway:
|
|
"""LLM 网关 - Provider 注册、模型别名解析、Fallback、Usage 追踪"""
|
|
|
|
def __init__(self, config: LLMConfig | None = None):
|
|
self._providers: dict[str, LLMProvider] = {}
|
|
self._usage_tracker = UsageTracker()
|
|
self._config = config or LLMConfig()
|
|
|
|
def register_provider(self, name: str, provider: LLMProvider) -> None:
|
|
"""注册 Provider"""
|
|
self._providers[name] = provider
|
|
logger.info(f"LLM provider '{name}' registered")
|
|
|
|
@property
|
|
def has_providers(self) -> bool:
|
|
"""Return True if at least one LLM provider is registered."""
|
|
return bool(self._providers)
|
|
|
|
async def chat(
|
|
self,
|
|
messages: list[dict[str, str]],
|
|
model: str,
|
|
agent_name: str = "",
|
|
task_type: str = "",
|
|
tools: list[dict] | None = None,
|
|
tool_choice: str = "auto",
|
|
**kwargs,
|
|
) -> LLMResponse:
|
|
"""发送 chat 请求，自动解析别名和 Fallback"""
|
|
resolved_model = self._resolve_model_alias(model)
|
|
|
|
if not self._providers:
|
|
raise LLMProviderError("", "No provider registered")
|
|
|
|
# Telemetry: start LLM span
|
|
_span_cm = None
|
|
_span = None
|
|
if _OTEL_AVAILABLE:
|
|
tracer = get_tracer()
|
|
if tracer is not None:
|
|
from opentelemetry.trace import SpanKind
|
|
_span_cm = tracer.start_as_current_span(
|
|
"gen_ai.chat",
|
|
kind=SpanKind.CLIENT,
|
|
attributes={
|
|
"gen_ai.system": resolved_model.split("/")[0] if "/" in resolved_model else "unknown",
|
|
"gen_ai.operation.name": "chat",
|
|
"gen_ai.request.model": resolved_model,
|
|
},
|
|
)
|
|
_span = _span_cm.__enter__()
|
|
|
|
start = time.monotonic()
|
|
models_to_try = self._get_models_to_try(resolved_model)
|
|
last_error: LLMProviderError | None = None
|
|
|
|
try:
|
|
for model_name in models_to_try:
|
|
try:
|
|
provider, actual_model = self._resolve_model(model_name)
|
|
except ModelNotFoundError:
|
|
continue
|
|
|
|
req = LLMRequest(
|
|
messages=messages,
|
|
model=actual_model,
|
|
tools=tools,
|
|
tool_choice=tool_choice,
|
|
**kwargs,
|
|
)
|
|
try:
|
|
response = await provider.chat(req)
|
|
break
|
|
except LLMProviderError as e:
|
|
last_error = e
|
|
logger.warning(f"Model '{model_name}' failed, trying next: {e}")
|
|
continue
|
|
else:
|
|
raise last_error or LLMProviderError("", f"All models failed for '{resolved_model}'")
|
|
|
|
latency_ms = (time.monotonic() - start) * 1000
|
|
|
|
# 计算成本
|
|
cost = self._calculate_cost(response.model, response.usage)
|
|
|
|
# 记录使用量
|
|
self._usage_tracker.record(
|
|
agent_name=agent_name,
|
|
model=response.model,
|
|
usage=response.usage,
|
|
cost=cost,
|
|
latency_ms=latency_ms,
|
|
)
|
|
|
|
# Telemetry: record token usage and end span
|
|
if _span is not None:
|
|
_span.set_attribute("gen_ai.usage.input_tokens", response.usage.prompt_tokens)
|
|
_span.set_attribute("gen_ai.usage.output_tokens", response.usage.completion_tokens)
|
|
_span.set_attribute("gen_ai.response.model", response.model)
|
|
_span.set_attribute("gen_ai.duration_ms", int(latency_ms))
|
|
llm_token_histogram().record(
|
|
response.usage.total_tokens,
|
|
{"gen_ai.request.model": resolved_model},
|
|
)
|
|
|
|
return response
|
|
finally:
|
|
if _span_cm is not None:
|
|
_span_cm.__exit__(None, None, None)
|
|
|
|
async def chat_stream(
|
|
self,
|
|
messages: list[dict[str, str]],
|
|
model: str,
|
|
agent_name: str = "",
|
|
task_type: str = "",
|
|
tools: list[dict] | None = None,
|
|
tool_choice: str = "auto",
|
|
**kwargs,
|
|
):
|
|
"""Stream chat response with fallback support.
|
|
|
|
If the primary model fails before any chunk is yielded, tries fallback
|
|
models. If it fails after chunks have been sent, yields an error chunk
|
|
and terminates (cannot switch mid-stream).
|
|
"""
|
|
resolved_model = self._resolve_model_alias(model)
|
|
|
|
if not self._providers:
|
|
raise LLMProviderError("", "No provider registered")
|
|
|
|
models_to_try = self._get_models_to_try(resolved_model)
|
|
last_error: Exception | None = None
|
|
|
|
for model_name in models_to_try:
|
|
try:
|
|
provider, actual_model = self._resolve_model(model_name)
|
|
except ModelNotFoundError:
|
|
continue
|
|
|
|
stream_request = LLMRequest(
|
|
messages=messages,
|
|
model=actual_model,
|
|
tools=tools,
|
|
tool_choice=tool_choice,
|
|
**kwargs,
|
|
)
|
|
|
|
chunk_yielded = False
|
|
start = time.monotonic()
|
|
total_content = ""
|
|
final_usage = None
|
|
final_model = model_name
|
|
|
|
try:
|
|
async for chunk in provider.chat_stream(stream_request):
|
|
chunk_yielded = True
|
|
if chunk.content:
|
|
total_content += chunk.content
|
|
if chunk.usage:
|
|
final_usage = chunk.usage
|
|
if chunk.model:
|
|
final_model = chunk.model
|
|
yield chunk
|
|
|
|
# Track usage after successful stream
|
|
latency_ms = (time.monotonic() - start) * 1000
|
|
if final_usage is None:
|
|
final_usage = TokenUsage()
|
|
cost = self._calculate_cost(final_model, final_usage)
|
|
self._usage_tracker.record(
|
|
agent_name=agent_name,
|
|
model=final_model,
|
|
usage=final_usage,
|
|
cost=cost,
|
|
latency_ms=latency_ms,
|
|
)
|
|
return # Success, done
|
|
except Exception as e:
|
|
last_error = e
|
|
if chunk_yielded:
|
|
# Can't switch mid-stream, terminate gracefully
|
|
logger.error(f"Stream failed after chunks sent for '{model_name}': {e}")
|
|
yield StreamChunk(
|
|
content="",
|
|
model=final_model,
|
|
usage=None,
|
|
is_final=True,
|
|
)
|
|
return
|
|
# No chunks yet, try next fallback
|
|
logger.warning(f"Stream failed for '{model_name}', trying fallback: {e}")
|
|
continue
|
|
|
|
# All models failed
|
|
raise last_error or LLMProviderError("", f"No provider available for streaming '{resolved_model}'")
|
|
|
|
def _get_models_to_try(self, resolved_model: str) -> list[str]:
|
|
"""Return [primary_model] + fallback_models for the given resolved model."""
|
|
fallback_models = self._config.fallbacks.get(resolved_model, [])
|
|
return [resolved_model] + fallback_models
|
|
|
|
def _resolve_model_alias(self, model: str) -> str:
|
|
"""解析模型别名"""
|
|
if model in self._config.model_aliases:
|
|
return self._config.model_aliases[model]
|
|
return model
|
|
|
|
def _resolve_model(self, model: str) -> tuple[LLMProvider, str]:
|
|
"""解析模型为 (provider, actual_model_name)"""
|
|
# model 格式: "provider/model_name" 或 "model_name"
|
|
if "/" in model:
|
|
provider_name, model_name = model.split("/", 1)
|
|
if provider_name not in self._providers:
|
|
raise ModelNotFoundError(model)
|
|
return self._providers[provider_name], model_name
|
|
|
|
# 无 "/" 前缀：仅当只有一个 provider 时自动匹配
|
|
if len(self._providers) == 1:
|
|
provider = next(iter(self._providers.values()))
|
|
return provider, model
|
|
|
|
raise ModelNotFoundError(model)
|
|
|
|
def _get_fallback_model(self, model: str) -> str | None:
|
|
"""获取 Fallback 模型"""
|
|
fallbacks = self._config.fallbacks.get(model, [])
|
|
return fallbacks[0] if fallbacks else None
|
|
|
|
def _calculate_cost(self, model: str, usage: TokenUsage) -> float:
|
|
"""计算成本"""
|
|
# 在 provider config 的 models 中查找成本配置
|
|
for provider_config in self._config.providers.values():
|
|
if model in provider_config.models:
|
|
model_conf = provider_config.models[model]
|
|
input_cost = usage.prompt_tokens * model_conf.get("cost_per_1k_input", 0) / 1000
|
|
output_cost = usage.completion_tokens * model_conf.get("cost_per_1k_output", 0) / 1000
|
|
return input_cost + output_cost
|
|
return 0.0
|
|
|
|
def get_usage(
|
|
self,
|
|
agent_name: str | None = None,
|
|
start_time=None,
|
|
end_time=None,
|
|
) -> UsageSummary:
|
|
"""查询使用量"""
|
|
return self._usage_tracker.get_usage(
|
|
agent_name=agent_name,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
)
|