# -*- 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"])]