173 lines
7.0 KiB
Python
173 lines
7.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Embedding providers — con thermal gate e keep_alive adattivo."""
|
|
|
|
import logging
|
|
import os
|
|
import time
|
|
from functools import lru_cache
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
|
|
from . import thermal
|
|
|
|
LOGGER = logging.getLogger("loogle_mcp.embeddings")
|
|
|
|
|
|
def embed_query(text: str) -> list[float]:
|
|
"""Embedding interattivo per una query di ricerca.
|
|
|
|
Ordine: OLLAMA_QUERY_URL (istanza dedicata su pi1, ~0.1s, mai in coda dietro
|
|
l'indicizzazione) → OLLAMA_URL (DS920, senza thermal gate) → OpenAI.
|
|
Timeout corti così il client MCP riceve un errore invece di un hang.
|
|
Cache LRU: gli agenti ripetono spesso le stesse query.
|
|
"""
|
|
return list(_embed_query_cached(text.strip(), os.environ.get("OLLAMA_EMBED_MODEL", "nomic-embed-text")))
|
|
|
|
|
|
@lru_cache(maxsize=256)
|
|
def _embed_query_cached(text: str, _model: str) -> tuple[float, ...]:
|
|
return tuple(_embed_query_uncached(text))
|
|
|
|
|
|
def _embed_query_uncached(text: str) -> list[float]:
|
|
query_url = os.environ.get("OLLAMA_QUERY_URL", "").strip()
|
|
ollama_url = os.environ.get("OLLAMA_URL", "").strip()
|
|
openai_key = os.environ.get("OPENAI_API_KEY", "").strip()
|
|
last_exc: Optional[Exception] = None
|
|
if query_url:
|
|
try:
|
|
# Istanza dedicata: options proprie, il modello resta caricato (keep_alive -1)
|
|
return _embed_ollama_query(
|
|
text, query_url, read_s=_float_env("OLLAMA_QUERY_URL_TIMEOUT_S", 10.0),
|
|
keep_alive=-1, options=_ollama_options("OLLAMA_QUERY_NUM_THREAD"),
|
|
)
|
|
except Exception as exc:
|
|
LOGGER.warning("Ollama query (%s) failed, fallback: %s", query_url, exc)
|
|
last_exc = exc
|
|
if ollama_url:
|
|
try:
|
|
return _embed_ollama_query(
|
|
text, ollama_url, read_s=_float_env("OLLAMA_QUERY_TIMEOUT_S", 20.0),
|
|
keep_alive=int(_float_env("OLLAMA_KEEP_ALIVE_COOL", 300)), options=_ollama_options(),
|
|
)
|
|
except Exception as exc:
|
|
LOGGER.warning("Ollama query embedding failed: %s", exc)
|
|
last_exc = exc
|
|
if openai_key:
|
|
return _embed_openai([text], openai_key, timeout=_float_env("OPENAI_QUERY_TIMEOUT_S", 20.0))[0]
|
|
if last_exc is not None:
|
|
raise RuntimeError(
|
|
f"Servizio embedding (Ollama) non disponibile o lento: {last_exc}. Riprova tra poco."
|
|
) from last_exc
|
|
raise RuntimeError("Nessun provider embedding configurato (OLLAMA_URL o OPENAI_API_KEY)")
|
|
|
|
|
|
def _float_env(name: str, default: float) -> float:
|
|
try:
|
|
return float(os.environ.get(name, str(default)))
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
def _ollama_options(env_name: str = "OLLAMA_NUM_THREAD") -> Optional[dict]:
|
|
# Sulla stessa istanza usare sempre le stesse options: se cambiano Ollama ricarica il modello
|
|
num_thread = os.environ.get(env_name, "").strip()
|
|
if num_thread:
|
|
try:
|
|
return {"num_thread": int(num_thread)}
|
|
except ValueError:
|
|
pass
|
|
return None
|
|
|
|
|
|
def _embed_ollama_query(
|
|
text: str, base_url: str, *, read_s: float, keep_alive: int, options: Optional[dict]
|
|
) -> list[float]:
|
|
model = os.environ.get("OLLAMA_EMBED_MODEL", "nomic-embed-text")
|
|
timeout = httpx.Timeout(connect=5.0, read=read_s, write=10.0, pool=5.0)
|
|
payload = {"model": model, "prompt": text, "keep_alive": keep_alive}
|
|
if options:
|
|
payload["options"] = options
|
|
with httpx.Client(timeout=timeout) as client:
|
|
resp = client.post(f"{base_url.rstrip('/')}/api/embeddings", json=payload)
|
|
resp.raise_for_status()
|
|
return resp.json()["embedding"]
|
|
|
|
|
|
def embed_texts(texts: list[str]) -> list[list[float]]:
|
|
if not texts:
|
|
return []
|
|
# OLLAMA_INDEX_URL: istanza dedicata su pi1 (condivisa con le query). Nessun fallback
|
|
# sul DS920 per l'indicizzazione: lì un chunk costa minuti, meglio riprovare al ciclo dopo
|
|
index_url = os.environ.get("OLLAMA_INDEX_URL", "").strip()
|
|
ollama_url = index_url or os.environ.get("OLLAMA_URL", "").strip()
|
|
if ollama_url:
|
|
try:
|
|
return _embed_ollama(texts, ollama_url, dedicated=bool(index_url))
|
|
except Exception as exc:
|
|
LOGGER.warning("Ollama embedding failed (%s): %s", ollama_url, exc)
|
|
openai_key = os.environ.get("OPENAI_API_KEY", "").strip()
|
|
if openai_key:
|
|
return _embed_openai(texts, openai_key)
|
|
if ollama_url:
|
|
raise RuntimeError(f"Ollama non disponibile per l'indicizzazione ({ollama_url})")
|
|
raise RuntimeError("Nessun provider embedding configurato (OLLAMA_URL o OPENAI_API_KEY)")
|
|
|
|
|
|
def _embed_ollama(texts: list[str], base_url: str, dedicated: bool = False) -> list[list[float]]:
|
|
"""dedicated=True: istanza condivisa con le query → stesse options e modello sempre caricato."""
|
|
model = os.environ.get("OLLAMA_EMBED_MODEL", "nomic-embed-text")
|
|
vectors = []
|
|
timeout = httpx.Timeout(connect=30.0, read=300.0, write=30.0, pool=30.0)
|
|
with httpx.Client(timeout=timeout) as client:
|
|
for i, text in enumerate(texts):
|
|
status = thermal.wait_for_headroom(context=f"embed:{i+1}/{len(texts)}")
|
|
keep_alive = -1 if dedicated else thermal.suggested_keep_alive(status)
|
|
delay = thermal.suggested_delay_s(status)
|
|
|
|
payload = {"model": model, "prompt": text, "keep_alive": keep_alive}
|
|
# options.num_thread limita i thread CPU lato Ollama (se supportato)
|
|
options = _ollama_options("OLLAMA_QUERY_NUM_THREAD" if dedicated else "OLLAMA_NUM_THREAD")
|
|
if options:
|
|
payload["options"] = options
|
|
|
|
resp = client.post(f"{base_url.rstrip('/')}/api/embeddings", json=payload)
|
|
if resp.status_code >= 400:
|
|
LOGGER.warning(
|
|
"Ollama embeddings HTTP %s: %s — payload keys=%s",
|
|
resp.status_code,
|
|
resp.text[:300],
|
|
list(payload.keys()),
|
|
)
|
|
resp.raise_for_status()
|
|
vectors.append(resp.json()["embedding"])
|
|
|
|
if delay > 0 and i + 1 < len(texts):
|
|
time.sleep(delay)
|
|
|
|
# Unload solo se esplicitamente richiesto (zona HARD) — evita spike da reload
|
|
if keep_alive == 0:
|
|
try:
|
|
client.post(
|
|
f"{base_url.rstrip('/')}/api/generate",
|
|
json={"model": model, "keep_alive": 0},
|
|
timeout=30.0,
|
|
)
|
|
except Exception:
|
|
pass
|
|
return vectors
|
|
|
|
|
|
def _embed_openai(texts: list[str], api_key: str, timeout: float = 120.0) -> list[list[float]]:
|
|
model = os.environ.get("OPENAI_EMBED_MODEL", "text-embedding-3-small")
|
|
with httpx.Client(timeout=timeout) as client:
|
|
resp = client.post(
|
|
"https://api.openai.com/v1/embeddings",
|
|
headers={"Authorization": f"Bearer {api_key}"},
|
|
json={"model": model, "input": texts},
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()["data"]
|
|
return [item["embedding"] for item in sorted(data, key=lambda x: x["index"])]
|