"""AI Provider 抽象层，支持 OpenAI 兼容 API 和 Anthropic 原生 API。 支持文本流式输出和 tool calling。""" import json from dataclasses import dataclass, field from typing import AsyncGenerator import httpx from ..models import AIConfig @dataclass class ToolCall: """一次工具调用""" id: str name: str arguments: dict class StreamEvent: """流式事件：文本chunk、usage信息、或工具调用""" def __init__( self, text: str = "", usage: dict | None = None, tool_calls: list[ToolCall] | None = None, ): self.text = text self.usage = usage self.tool_calls = tool_calls async def stream_chat( config: AIConfig, messages: list[dict], system_prompt: str = "", tools: list[dict] | None = None, ) -> AsyncGenerator[StreamEvent, None]: """根据 provider 类型调用对应的流式 API，返回 StreamEvent。""" if config.provider == "anthropic": async for event in _stream_anthropic(config, messages, system_prompt, tools): yield event else: async for event in _stream_openai(config, messages, system_prompt, tools): yield event # ── 消息格式构建（用于 tool call 循环中追加消息）── def build_assistant_message(provider: str, text: str, tool_calls: list[ToolCall]) -> dict: """构建包含工具调用的 assistant 消息""" if provider == "anthropic": content = [] if text: content.append({"type": "text", "text": text}) for tc in tool_calls: content.append({ "type": "tool_use", "id": tc.id, "name": tc.name, "input": tc.arguments, }) return {"role": "assistant", "content": content} else: # OpenAI 格式 msg: dict = { "role": "assistant", "content": text or None, "tool_calls": [ { "id": tc.id, "type": "function", "function": { "name": tc.name, "arguments": json.dumps(tc.arguments, ensure_ascii=False), }, } for tc in tool_calls ], } return msg def build_tool_results_messages( provider: str, results: list[dict], ) -> list[dict]: """构建工具结果消息。results: [{"id": str, "result": str}] OpenAI: 返回多条 tool message Anthropic: 返回一条 user message，content 是 tool_result 数组 """ if provider == "anthropic": content = [ {"type": "tool_result", "tool_use_id": r["id"], "content": r["result"]} for r in results ] return [{"role": "user", "content": content}] else: return [ {"role": "tool", "tool_call_id": r["id"], "content": r["result"]} for r in results ] # ── OpenAI 兼容 ── async def _stream_openai( config: AIConfig, messages: list[dict], system_prompt: str, tools: list[dict] | None = None, ) -> AsyncGenerator[StreamEvent, None]: """调用 OpenAI 兼容 API（stream=true），支持 tool calling。""" url = f"{config.base_url.rstrip('/')}/chat/completions" payload_messages = [] if system_prompt: payload_messages.append({"role": "system", "content": system_prompt}) payload_messages.extend(messages) body: dict = { "model": config.model, "messages": payload_messages, "stream": True, "stream_options": {"include_usage": True}, } if tools: body["tools"] = tools async with httpx.AsyncClient(timeout=120) as client: async with client.stream( "POST", url, headers={ "Authorization": f"Bearer {config.api_key}", "Content-Type": "application/json", }, json=body, ) as resp: if resp.status_code >= 400: body_text = await resp.aread() try: err = json.loads(body_text) msg = err.get("error", {}).get("message", "") or str(err) except Exception: msg = body_text.decode(errors="replace")[:200] raise RuntimeError(f"API错误 ({resp.status_code}): {msg}") ct = resp.headers.get("content-type", "") if "text/html" in ct: raise RuntimeError(f"API地址返回了HTML页面，请检查API地址是否正确（当前: {url}）") # 累积工具调用 pending_tool_calls: dict[int, dict] = {} # index -> {id, name, arguments} async for line in resp.aiter_lines(): if not line.startswith("data: "): continue data = line[6:] if data.strip() == "[DONE]": break try: chunk = json.loads(data) # 提取 usage usage = chunk.get("usage") if usage: yield StreamEvent(usage={ "prompt_tokens": usage.get("prompt_tokens", 0), "completion_tokens": usage.get("completion_tokens", 0), "total_tokens": usage.get("total_tokens", 0), }) choices = chunk.get("choices", []) if not choices: continue choice = choices[0] delta = choice.get("delta", {}) finish_reason = choice.get("finish_reason") # 文本内容 content = delta.get("content") if content: yield StreamEvent(text=content) # 工具调用 chunk for tc_delta in delta.get("tool_calls", []): idx = tc_delta["index"] if idx not in pending_tool_calls: pending_tool_calls[idx] = {"id": "", "name": "", "arguments": ""} entry = pending_tool_calls[idx] if "id" in tc_delta: entry["id"] = tc_delta["id"] func = tc_delta.get("function", {}) if "name" in func: entry["name"] = func["name"] if "arguments" in func: entry["arguments"] += func["arguments"] # 结束：如果有工具调用，发出事件 if finish_reason == "tool_calls" and pending_tool_calls: tool_calls = [] for tc_data in pending_tool_calls.values(): try: args = json.loads(tc_data["arguments"]) if tc_data["arguments"] else {} except json.JSONDecodeError: args = {} tool_calls.append(ToolCall( id=tc_data["id"], name=tc_data["name"], arguments=args, )) yield StreamEvent(tool_calls=tool_calls) except (json.JSONDecodeError, KeyError, IndexError): continue # ── Anthropic 原生 ── async def _stream_anthropic( config: AIConfig, messages: list[dict], system_prompt: str, tools: list[dict] | None = None, ) -> AsyncGenerator[StreamEvent, None]: """调用 Anthropic Messages API（stream=true），支持 tool calling。""" url = f"{config.base_url.rstrip('/')}/messages" body: dict = { "model": config.model, "max_tokens": 4096, "messages": messages, "stream": True, } if system_prompt: body["system"] = system_prompt if tools: body["tools"] = tools async with httpx.AsyncClient(timeout=120) as client: async with client.stream( "POST", url, headers={ "x-api-key": config.api_key, "anthropic-version": "2023-06-01", "content-type": "application/json", }, json=body, ) as resp: if resp.status_code >= 400: body_text = await resp.aread() try: err = json.loads(body_text) msg = err.get("error", {}).get("message", "") or str(err) except Exception: msg = body_text.decode(errors="replace")[:200] raise RuntimeError(f"API错误 ({resp.status_code}): {msg}") ct = resp.headers.get("content-type", "") if "text/html" in ct: raise RuntimeError(f"API地址返回了HTML页面，请检查API地址是否正确（当前: {url}）") # 累积工具调用 pending_tool_uses: dict[int, dict] = {} # block_index -> {id, name, input_json} stop_reason = None async for line in resp.aiter_lines(): if not line.startswith("data: "): continue try: event = json.loads(line[6:]) event_type = event.get("type") if event_type == "content_block_start": block = event.get("content_block", {}) idx = event.get("index", 0) if block.get("type") == "tool_use": pending_tool_uses[idx] = { "id": block["id"], "name": block["name"], "input_json": "", } elif event_type == "content_block_delta": idx = event.get("index", 0) delta = event.get("delta", {}) if delta.get("type") == "text_delta": text = delta.get("text", "") if text: yield StreamEvent(text=text) elif delta.get("type") == "input_json_delta": if idx in pending_tool_uses: pending_tool_uses[idx]["input_json"] += delta.get("partial_json", "") elif event_type == "message_delta": usage = event.get("usage", {}) if usage: yield StreamEvent(usage={ "prompt_tokens": 0, "completion_tokens": usage.get("output_tokens", 0), "total_tokens": 0, }) stop_reason = event.get("delta", {}).get("stop_reason") elif event_type == "message_start": msg = event.get("message", {}) usage = msg.get("usage", {}) if usage: yield StreamEvent(usage={ "prompt_tokens": usage.get("input_tokens", 0), "completion_tokens": 0, "total_tokens": 0, }) elif event_type == "message_stop": break except json.JSONDecodeError: continue # 如果有工具调用，发出事件 if stop_reason == "tool_use" and pending_tool_uses: tool_calls = [] for tu_data in pending_tool_uses.values(): try: args = json.loads(tu_data["input_json"]) if tu_data["input_json"] else {} except json.JSONDecodeError: args = {} tool_calls.append(ToolCall( id=tu_data["id"], name=tu_data["name"], arguments=args, )) yield StreamEvent(tool_calls=tool_calls)