185 lines
6.5 KiB
Python
185 lines
6.5 KiB
Python
"""
|
||
统一 LLM 客户端 - 支持多种提供商
|
||
|
||
支持的提供商:
|
||
- ollama: 本地运行 (免费，但占用 CPU/内存)
|
||
- deepseek: 云端 API (便宜，¥1/百万 tokens)
|
||
- openai: 云端 API (GPT-4)
|
||
- qwen: 阿里云通义千问 (便宜，国内稳定)
|
||
"""
|
||
|
||
import httpx
|
||
from typing import Dict, List, Optional, Any
|
||
import os
|
||
import json
|
||
|
||
|
||
class LLMClient:
|
||
"""统一 LLM 客户端"""
|
||
|
||
def __init__(self):
|
||
self.provider = os.getenv("LLM_PROVIDER", "deepseek").lower()
|
||
self.client = self._create_client()
|
||
|
||
def _create_client(self) -> 'LLMProvider':
|
||
"""根据配置创建对应的提供商客户端"""
|
||
if self.provider == "ollama":
|
||
return OllamaClient()
|
||
elif self.provider == "deepseek":
|
||
return OpenAICompatibleClient(
|
||
api_key=os.getenv("DEEPSEEK_API_KEY"),
|
||
base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"),
|
||
model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat")
|
||
)
|
||
elif self.provider == "openai":
|
||
return OpenAICompatibleClient(
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
base_url=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
|
||
model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")
|
||
)
|
||
elif self.provider == "qwen":
|
||
return OpenAICompatibleClient(
|
||
api_key=os.getenv("QWEN_API_KEY"),
|
||
base_url=os.getenv("QWEN_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1"),
|
||
model=os.getenv("QWEN_MODEL", "qwen-plus")
|
||
)
|
||
else:
|
||
raise ValueError(f"不支持的 LLM 提供商: {self.provider}")
|
||
|
||
def generate(self, prompt: str, **kwargs) -> str:
|
||
"""生成文本"""
|
||
return self.client.generate(prompt, **kwargs)
|
||
|
||
def generate_json(self, prompt: str, **kwargs) -> Dict:
|
||
"""生成 JSON 格式响应"""
|
||
return self.client.generate_json(prompt, **kwargs)
|
||
|
||
def get_embedding(self, text: str) -> List[float]:
|
||
"""获取向量嵌入"""
|
||
return self.client.get_embedding(text)
|
||
|
||
|
||
class LLMProvider:
|
||
"""LLM 提供商基类"""
|
||
def generate(self, prompt: str, **kwargs) -> str:
|
||
raise NotImplementedError
|
||
|
||
def generate_json(self, prompt: str, **kwargs) -> Dict:
|
||
raise NotImplementedError
|
||
|
||
def get_embedding(self, text: str) -> List[float]:
|
||
raise NotImplementedError
|
||
|
||
|
||
class OllamaClient(LLMProvider):
|
||
"""Ollama 本地客户端"""
|
||
|
||
def __init__(self):
|
||
self.host = os.getenv("OLLAMA_HOST", "http://127.0.0.1:11434").rstrip('/')
|
||
self.chat_model = os.getenv("OLLAMA_CHAT_MODEL", "llama3:8b")
|
||
self.embedding_model = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")
|
||
self.client = httpx.Client(base_url=self.host, timeout=300.0, trust_env=False)
|
||
|
||
def generate(self, prompt: str, **kwargs) -> str:
|
||
response = self.client.post("/api/generate", json={
|
||
"model": self.chat_model,
|
||
"prompt": prompt,
|
||
"stream": False,
|
||
**kwargs
|
||
})
|
||
response.raise_for_status()
|
||
return response.json()["response"]
|
||
|
||
def generate_json(self, prompt: str, **kwargs) -> Dict:
|
||
response = self.client.post("/api/generate", json={
|
||
"model": self.chat_model,
|
||
"prompt": prompt,
|
||
"stream": False,
|
||
"format": "json",
|
||
**kwargs
|
||
})
|
||
response.raise_for_status()
|
||
content = response.json()["response"]
|
||
# 清理 markdown 代码块
|
||
if content.startswith("```json"):
|
||
content = content[7:]
|
||
if content.startswith("```"):
|
||
content = content[3:]
|
||
if content.endswith("```"):
|
||
content = content[:-3]
|
||
content = content.strip()
|
||
return json.loads(content)
|
||
|
||
def get_embedding(self, text: str) -> List[float]:
|
||
response = self.client.post("/api/embeddings", json={
|
||
"model": self.embedding_model,
|
||
"prompt": text
|
||
})
|
||
response.raise_for_status()
|
||
return response.json().get("embedding", [])
|
||
|
||
|
||
class OpenAICompatibleClient(LLMProvider):
|
||
"""OpenAI 兼容客户端 (适用于 DeepSeek, Qwen 等)"""
|
||
|
||
def __init__(self, api_key: str, base_url: str, model: str):
|
||
self.model = model
|
||
self.base_url = base_url
|
||
self.api_key = api_key
|
||
# 延迟初始化客户端，在实际调用时才检查 API_KEY
|
||
self._client = None
|
||
|
||
@property
|
||
def client(self):
|
||
"""延迟初始化客户端"""
|
||
if self._client is None:
|
||
if not self.api_key:
|
||
raise ValueError(f"请设置 {self.base_url} 的 API_KEY")
|
||
self._client = httpx.Client(
|
||
base_url=self.base_url,
|
||
timeout=120.0,
|
||
headers={"Authorization": f"Bearer {self.api_key}"}
|
||
)
|
||
return self._client
|
||
|
||
def generate(self, prompt: str, **kwargs) -> str:
|
||
response = self.client.post("/chat/completions", json={
|
||
"model": self.model,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
**kwargs
|
||
})
|
||
response.raise_for_status()
|
||
return response.json()["choices"][0]["message"]["content"]
|
||
|
||
def generate_json(self, prompt: str, **kwargs) -> Dict:
|
||
response = self.client.post("/chat/completions", json={
|
||
"model": self.model,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"response_format": {"type": "json_object"},
|
||
**kwargs
|
||
})
|
||
response.raise_for_status()
|
||
content = response.json()["choices"][0]["message"]["content"]
|
||
return json.loads(content)
|
||
|
||
def get_embedding(self, text: str) -> List[float]:
|
||
# DeepSeek 不支持 embedding，使用替代方案或返回空
|
||
if "deepseek" in self.base_url:
|
||
print("[WARNING] DeepSeek 不支持 embedding API，使用简化向量")
|
||
# 简化方案：基于文本内容生成伪向量（实际使用时应该用其他服务）
|
||
import hashlib
|
||
hash_obj = hashlib.md5(text.encode())
|
||
# 生成 768 维伪向量（与 nomic-embed-text 维度一致）
|
||
return [float((hash_obj.digest()[i % 16] / 255.0) * 2 - 1) for _ in range(768)]
|
||
|
||
response = self.client.post("/embeddings", json={
|
||
"model": self.model,
|
||
"input": text
|
||
})
|
||
response.raise_for_status()
|
||
return response.json()["data"][0]["embedding"]
|
||
|
||
|
||
# 单例实例
|
||
llm_client = LLMClient()
|