94 lines
2.6 KiB
Python
94 lines
2.6 KiB
Python
"""Local Ollama client for mail analysis."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from config import LLM_BASE_URL, LLM_MODEL, LLM_TIMEOUT_SECONDS
|
|
|
|
logger = logging.getLogger("mail_summary")
|
|
|
|
_JSON_BLOCK_RE = re.compile(r"\{.*\}", re.DOTALL)
|
|
|
|
|
|
class LlmError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def extract_json_object(text: str) -> dict[str, Any]:
|
|
text = (text or "").strip()
|
|
if not text:
|
|
raise LlmError("empty_llm_response")
|
|
try:
|
|
data = json.loads(text)
|
|
if isinstance(data, dict):
|
|
return data
|
|
except json.JSONDecodeError:
|
|
pass
|
|
match = _JSON_BLOCK_RE.search(text)
|
|
if not match:
|
|
raise LlmError("llm_response_not_json")
|
|
data = json.loads(match.group(0))
|
|
if not isinstance(data, dict):
|
|
raise LlmError("llm_response_not_object")
|
|
return data
|
|
|
|
|
|
async def chat_json(
|
|
*,
|
|
system: str,
|
|
user: str,
|
|
model: str | None = None,
|
|
temperature: float = 0.2,
|
|
) -> dict[str, Any]:
|
|
selected = (model or LLM_MODEL).strip() or LLM_MODEL
|
|
payload = {
|
|
"model": selected,
|
|
"stream": False,
|
|
"format": "json",
|
|
"options": {"temperature": temperature},
|
|
"messages": [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": user},
|
|
],
|
|
}
|
|
url = f"{LLM_BASE_URL}/api/chat"
|
|
try:
|
|
async with httpx.AsyncClient(timeout=LLM_TIMEOUT_SECONDS) as client:
|
|
response = await client.post(url, json=payload)
|
|
response.raise_for_status()
|
|
body = response.json()
|
|
except httpx.HTTPError as exc:
|
|
logger.exception("Ollama request failed")
|
|
raise LlmError(f"ollama_request_failed: {exc}") from exc
|
|
|
|
content = (
|
|
(body.get("message") or {}).get("content")
|
|
or body.get("response")
|
|
or ""
|
|
)
|
|
return extract_json_object(content)
|
|
|
|
|
|
async def ping_llm(model: str | None = None) -> dict[str, Any]:
|
|
selected = (model or LLM_MODEL).strip() or LLM_MODEL
|
|
try:
|
|
async with httpx.AsyncClient(timeout=3.0) as client:
|
|
tags = await client.get(f"{LLM_BASE_URL}/api/tags")
|
|
tags.raise_for_status()
|
|
data = tags.json()
|
|
except Exception as exc: # noqa: BLE001
|
|
return {"ok": False, "error": str(exc), "model": selected}
|
|
names = [m.get("name", "") for m in data.get("models", [])]
|
|
return {
|
|
"ok": True,
|
|
"model": selected,
|
|
"model_present": any(selected in name or name in selected for name in names),
|
|
"models": names[:20],
|
|
}
|