55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
import json
|
|
import httpx
|
|
|
|
|
|
class LLMClient:
|
|
"""统一 LLM 接口，支持 OpenAI 兼容接口和 Anthropic 原生接口。"""
|
|
|
|
def __init__(self, provider: str, api_key: str, base_url: str, model: str):
|
|
self.provider = provider.lower()
|
|
self.api_key = api_key
|
|
self.base_url = base_url.rstrip("/")
|
|
self.model = model
|
|
|
|
async def complete(self, system_prompt: str, user_prompt: str) -> str:
|
|
if self.provider == "anthropic":
|
|
return await self._call_anthropic(system_prompt, user_prompt)
|
|
return await self._call_openai_compat(system_prompt, user_prompt)
|
|
|
|
async def _call_openai_compat(self, system_prompt: str, user_prompt: str) -> str:
|
|
"""适配 DeepSeek / 通义千问 / OpenAI 等兼容 /v1/chat/completions 的接口。"""
|
|
async with httpx.AsyncClient(timeout=90) as client:
|
|
resp = await client.post(
|
|
f"{self.base_url}/v1/chat/completions",
|
|
headers={"Authorization": f"Bearer {self.api_key}"},
|
|
json={
|
|
"model": self.model,
|
|
"temperature": 0.2,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
],
|
|
},
|
|
)
|
|
resp.raise_for_status()
|
|
return resp.json()["choices"][0]["message"]["content"]
|
|
|
|
async def _call_anthropic(self, system_prompt: str, user_prompt: str) -> str:
|
|
async with httpx.AsyncClient(timeout=90) as client:
|
|
resp = await client.post(
|
|
f"{self.base_url}/v1/messages",
|
|
headers={
|
|
"x-api-key": self.api_key,
|
|
"anthropic-version": "2023-06-01",
|
|
"content-type": "application/json",
|
|
},
|
|
json={
|
|
"model": self.model,
|
|
"max_tokens": 2048,
|
|
"system": system_prompt,
|
|
"messages": [{"role": "user", "content": user_prompt}],
|
|
},
|
|
)
|
|
resp.raise_for_status()
|
|
return resp.json()["content"][0]["text"]
|