"""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, )