Files
2026-06-08 15:05:57 +08:00

70 lines
2.7 KiB
Python

import json
from collections.abc import AsyncGenerator
import httpx
from sqlalchemy.ext.asyncio import AsyncSession
from app.models import Agent, LLMModel, LLMProvider
class LLMService:
def __init__(self, db: AsyncSession):
self.db = db
async def complete(self, agent: Agent, user_message: str, history: list[dict] | None = None) -> str:
if not agent.model_id:
return self._fallback_answer(agent, user_message)
model = await self.db.get(LLMModel, agent.model_id)
if not model:
return self._fallback_answer(agent, user_message)
provider = await self.db.get(LLMProvider, model.provider_id)
if not provider or provider.status != "enabled" or not provider.api_key or not provider.base_url:
return self._fallback_answer(agent, user_message)
messages = [{"role": "system", "content": agent.system_prompt or f"你是 {agent.name}"}]
messages.extend(history or [])
messages.append({"role": "user", "content": user_message})
url = provider.base_url.rstrip("/") + "/chat/completions"
payload = {
"model": model.name,
"messages": messages,
"temperature": agent.temperature,
"max_tokens": agent.max_tokens,
"stream": False,
}
async with httpx.AsyncClient(timeout=60) as client:
response = await client.post(
url,
headers={"Authorization": f"Bearer {provider.api_key}", "Content-Type": "application/json"},
json=payload,
)
response.raise_for_status()
data = response.json()
return data.get("choices", [{}])[0].get("message", {}).get("content") or ""
async def stream(self, agent: Agent, user_message: str, history: list[dict] | None = None) -> AsyncGenerator[str, None]:
answer = await self.complete(agent, user_message, history)
for chunk in self._chunk_text(answer):
yield chunk
def _fallback_answer(self, agent: Agent, user_message: str) -> str:
prompt = agent.system_prompt.strip() or "轻量 AI Agent"
return (
f"{agent.name} 已收到任务:{user_message}\n\n"
f"当前使用本地占位响应。配置可用的 LLM Provider 和 Model 后,将自动调用真实模型。\n\n"
f"系统提示词摘要:{prompt[:160]}"
)
def _chunk_text(self, text: str) -> list[str]:
if not text:
return [""]
return [text[i : i + 24] for i in range(0, len(text), 24)]
def sse(data: dict | str) -> str:
if isinstance(data, str):
payload = data
else:
payload = json.dumps(data, ensure_ascii=False, default=str)
return f"data: {payload}\n\n"