35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
from app.config import settings
|
||
|
|
|
||
|
|
|
||
|
|
async def generate(prompt: str, system: str | None = None, timeout: float = 300.0) -> str:
|
||
|
|
messages: list[dict[str, str]] = []
|
||
|
|
if system:
|
||
|
|
messages.append({"role": "system", "content": system})
|
||
|
|
messages.append({"role": "user", "content": prompt})
|
||
|
|
return await chat_messages(messages, timeout=timeout)
|
||
|
|
|
||
|
|
|
||
|
|
async def chat_messages(messages: list[dict[str, str]], timeout: float = 300.0) -> str:
|
||
|
|
url = f"{settings.OLLAMA_URL.rstrip('/')}/api/chat"
|
||
|
|
payload = {
|
||
|
|
"model": settings.OLLAMA_MODEL,
|
||
|
|
"messages": messages,
|
||
|
|
"think": False,
|
||
|
|
"stream": False,
|
||
|
|
"keep_alive": "30m",
|
||
|
|
"options": {"num_predict": 280, "temperature": 0.4},
|
||
|
|
}
|
||
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
||
|
|
resp = await client.post(url, json=payload)
|
||
|
|
resp.raise_for_status()
|
||
|
|
msg = resp.json().get("message") or {}
|
||
|
|
content = (msg.get("content") or "").strip()
|
||
|
|
if content:
|
||
|
|
return content
|
||
|
|
thinking = (msg.get("thinking") or "").strip()
|
||
|
|
return thinking[:2000] if thinking else ""
|