Upload 3 files

app.py

@@ -12,11 +12,9 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
-import faiss
 import gradio as gr
-from sentence_transformers import SentenceTransformer
 
-from extracted_phase2_core import AgenticSelfRAG, Chunk, K_PASSAGES
+from selfrag_core import AgenticSelfRAG, Chunk, K_PASSAGES, LightweightRetriever
 
 
 APP_NAME = "SourceTruth"
@@ -163,44 +161,15 @@ class EmptyRetriever:
         return []
 
 
-class SessionRetriever:
-    def __init__(self, chunks: List[Chunk], encoder: SentenceTransformer):
-        self.chunks = chunks
-        self._encoder = encoder
-        self.index = None
-        self._build_index()
-
-    def _build_index(self):
-        if not self.chunks:
-            return
-        texts = [f"{chunk.source_file} {chunk.text}" for chunk in self.chunks]
-        embeddings = self._encoder.encode(
-            texts,
-            convert_to_numpy=True,
-            normalize_embeddings=True,
-            show_progress_bar=False,
-        ).astype("float32")
-        self.index = faiss.IndexFlatIP(embeddings.shape[1])
-        self.index.add(embeddings)
-
-    def retrieve(self, query: str, k: int = K_PASSAGES) -> List[Chunk]:
-        if self.index is None:
-            return []
-        query_embedding = self._encoder.encode(
-            [query],
-            convert_to_numpy=True,
-            normalize_embeddings=True,
-            show_progress_bar=False,
-        ).astype("float32")
-        _, indices = self.index.search(query_embedding, min(k, len(self.chunks)))
-        return [self.chunks[i] for i in indices[0] if 0 <= i < len(self.chunks)]
+class SessionRetriever(LightweightRetriever):
+    def __init__(self, chunks: List[Chunk]):
+        super().__init__(device="cpu")
+        self.index_chunks(chunks)
 
 
 SESSIONS: Dict[str, SessionData] = {}
 SESSIONS_LOCK = threading.Lock()
 MODEL_LOCK = threading.Lock()
-EMBEDDER_LOCK = threading.Lock()
-GLOBAL_EMBEDDER: Optional[SentenceTransformer] = None
 GLOBAL_AGENT_TEMPLATE: Optional[AgenticSelfRAG] = None
 
 
@@ -388,14 +357,6 @@ def cleanup_expired_sessions():
         log_event("session_expired", session_id=session_id)
 
 
-def get_embedder() -> SentenceTransformer:
-    global GLOBAL_EMBEDDER
-    with EMBEDDER_LOCK:
-        if GLOBAL_EMBEDDER is None:
-            GLOBAL_EMBEDDER = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
-        return GLOBAL_EMBEDDER
-
-
 def get_agent_template() -> AgenticSelfRAG:
     global GLOBAL_AGENT_TEMPLATE
     with MODEL_LOCK:
@@ -935,7 +896,7 @@ def build_session(file_path: str) -> SessionData:
 
         file_hash = sha256_file(dest_path)
         file_size_bytes = os.path.getsize(dest_path)
-        retriever = SessionRetriever(chunks, get_embedder())
+        retriever = SessionRetriever(chunks)
         agent = build_session_agent(retriever)
 
         return SessionData(
@@ -1091,7 +1052,7 @@ def build_corpus_session() -> SessionData:
         if file_name == "02_Validation_Master_Plan.pdf":
             structured["vmp_table"] = parse_vmp_table(doc_pages)
 
-    retriever = SessionRetriever(chunks, get_embedder())
+    retriever = SessionRetriever(chunks)
     agent = build_session_agent(retriever)
     corpus_hash = hashlib.sha256("|".join(file_hash_parts).encode("utf-8")).hexdigest()
     return SessionData(

requirements.txt

@@ -6,6 +6,5 @@ bitsandbytes
 sentence-transformers
 sentencepiece
 protobuf
-faiss-cpu
 pymupdf
 pypdf