# -*- coding: utf-8 -*- from __future__ import annotations import os, time, uuid, logging from typing import List, Optional, Dict, Any, Tuple import requests import numpy as np from fastapi import FastAPI, BackgroundTasks, Header, HTTPException from pydantic import BaseModel, Field from qdrant_client import QdrantClient from qdrant_client.http.models import VectorParams, Distance, PointStruct logging.basicConfig(level=logging.INFO) LOG = logging.getLogger("remote_indexer") # ---------- ENV ---------- AUTH_TOKEN = os.getenv("REMOTE_INDEX_TOKEN", "").strip() # simple header auth HF_TOKEN = os.getenv("HF_API_TOKEN", "").strip() HF_MODEL = os.getenv("HF_EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2") HF_URL = os.getenv("HF_API_URL", "").strip() or f"https://api-inference.huggingface.co/pipeline/feature-extraction/{HF_MODEL}" QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") QDRANT_API = os.getenv("QDRANT_API_KEY", "").strip() if not HF_TOKEN: LOG.warning("HF_API_TOKEN manquant — le service refusera /index et /query.") # ---------- Clients ---------- qdr = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API if QDRANT_API else None) # ---------- Pydantic ---------- class FileIn(BaseModel): path: str text: str class IndexRequest(BaseModel): project_id: str = Field(..., min_length=1) files: List[FileIn] chunk_size: int = 1200 overlap: int = 200 batch_size: int = 8 store_text: bool = True class QueryRequest(BaseModel): project_id: str query: str top_k: int = 6 # ---------- Jobs store (en mémoire) ---------- JOBS: Dict[str, Dict[str, Any]] = {} # {job_id: {"status": "...", "logs": [...], "created": ts}} # ---------- Utils ---------- def _auth(x_auth: Optional[str]): if AUTH_TOKEN and (x_auth or "") != AUTH_TOKEN: raise HTTPException(status_code=401, detail="Unauthorized") def _post_embeddings(batch: List[str]) -> Tuple[np.ndarray, int]: if not HF_TOKEN: raise RuntimeError("HF_API_TOKEN manquant (server).") headers = {"Authorization": f"Bearer {HF_TOKEN}"} r = requests.post(HF_URL, headers=headers, json=batch, timeout=120) size = int(r.headers.get("Content-Length", "0")) r.raise_for_status() data = r.json() arr = np.array(data, dtype=np.float32) # arr: [batch, dim] (sentence-transformers) # ou [batch, tokens, dim] -> mean pooling if arr.ndim == 3: arr = arr.mean(axis=1) if arr.ndim != 2: raise RuntimeError(f"Unexpected embeddings shape: {arr.shape}") # normalisation norms = np.linalg.norm(arr, axis=1, keepdims=True) + 1e-12 arr = arr / norms return arr.astype(np.float32), size def _ensure_collection(name: str, dim: int): try: qdr.get_collection(name) return except Exception: pass qdr.create_collection( collection_name=name, vectors_config=VectorParams(size=dim, distance=Distance.COSINE), ) def _chunk_with_spans(text: str, size: int, overlap: int): n = len(text) if size <= 0: yield (0, n, text) return i = 0 while i < n: j = min(n, i + size) yield (i, j, text[i:j]) i = max(0, j - overlap) if i >= n: break def _append_log(job_id: str, line: str): job = JOBS.get(job_id) if not job: return job["logs"].append(line) def _set_status(job_id: str, status: str): job = JOBS.get(job_id) if not job: return job["status"] = status # ---------- Background task ---------- def run_index_job(job_id: str, req: IndexRequest): try: _set_status(job_id, "running") total_chunks = 0 LOG.info(f"[{job_id}] Index start project={req.project_id} files={len(req.files)}") _append_log(job_id, f"Start project={req.project_id} files={len(req.files)}") # premier batch pour récupérer la dimension # on prépare un mini lot warmup = [] for f in req.files[:1]: warmup.append(next(_chunk_with_spans(f.text, req.chunk_size, req.overlap))[2]) embs, sz = _post_embeddings(warmup) dim = embs.shape[1] col = f"proj_{req.project_id}" _ensure_collection(col, dim) _append_log(job_id, f"Collection ready: {col} (dim={dim})") points_buffer: List[PointStruct] = [] point_id = 0 def flush_points(): nonlocal points_buffer if points_buffer: qdr.upsert(collection_name=col, points=points_buffer) points_buffer = [] # boucle fichiers for fi, f in enumerate(req.files, 1): chunks, metas = [], [] for ci, (start, end, chunk_txt) in enumerate(_chunk_with_spans(f.text, req.chunk_size, req.overlap)): chunks.append(chunk_txt) payload = {"path": f.path, "chunk": ci, "start": start, "end": end} if req.store_text: payload["text"] = chunk_txt metas.append(payload) if len(chunks) >= req.batch_size: vecs, sz = _post_embeddings(chunks) batch_points = [] for k, vec in enumerate(vecs): batch_points.append(PointStruct(id=point_id, vector=vec.tolist(), payload=metas[k])) point_id += 1 qdr.upsert(collection_name=col, points=batch_points) total_chunks += len(chunks) _append_log(job_id, f"file {fi}/{len(req.files)}: +{len(chunks)} chunks (total={total_chunks}) ~{sz/1024:.1f}KiB") chunks, metas = [], [] # flush fin de fichier if chunks: vecs, sz = _post_embeddings(chunks) batch_points = [] for k, vec in enumerate(vecs): batch_points.append(PointStruct(id=point_id, vector=vec.tolist(), payload=metas[k])) point_id += 1 qdr.upsert(collection_name=col, points=batch_points) total_chunks += len(chunks) _append_log(job_id, f"file {fi}/{len(req.files)}: +{len(chunks)} chunks (total={total_chunks}) ~{sz/1024:.1f}KiB") flush_points() _append_log(job_id, f"Done. chunks={total_chunks}") _set_status(job_id, "done") LOG.info(f"[{job_id}] Index finished. chunks={total_chunks}") except Exception as e: LOG.exception("Index job failed") _append_log(job_id, f"ERROR: {e}") _set_status(job_id, "error") # ---------- API ---------- app = FastAPI() @app.get("/health") def health(): return {"ok": True} @app.post("/index") def start_index(req: IndexRequest, background_tasks: BackgroundTasks, x_auth_token: Optional[str] = Header(default=None)): _auth(x_auth_token) if not HF_TOKEN: raise HTTPException(400, "HF_API_TOKEN manquant côté serveur.") job_id = uuid.uuid4().hex[:12] JOBS[job_id] = {"status": "queued", "logs": [], "created": time.time()} background_tasks.add_task(run_index_job, job_id, req) return {"job_id": job_id} @app.get("/status/{job_id}") def status(job_id: str, x_auth_token: Optional[str] = Header(default=None)): _auth(x_auth_token) j = JOBS.get(job_id) if not j: raise HTTPException(404, "job inconnu") return {"status": j["status"], "logs": j["logs"][-800:]} @app.post("/query") def query(req: QueryRequest, x_auth_token: Optional[str] = Header(default=None)): _auth(x_auth_token) if not HF_TOKEN: raise HTTPException(400, "HF_API_TOKEN manquant côté serveur.") vec, _ = _post_embeddings([req.query]) vec = vec[0].tolist() col = f"proj_{req.project_id}" try: res = qdr.search(collection_name=col, query_vector=vec, limit=int(req.top_k)) except Exception as e: raise HTTPException(400, f"Search failed: {e}") out = [] for p in res: pl = p.payload or {} txt = pl.get("text") # hard cap snippet size if txt and len(txt) > 800: txt = txt[:800] + "..." out.append({"path": pl.get("path"), "chunk": pl.get("chunk"), "start": pl.get("start"), "end": pl.get("end"), "text": txt}) return {"results": out} @app.post("/wipe") def wipe_collection(project_id: str, x_auth_token: Optional[str] = Header(default=None)): _auth(x_auth_token) col = f"proj_{project_id}" try: qdr.delete_collection(col) return {"ok": True} except Exception as e: raise HTTPException(400, f"wipe failed: {e}")