Matteo Benedetto
4776b4b496
- gateway: separate config, models, state, audit, guardrail, embeddings, store, metrics, cleanup and routes; keep main.py as FastAPI bootstrap - extension: split client/config, six tools, config command and rules; preserve jiti entrypoint and registrations - Dockerfile copies the complete gateway module set - tests: update monkeypatch boundaries for modular config/state
89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
"""Embedding denso (Ollama/llama.cpp) e sparse BM25."""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Optional
|
|
|
|
import httpx
|
|
from qdrant_client.http import models as qm
|
|
|
|
from config import EMBED_API, EMBED_API_KEY, EMBED_MODEL, EMBED_URL, SPARSE_VECTOR_NAME, log
|
|
|
|
try:
|
|
from fastembed import SparseTextEmbedding
|
|
_sparse_model: Optional[SparseTextEmbedding] = None
|
|
SPARSE_AVAILABLE = True
|
|
except Exception: # noqa: BLE001
|
|
_sparse_model = None
|
|
SPARSE_AVAILABLE = False
|
|
log.warning("fastembed non disponibile: hybrid retrieval disattivato")
|
|
|
|
_http: Optional[httpx.AsyncClient] = None
|
|
|
|
|
|
def get_http() -> httpx.AsyncClient:
|
|
global _http
|
|
if _http is None:
|
|
_http = httpx.AsyncClient(timeout=30)
|
|
return _http
|
|
|
|
|
|
async def close_http() -> None:
|
|
global _http
|
|
if _http is not None:
|
|
await _http.aclose()
|
|
_http = None
|
|
|
|
|
|
async def embed(text: str) -> list[float]:
|
|
if EMBED_API == "llamacpp":
|
|
headers = {"Content-Type": "application/json"}
|
|
if EMBED_API_KEY:
|
|
headers["Authorization"] = f"Bearer {EMBED_API_KEY}"
|
|
response = await get_http().post(f"{EMBED_URL}/v1/embeddings", json={"model": EMBED_MODEL, "input": text}, headers=headers)
|
|
response.raise_for_status()
|
|
return response.json()["data"][0]["embedding"]
|
|
response = await get_http().post(f"{EMBED_URL}/api/embed", json={"model": EMBED_MODEL, "input": text})
|
|
response.raise_for_status()
|
|
return response.json()["embeddings"][0]
|
|
|
|
|
|
def get_sparse_model():
|
|
global _sparse_model
|
|
if _sparse_model is None and SPARSE_AVAILABLE:
|
|
_sparse_model = SparseTextEmbedding(model_name="Qdrant/bm25")
|
|
return _sparse_model
|
|
|
|
|
|
def sparse_encode(text: str) -> Optional[qm.SparseVector]:
|
|
model = get_sparse_model()
|
|
if model is None:
|
|
return None
|
|
emb = next(model.embed(text))
|
|
return qm.SparseVector(indices=emb.indices.tolist(), values=emb.values.tolist())
|
|
|
|
|
|
def backfill_sparse(qdrant: Any, collection: str) -> None:
|
|
if not SPARSE_AVAILABLE:
|
|
return
|
|
offset: Any = None
|
|
updated = 0
|
|
while True:
|
|
points, next_offset = qdrant.scroll(collection_name=collection, limit=100, with_payload=["text"], with_vectors=True, offset=offset)
|
|
batch: list[qm.PointStruct] = []
|
|
for point in points:
|
|
vectors = point.vector or {}
|
|
if SPARSE_VECTOR_NAME in vectors:
|
|
continue
|
|
text = (point.payload or {}).get("text", "")
|
|
sparse = sparse_encode(text) if text else None
|
|
if sparse is not None:
|
|
batch.append(qm.PointStruct(id=point.id, vector={SPARSE_VECTOR_NAME: sparse}))
|
|
if batch:
|
|
qdrant.update_vectors(collection_name=collection, points=batch)
|
|
updated += len(batch)
|
|
if not next_offset:
|
|
break
|
|
offset = next_offset
|
|
if updated:
|
|
log.info("backfill sparse: %d record aggiornati", updated)
|