Update app.py

app.py

@@ -1,23 +1,22 @@
 """
 KG Embedding Server — FastAPI on HuggingFace Spaces
-يشتغل كـ REST API، يحمل الموديل مرة واحدة بس
+Session-based: client syncs once → server holds FAISS index in RAM
 """
 
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from pydantic import BaseModel
-from typing import Optional
-import math
-import re
-import contextlib
+from typing import Dict, Any, Optional
+import math, contextlib, uuid
 from io import StringIO
 
 app = FastAPI(title="KG Embedding Server")
 
 # ══════════════════════════════════════════════════════
-#  GLOBALS — الموديل بيتحمل مرة واحدة عند startup
+#  GLOBALS
 # ══════════════════════════════════════════════════════
 _model = None
 _use_st = False
+_sessions: Dict[str, Dict[str, Any]] = {}   # session_id → {nodes, index}
 
 @app.on_event("startup")
 def load_model():
@@ -30,24 +29,35 @@ def load_model():
         print("[Server] sentence-transformers loaded ✓")
     except Exception as e:
         print(f"[Server] ST unavailable: {e}")
-        _use_st = False
 
 # ══════════════════════════════════════════════════════
-#  REQUEST / RESPONSE MODELS
+#  MODELS
 # ══════════════════════════════════════════════════════
-class EmbedRequest(BaseModel):
-    texts: list[str]
+class NodePayload(BaseModel):
+    id: int
+    title: str
+    content: str
 
-class EmbedResponse(BaseModel):
-    embeddings: list[list[float]]
-    model: str = "all-MiniLM-L6-v2"
+class SyncRequest(BaseModel):
+    session_id: str
+    nodes: list[NodePayload]
 
-class SimilarityRequest(BaseModel):
+class SyncResponse(BaseModel):
+    status: str
+    count: int
+
+class SearchRequest(BaseModel):
+    session_id: str
     query: str
-    candidates: list[str]   # النصوص اللي هنقيس عليها
+    top_k: int = 8
+
+class SearchResult(BaseModel):
+    node_id: int
+    score: float
 
-class SimilarityResponse(BaseModel):
-    scores: list[float]     # cosine similarity لكل candidate
+class SearchResponse(BaseModel):
+    results: list[SearchResult]
+    session_missing: bool = False
 
 class HealthResponse(BaseModel):
     status: str
@@ -61,37 +71,66 @@ class HealthResponse(BaseModel):
 def health():
     return {"status": "ok", "model_loaded": _use_st}
 
-@app.post("/embed", response_model=EmbedResponse)
-def embed(req: EmbedRequest):
+@app.post("/sync", response_model=SyncResponse)
+def sync(req: SyncRequest):
+    """Client uploads nodes once → server builds FAISS index in RAM."""
+    try:
+        import faiss
+        import numpy as np
+    except ImportError:
+        raise Exception("faiss not installed on server")
+
     if not _use_st or _model is None:
-        raise HTTPException(503, "Model not loaded")
-    if not req.texts:
-        return EmbedResponse(embeddings=[])
+        return SyncResponse(status="no_model", count=0)
+
+    texts = [f"{n.title} {n.content}" for n in req.nodes]
+    if not texts:
+        _sessions[req.session_id] = {"nodes": [], "index": None}
+        return SyncResponse(status="ok", count=0)
+
     with contextlib.redirect_stdout(StringIO()), contextlib.redirect_stderr(StringIO()):
-        vecs = _model.encode(req.texts, show_progress_bar=False)
-    return EmbedResponse(embeddings=[v.tolist() for v in vecs])
+        vecs = _model.encode(texts, show_progress_bar=False, normalize_embeddings=True)
+
+    vecs = np.array(vecs, dtype="float32")
+    dim  = vecs.shape[1]
+
+    index = faiss.IndexFlatIP(dim)   # cosine sim (vectors already normalized)
+    index.add(vecs)
+
+    _sessions[req.session_id] = {
+        "nodes": [n.dict() for n in req.nodes],
+        "index": index,
+    }
+    print(f"[Server] Session {req.session_id[:8]}… → {len(texts)} nodes indexed ✓")
+    return SyncResponse(status="ok", count=len(texts))
+
+
+@app.post("/search", response_model=SearchResponse)
+def search(req: SearchRequest):
+    """Embed query only, FAISS search against cached index."""
+    import numpy as np
+
+    session = _sessions.get(req.session_id)
+    if not session or session["index"] is None:
+        return SearchResponse(results=[], session_missing=True)
 
-@app.post("/similarity", response_model=SimilarityResponse)
-def similarity(req: SimilarityRequest):
-    """Query + candidates → cosine scores (الأسرع لو عايز ترتب نتايج)"""
-    if not _use_st or _model is None:
-        raise HTTPException(503, "Model not loaded")
-    texts = [req.query] + req.candidates
     with contextlib.redirect_stdout(StringIO()), contextlib.redirect_stderr(StringIO()):
-        vecs = _model.encode(texts, show_progress_bar=False)
-    qvec = vecs[0]
-    scores = [_cosine(qvec, v) for v in vecs[1:]]
-    return SimilarityResponse(scores=scores)
+        qvec = _model.encode([req.query], normalize_embeddings=True)
+
+    qvec = np.array(qvec, dtype="float32")
+    k    = min(req.top_k, len(session["nodes"]))
+    scores, indices = session["index"].search(qvec, k)
+
+    results = []
+    for score, idx in zip(scores[0], indices[0]):
+        if idx >= 0:
+            results.append(SearchResult(
+                node_id=session["nodes"][idx]["id"],
+                score=float(score)
+            ))
+    return SearchResponse(results=results)
 
-def _cosine(a, b) -> float:
-    dot = sum(x * y for x, y in zip(a, b))
-    na = math.sqrt(sum(x * x for x in a)) or 1e-9
-    nb = math.sqrt(sum(x * x for x in b)) or 1e-9
-    return dot / (na * nb)
 
-# ══════════════════════════════════════════════════════
-#  PING — لمنع الـ Space من النوم (اعمل cron job يبعته كل 4 دقايق)
-# ══════════════════════════════════════════════════════
 @app.get("/ping")
 def ping():
     return {"pong": True}