70 lines
2.7 KiB
Python
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"
|