"""OpenAI Compatible Provider - 支持 OpenAI/DeepSeek/Anthropic 等兼容 API""" import json import logging import time import httpx from agentkit.core.exceptions import LLMProviderError from agentkit.llm.protocol import LLMProvider, LLMRequest, LLMResponse, StreamChunk, TokenUsage, ToolCall from agentkit.llm.retry import ( CircuitBreaker, CircuitBreakerConfig, RetryConfig, RetryPolicy, ) logger = logging.getLogger(__name__) class _StreamContext: """Wraps an httpx streaming response context manager for use with retry/circuit breaker. The ``__aenter__`` returns the httpx response so callers can use ``async with ctx as response:`` naturally. """ def __init__(self, response_ctx, response): self._response_ctx = response_ctx self._response = response async def __aenter__(self): return self._response async def __aexit__(self, exc_type, exc_val, exc_tb): return await self._response_ctx.__aexit__(exc_type, exc_val, exc_tb) class OpenAICompatibleProvider(LLMProvider): """OpenAI 兼容 API Provider""" def __init__( self, api_key: str, base_url: str = "https://api.openai.com/v1", default_model: str = "gpt-4o-mini", retry_config: RetryConfig | None = None, circuit_breaker_config: CircuitBreakerConfig | None = None, ): self._api_key = api_key self._base_url = base_url.rstrip("/") self._default_model = default_model self._client = httpx.AsyncClient(timeout=60.0) self._retry_policy = RetryPolicy(retry_config) if retry_config else None self._circuit_breaker = ( CircuitBreaker(circuit_breaker_config, provider="openai") if circuit_breaker_config else None ) async def close(self) -> None: """关闭 HTTP 客户端连接池""" await self._client.aclose() async def chat(self, request: LLMRequest) -> LLMResponse: """发送 chat 请求（带 retry + circuit breaker）""" if self._circuit_breaker and self._retry_policy: return await self._circuit_breaker.execute( self._retry_policy.execute, self._chat_impl, request ) if self._retry_policy: return await self._retry_policy.execute(self._chat_impl, request) if self._circuit_breaker: return await self._circuit_breaker.execute(self._chat_impl, request) return await self._chat_impl(request) async def _chat_impl(self, request: LLMRequest) -> LLMResponse: """发送 chat 请求""" url = f"{self._base_url}/chat/completions" headers = { "Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json", } payload: dict = { "model": request.model, "messages": request.messages, "temperature": request.temperature, "max_tokens": request.max_tokens, } if request.tools: payload["tools"] = request.tools payload["tool_choice"] = request.tool_choice logger.debug(f"Chat request to {url}: model={request.model}, messages={len(request.messages)}, tools={len(request.tools or [])}") start = time.monotonic() try: resp = await self._client.post(url, json=payload, headers=headers) except httpx.HTTPError as e: raise LLMProviderError("openai", str(e)) from e latency_ms = (time.monotonic() - start) * 1000 if resp.status_code != 200: try: error_body = resp.json() error_msg = error_body.get("error", {}).get("message", "Request failed") except Exception: error_msg = f"HTTP {resp.status_code}" logger.error(f"Chat request failed: HTTP {resp.status_code}, error: {error_msg}") # 不在错误消息中暴露完整响应体，防止 API Key 泄露 raise LLMProviderError("openai", f"HTTP {resp.status_code}: {error_msg}") data = resp.json() choice = data["choices"][0] message = choice["message"] usage_data = data.get("usage", {}) usage = TokenUsage( prompt_tokens=usage_data.get("prompt_tokens", 0), completion_tokens=usage_data.get("completion_tokens", 0), ) tool_calls: list[ToolCall] = [] raw_tool_calls = message.get("tool_calls") if raw_tool_calls: for tc in raw_tool_calls: func = tc["function"] arguments = json.loads(func["arguments"]) if isinstance(func["arguments"], str) else func["arguments"] tool_calls.append( ToolCall( id=tc["id"], name=func["name"], arguments=arguments, ) ) content = message.get("content") or "" return LLMResponse( content=content, model=data.get("model", request.model), usage=usage, tool_calls=tool_calls, latency_ms=latency_ms, ) async def chat_stream(self, request: LLMRequest): """Stream chat response using SSE（带 retry + circuit breaker）""" # For streaming, retry/circuit breaker only protect the connection phase. # Once the stream is open, we iterate without retry. if self._circuit_breaker and self._retry_policy: ctx = await self._circuit_breaker.execute( self._retry_policy.execute, self._open_stream, request ) elif self._retry_policy: ctx = await self._retry_policy.execute(self._open_stream, request) elif self._circuit_breaker: ctx = await self._circuit_breaker.execute(self._open_stream, request) else: ctx = await self._open_stream(request) async with ctx as response: async for chunk in self._iterate_stream(response, request): yield chunk async def _open_stream(self, request: LLMRequest): """Open the streaming HTTP connection; returns a _StreamContext.""" url = f"{self._base_url}/chat/completions" headers = { "Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json", } payload: dict = { "model": request.model, "messages": request.messages, "temperature": request.temperature, "max_tokens": request.max_tokens, "stream": True, } if request.tools: payload["tools"] = request.tools payload["tool_choice"] = request.tool_choice tool_names = [t.get("function", {}).get("name", "?") for t in request.tools] logger.info(f"OpenAIProvider stream: model={request.model}, tools={len(request.tools)} {tool_names}") else: logger.info(f"OpenAIProvider stream: model={request.model}, NO tools") response_ctx = self._client.stream("POST", url, json=payload, headers=headers) response = await response_ctx.__aenter__() if response.status_code != 200: await response.aread() await response_ctx.__aexit__(None, None, None) # Parse error body for detailed message try: error_body = response.json() error_msg = error_body.get("error", {}).get("message", f"HTTP {response.status_code}") except Exception: error_msg = f"HTTP {response.status_code}" logger.error(f"Stream request failed: HTTP {response.status_code}, error: {error_msg}") raise LLMProviderError("openai", f"HTTP {response.status_code}: {error_msg}") return _StreamContext(response_ctx, response) async def _iterate_stream(self, response, request: LLMRequest): """Iterate over an already-open SSE stream and yield StreamChunks.""" accumulated_tool_calls: dict[int, dict] = {} # index -> {id, name, arguments_str} async for line in response.aiter_lines(): line = line.strip() if not line or not line.startswith("data: "): continue data_str = line[6:] # Remove "data: " prefix if data_str == "[DONE]": break try: data = json.loads(data_str) except json.JSONDecodeError: continue choices = data.get("choices", []) if not choices: # Usage-only chunk usage_data = data.get("usage") if usage_data: yield StreamChunk( content="", model=data.get("model", request.model), usage=TokenUsage( prompt_tokens=usage_data.get("prompt_tokens", 0), completion_tokens=usage_data.get("completion_tokens", 0), ), is_final=True, ) continue delta = choices[0].get("delta", {}) content = delta.get("content", "") # Accumulate tool calls from streaming raw_tool_calls = delta.get("tool_calls") if raw_tool_calls: for tc in raw_tool_calls: idx = tc.get("index", 0) if idx not in accumulated_tool_calls: accumulated_tool_calls[idx] = { "id": tc.get("id", ""), "name": "", "arguments_str": "", } if tc.get("id"): accumulated_tool_calls[idx]["id"] = tc["id"] func = tc.get("function", {}) if func.get("name"): accumulated_tool_calls[idx]["name"] = func["name"] if func.get("arguments"): accumulated_tool_calls[idx]["arguments_str"] += func["arguments"] # Only yield content chunks (not empty deltas) if content: yield StreamChunk( content=content, model=data.get("model", request.model), ) # If we accumulated tool calls, yield them as a final chunk if accumulated_tool_calls: tool_calls = [] for idx in sorted(accumulated_tool_calls.keys()): tc_data = accumulated_tool_calls[idx] try: arguments = json.loads(tc_data["arguments_str"]) if tc_data["arguments_str"] else {} except json.JSONDecodeError: arguments = {"raw": tc_data["arguments_str"]} tool_calls.append(ToolCall( id=tc_data["id"], name=tc_data["name"], arguments=arguments, )) yield StreamChunk( content="", model=request.model, tool_calls=tool_calls, is_final=True, )