zyi
b19a2415f9
- 新增 app/api、app/service 数据分析 Agent 全套服务与接口 - schemas.py 重构为 schemas 包（analyst schema） - 新增 SQL 防注入、guardrail、缓存、字典、LLM 等服务 - 新增 tests 测试套件与 scripts/dev、scripts/setup 脚本 - 补充需求规格、架构说明书、开发清单、表设计等文档
74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
"""DeepSeek LLM 客户端（对话 / NL2SQL / 解读）。"""
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
|
|
import httpx
|
|
|
|
from app.config.settings import settings
|
|
|
|
|
|
class LLMError(Exception):
|
|
pass
|
|
|
|
|
|
class DeepSeekLLM:
|
|
def __init__(
|
|
self,
|
|
api_key: str | None = None,
|
|
base_url: str | None = None,
|
|
model: str = "deepseek-chat",
|
|
) -> None:
|
|
self.api_key = api_key or settings.deepseek_api_key
|
|
self.base_url = (base_url or settings.deepseek_base_url).rstrip("/")
|
|
self.model = model
|
|
|
|
def complete(
|
|
self,
|
|
messages: list[dict],
|
|
temperature: float = 0.0,
|
|
max_tokens: int = 2048,
|
|
) -> tuple[str, dict]:
|
|
"""返回 (文本, usage)。"""
|
|
if not self.api_key:
|
|
raise LLMError("缺少 DEEPSEEK_API_KEY")
|
|
r = httpx.post(
|
|
f"{self.base_url}/chat/completions",
|
|
headers={
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json={
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
},
|
|
timeout=60,
|
|
)
|
|
if r.status_code != 200:
|
|
raise LLMError(f"LLM HTTP {r.status_code}: {r.text[:300]}")
|
|
data = r.json()
|
|
content = data["choices"][0]["message"]["content"]
|
|
usage = data.get("usage", {})
|
|
return content, usage
|
|
|
|
def chat(self, messages: list[dict], temperature: float = 0.0, max_tokens: int = 2048) -> str:
|
|
text, _ = self.complete(messages, temperature, max_tokens)
|
|
return text
|
|
|
|
|
|
def extract_sql(text: str) -> str:
|
|
"""从 LLM 回复中提取 SQL（去掉 ```sql 代码块等）。"""
|
|
m = re.search(r"```(?:sql)?\s*(.*?)```", text, re.IGNORECASE | re.DOTALL)
|
|
if m:
|
|
return m.group(1).strip()
|
|
return text.strip()
|
|
|
|
|
|
def estimate_cost(usage: dict) -> float:
|
|
"""粗略成本估算（元），按 DeepSeek 通用单价量级。"""
|
|
prompt = int(usage.get("prompt_tokens", 0) or 0)
|
|
completion = int(usage.get("completion_tokens", 0) or 0)
|
|
return round((prompt * 1.0 + completion * 2.0) / 1_000_000, 6)
|