103 lines
3.1 KiB
Python
103 lines
3.1 KiB
Python
|
"""HTTP client for a local Ollama generate API.
|
||
|
|
|
||
|
|
Matches the newsPipeline local_llm summarizer: POST /api/generate,
|
||
|
|
GET /api/tags. No API key. ``LLM_URL`` may be a host, ``/api``, or
|
||
|
|
``/api/generate`` — we normalize to the origin.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import os
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
_DEFAULT_UA = "osint-dashboard-news-summarizer"
|
||
|
|
_DEFAULT_BASE = "http://127.0.0.1:11434"
|
||
|
|
_GENERATE_TIMEOUT = 300.0
|
||
|
|
_TAGS_TIMEOUT = 8.0
|
||
|
|
|
||
|
|
_GENERATE_OPTIONS = {
|
||
|
|
"num_predict": 4096,
|
||
|
|
"temperature": 0.6,
|
||
|
|
"top_p": 0.9,
|
||
|
|
"top_k": 40,
|
||
|
|
"num_ctx": 32768,
|
||
|
|
"repeat_penalty": 1.1,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_base(url: str | None) -> str:
|
||
|
|
"""Strip path suffixes so LLM_URL variants share one origin."""
|
||
|
|
raw = (url or os.environ.get("LLM_URL", "") or _DEFAULT_BASE).strip()
|
||
|
|
raw = raw.rstrip("/")
|
||
|
|
for suffix in ("/api/generate", "/api/chat", "/api", "/v1"):
|
||
|
|
if raw.endswith(suffix):
|
||
|
|
raw = raw[: -len(suffix)].rstrip("/")
|
||
|
|
break
|
||
|
|
return raw or _DEFAULT_BASE
|
||
|
|
|
||
|
|
|
||
|
|
def _headers() -> dict[str, str]:
|
||
|
|
return {
|
||
|
|
"User-Agent": os.environ.get("OSINT_USER_AGENT") or _DEFAULT_UA,
|
||
|
|
"Accept": "application/json",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def chat(prompt: str, *, model: str, base_url: str | None = None, json_mode: bool = False) -> str:
|
||
|
|
"""POST /api/generate. Returns response text, or \"\" on any failure."""
|
||
|
|
if not (model or "").strip():
|
||
|
|
return ""
|
||
|
|
origin = normalize_base(base_url)
|
||
|
|
url = f"{origin}/api/generate"
|
||
|
|
payload: dict = {
|
||
|
|
"model": model.strip(),
|
||
|
|
"prompt": prompt,
|
||
|
|
"stream": False,
|
||
|
|
"options": dict(_GENERATE_OPTIONS),
|
||
|
|
}
|
||
|
|
if json_mode:
|
||
|
|
payload["format"] = "json"
|
||
|
|
try:
|
||
|
|
with httpx.Client(timeout=_GENERATE_TIMEOUT) as client:
|
||
|
|
resp = client.post(url, headers=_headers(), json=payload)
|
||
|
|
if resp.status_code == 401 or resp.status_code >= 500:
|
||
|
|
return ""
|
||
|
|
data = resp.json()
|
||
|
|
text = data.get("response") if isinstance(data, dict) else None
|
||
|
|
return text if isinstance(text, str) else ""
|
||
|
|
except Exception:
|
||
|
|
return ""
|
||
|
|
|
||
|
|
|
||
|
|
def list_tags(base_url: str | None = None) -> list[str]:
|
||
|
|
"""GET /api/tags model names. Empty list on failure."""
|
||
|
|
origin = normalize_base(base_url)
|
||
|
|
url = f"{origin}/api/tags"
|
||
|
|
try:
|
||
|
|
with httpx.Client(timeout=_TAGS_TIMEOUT) as client:
|
||
|
|
resp = client.get(url, headers=_headers())
|
||
|
|
if resp.status_code != 200:
|
||
|
|
return []
|
||
|
|
models = (resp.json() or {}).get("models") or []
|
||
|
|
names: list[str] = []
|
||
|
|
for item in models:
|
||
|
|
if isinstance(item, dict):
|
||
|
|
name = item.get("name") or item.get("model") or ""
|
||
|
|
if name:
|
||
|
|
names.append(str(name))
|
||
|
|
return names
|
||
|
|
except Exception:
|
||
|
|
return []
|
||
|
|
|
||
|
|
|
||
|
|
def ollama_ready(model: str, *, base_url: str | None = None) -> bool:
|
||
|
|
"""True when /api/tags is up and ``model`` is present (prefix match)."""
|
||
|
|
if not (model or "").strip():
|
||
|
|
return False
|
||
|
|
names = list_tags(base_url)
|
||
|
|
if not names:
|
||
|
|
return False
|
||
|
|
want = model.strip()
|
||
|
|
return any(n == want or n.startswith(want) for n in names)
|