252 lines
8.9 KiB
Python
252 lines
8.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from app.ai_provider.config.models import ModelConfig, ProviderConfig
|
|
from app.ai_provider.pipeline import ParameterPipeline, PipelineContext, PipelineFactory
|
|
from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
|
|
from app.ai_provider.types.errors import AIProviderError
|
|
from app.ai_provider.types.messages import AIMessage, BaseMessage, ToolMessage
|
|
from app.core.logging import logger
|
|
|
|
|
|
class OpenAIFormatAdapter:
|
|
"""兼容 OpenAI API 格式的适配器（覆盖 deepseek、xiaomi 等）。"""
|
|
|
|
def __init__(
|
|
self,
|
|
provider_config: ProviderConfig,
|
|
pipeline: ParameterPipeline | None = None,
|
|
) -> None:
|
|
self._provider = provider_config
|
|
self._pipeline = pipeline or PipelineFactory.create(provider_config)
|
|
self._client: httpx.AsyncClient | None = None
|
|
|
|
@property
|
|
def _http(self) -> httpx.AsyncClient:
|
|
if self._client is None:
|
|
self._client = httpx.AsyncClient(
|
|
base_url=self._provider.base_url or "",
|
|
timeout=self._provider.timeout_seconds,
|
|
)
|
|
return self._client
|
|
|
|
def _convert_message(self, msg: BaseMessage) -> dict[str, Any]:
|
|
if isinstance(msg, ToolMessage):
|
|
return {
|
|
"role": "tool",
|
|
"content": msg.content,
|
|
"tool_call_id": msg.tool_call_id,
|
|
}
|
|
if isinstance(msg, AIMessage):
|
|
d: dict[str, Any] = {"role": "assistant", "content": msg.content}
|
|
if msg.tool_calls:
|
|
d["tool_calls"] = [
|
|
{
|
|
"id": tc.get("id", f"call_{i}"),
|
|
"type": "function",
|
|
"function": {
|
|
"name": tc["name"],
|
|
"arguments": json.dumps(tc.get("args", {})),
|
|
},
|
|
}
|
|
for i, tc in enumerate(msg.tool_calls)
|
|
]
|
|
return d
|
|
return {"role": msg.role, "content": msg.content}
|
|
|
|
def build_request(
|
|
self,
|
|
model: ModelConfig,
|
|
messages: list[BaseMessage],
|
|
params: dict[str, Any] | None = None,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
stream: bool = False,
|
|
) -> dict[str, Any]:
|
|
ctx = PipelineContext(
|
|
params=params or {},
|
|
model_config=model,
|
|
provider_config=self._provider,
|
|
)
|
|
pipe_result = self._pipeline.run(ctx)
|
|
|
|
body: dict[str, Any] = {
|
|
"model": model.model,
|
|
"messages": [self._convert_message(m) for m in messages],
|
|
"stream": stream,
|
|
}
|
|
body.update(pipe_result.params)
|
|
|
|
if tools:
|
|
body["tools"] = tools
|
|
if tool_choice:
|
|
body["tool_choice"] = tool_choice
|
|
|
|
return body
|
|
|
|
def bind_tools(
|
|
self,
|
|
request_body: dict[str, Any],
|
|
tools: list[dict[str, Any]],
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
) -> None:
|
|
request_body["tools"] = tools
|
|
if tool_choice:
|
|
request_body["tool_choice"] = tool_choice
|
|
|
|
def parse_chunk(self, raw: dict[str, Any]) -> ChatChunk:
|
|
chunk = ChatChunk()
|
|
choices = raw.get("choices", [])
|
|
if choices:
|
|
delta = choices[0].get("delta", {})
|
|
chunk.delta_content = delta.get("content") or ""
|
|
chunk.delta_reasoning = delta.get("reasoning_content") or ""
|
|
chunk.finish_reason = choices[0].get("finish_reason")
|
|
|
|
if delta.get("tool_calls"):
|
|
for tc in delta["tool_calls"]:
|
|
fn = tc.get("function", {})
|
|
chunk.tool_call_chunks.append(
|
|
{
|
|
"index": tc.get("index", 0),
|
|
"id": tc.get("id"),
|
|
"name": fn.get("name"),
|
|
"args": fn.get("arguments", ""),
|
|
}
|
|
)
|
|
usage = raw.get("usage")
|
|
if usage:
|
|
chunk.usage = TokenUsage(
|
|
prompt_tokens=usage.get("prompt_tokens", 0),
|
|
completion_tokens=usage.get("completion_tokens", 0),
|
|
total_tokens=usage.get("total_tokens", 0),
|
|
prompt_tokens_cached=usage.get("prompt_tokens_cached", 0),
|
|
)
|
|
return chunk
|
|
|
|
def parse_nonstreaming_response(
|
|
self,
|
|
raw: dict[str, Any],
|
|
model_name: str,
|
|
) -> ChatResult:
|
|
choice = raw.get("choices", [{}])[0]
|
|
message = choice.get("message", {})
|
|
content = message.get("content", "")
|
|
tool_calls_raw = message.get("tool_calls", [])
|
|
|
|
result = ChatResult(
|
|
content=content,
|
|
reasoning=message.get("reasoning_content", ""),
|
|
model_name=model_name,
|
|
)
|
|
if tool_calls_raw:
|
|
result.tool_calls = [
|
|
{
|
|
"name": tc["function"]["name"],
|
|
"args": json.loads(tc["function"].get("arguments", "{}")),
|
|
"id": tc["id"],
|
|
}
|
|
for tc in tool_calls_raw
|
|
]
|
|
|
|
usage = raw.get("usage")
|
|
if usage:
|
|
result.usage = TokenUsage(
|
|
prompt_tokens=usage.get("prompt_tokens", 0),
|
|
completion_tokens=usage.get("completion_tokens", 0),
|
|
total_tokens=usage.get("total_tokens", 0),
|
|
prompt_tokens_cached=usage.get("prompt_tokens_cached", 0),
|
|
)
|
|
return result
|
|
|
|
async def chat(
|
|
self,
|
|
model: ModelConfig,
|
|
messages: list[BaseMessage],
|
|
api_key: str,
|
|
base_url: str | None = None,
|
|
params: dict[str, Any] | None = None,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
) -> ChatResult:
|
|
url = f"{base_url or self._provider.base_url}/chat/completions"
|
|
body = self.build_request(
|
|
model, messages, params=params, tools=tools, tool_choice=tool_choice
|
|
)
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
try:
|
|
response = await self._http.post(url, json=body, headers=headers)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
return self.parse_nonstreaming_response(data, model.name)
|
|
except httpx.HTTPStatusError as exc:
|
|
raise AIProviderError.from_provider(
|
|
message=f"OpenAI API error: {exc.response.status_code} {exc.response.text[:200]}",
|
|
code="PROVIDER_ERROR",
|
|
provider=self._provider.name,
|
|
original=exc,
|
|
) from exc
|
|
|
|
async def chat_stream(
|
|
self,
|
|
model: ModelConfig,
|
|
messages: list[BaseMessage],
|
|
api_key: str,
|
|
base_url: str | None = None,
|
|
params: dict[str, Any] | None = None,
|
|
tools: list[dict[str, Any]] | None = None,
|
|
tool_choice: str | dict[str, Any] | None = None,
|
|
) -> AsyncIterator[ChatChunk]:
|
|
url = f"{base_url or self._provider.base_url}/chat/completions"
|
|
body = self.build_request(
|
|
model, messages, params=params, tools=tools, tool_choice=tool_choice, stream=True
|
|
)
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
async with self._http.stream("POST", url, json=body, headers=headers) as response:
|
|
if response.status_code != 200:
|
|
error_text = await response.aread()
|
|
raise AIProviderError.from_provider(
|
|
message=f"OpenAI stream error: {response.status_code} {error_text[:200]}",
|
|
code="PROVIDER_ERROR",
|
|
provider=self._provider.name,
|
|
)
|
|
async for line in response.aiter_lines():
|
|
if line.startswith("data: "):
|
|
data_str = line[6:].strip()
|
|
if data_str == "[DONE]":
|
|
break
|
|
if data_str:
|
|
try:
|
|
raw = json.loads(data_str)
|
|
yield self.parse_chunk(raw)
|
|
except json.JSONDecodeError:
|
|
logger.warning("openai_stream_parse_error", raw=data_str[:100])
|
|
|
|
async def close(self) -> None:
|
|
if self._client:
|
|
await self._client.aclose()
|
|
self._client = None
|
|
|
|
|
|
class OpenAIFormatAdapterFactory:
|
|
"""工厂：创建 OpenAI 格式适配器。"""
|
|
|
|
@staticmethod
|
|
def create(provider_config: ProviderConfig) -> OpenAIFormatAdapter:
|
|
return OpenAIFormatAdapter(provider_config=provider_config)
|