Update src/services/rag_service.py

src/services/rag_service.py

@@ -1,4 +1,4 @@
-import os
+import re
 import time
 import logging
 from pathlib import Path
@@ -8,8 +8,8 @@ import numpy as np
 logger = logging.getLogger(__name__)
 
 KNOWLEDGE_DIR = Path("/app/data/knowledge")
-CHUNK_SIZE = 400        # tokens approx — ~1600 chars
-CHUNK_OVERLAP = 80      # chars overlap between chunks
+CHUNK_MAX_CHARS = 1600   # ~400 tokens — fits well in context without over-chunking
+CHUNK_OVERLAP_CHARS = 200
 
 
 class RAGService:
@@ -17,7 +17,7 @@ class RAGService:
         self.vectorstore = None
         self.bm25 = None
         self.bm25_corpus: List[str] = []
-        self.bm25_docs: List[dict] = []   # parallel to bm25_corpus
+        self.bm25_meta: List[dict] = []
         self._embeddings = None
         self._ready = False
 
@@ -27,7 +27,11 @@ class RAGService:
     # ── Startup ────────────────────────────────────────────────────────────
 
     async def initialize(self, openai_api_key: str):
-        """Call once from lifespan. Loads files, chunks, vectorizes."""
+        """
+        Called once from app.py lifespan at startup.
+        Loads all .txt files from KNOWLEDGE_DIR, chunks them,
+        builds FAISS (dense) + BM25 (sparse) indexes.
+        """
         t0 = time.time()
         logger.info("RAGService: starting initialization...")
 
@@ -37,34 +41,38 @@ class RAGService:
             from rank_bm25 import BM25Okapi
 
             self._embeddings = OpenAIEmbeddings(
-                model="text-embedding-3-small",   # cheaper + fast, plenty for KB this size
-                openai_api_key=openai_api_key
+                model="text-embedding-3-small",
+                openai_api_key=openai_api_key,
             )
 
-            # 1. Load and chunk all .txt files
+            # 1. Load and chunk
             chunks, metadatas = self._load_and_chunk_all()
             if not chunks:
-                logger.error("RAGService: no chunks loaded — check KNOWLEDGE_DIR path")
+                logger.error(
+                    "RAGService: no chunks loaded. "
+                    f"Check that {KNOWLEDGE_DIR} exists and contains .txt files."
+                )
                 return
 
-            logger.info(f"RAGService: loaded {len(chunks)} chunks from {KNOWLEDGE_DIR}")
+            logger.info(f"RAGService: {len(chunks)} chunks from {KNOWLEDGE_DIR}")
 
-            # 2. Build FAISS index
+            # 2. FAISS index (calls OpenAI embeddings API)
             self.vectorstore = await FAISS.afrom_texts(
                 chunks,
                 self._embeddings,
-                metadatas=metadatas
+                metadatas=metadatas,
             )
 
-            # 3. Build BM25 index (parallel corpus)
+            # 3. BM25 index (in-memory, instant)
             self.bm25_corpus = chunks
-            self.bm25_docs = metadatas
+            self.bm25_meta = metadatas
             tokenized = [c.lower().split() for c in chunks]
             self.bm25 = BM25Okapi(tokenized)
 
             self._ready = True
+            elapsed = time.time() - t0
             logger.info(
-                f"RAGService: ready — {len(chunks)} chunks indexed in {time.time()-t0:.1f}s"
+                f"RAGService: ✅ ready — {len(chunks)} chunks indexed in {elapsed:.1f}s"
             )
 
         except Exception as e:
@@ -73,9 +81,8 @@ class RAGService:
     # ── Document Loading & Chunking ────────────────────────────────────────
 
     def _load_and_chunk_all(self) -> tuple[List[str], List[dict]]:
-        """Load all .txt files and split into overlapping chunks."""
-        all_chunks = []
-        all_meta = []
+        all_chunks: List[str] = []
+        all_meta: List[dict] = []
 
         if not KNOWLEDGE_DIR.exists():
             logger.error(f"KNOWLEDGE_DIR not found: {KNOWLEDGE_DIR}")
@@ -90,17 +97,18 @@ class RAGService:
                 if not text:
                     continue
 
-                # Extract SOURCE and TITLE from header if present
+                # Extract SOURCE and TITLE from header lines
                 source_url = ""
                 title = fpath.stem.replace("_", " ").replace("-", " ")
-                for line in text.splitlines()[:5]:
+                for line in text.splitlines()[:6]:
+                    line = line.strip()
                     if line.startswith("SOURCE:"):
                         source_url = line.replace("SOURCE:", "").strip()
                     elif line.startswith("TITLE:"):
                         title = line.replace("TITLE:", "").strip()
 
-                chunks = self._chunk_text(text, fpath.name)
-                for chunk in chunks:
+                file_chunks = self._chunk_text(text)
+                for chunk in file_chunks:
                     all_chunks.append(chunk)
                     all_meta.append({
                         "source": source_url,
@@ -108,39 +116,45 @@ class RAGService:
                         "filename": fpath.name,
                     })
 
+                logger.info(f"  {fpath.name}: {len(file_chunks)} chunks")
+
             except Exception as e:
                 logger.warning(f"Failed to load {fpath.name}: {e}")
 
         return all_chunks, all_meta
 
-    def _chunk_text(self, text: str, filename: str) -> List[str]:
+    def _chunk_text(self, text: str) -> List[str]:
         """
-        Split on TOPIC: headers first (respects your file structure),
-        then hard-split oversized sections by character count.
+        Split on TOPIC: headers (matching the structure of your .txt files),
+        then hard-split oversized sections with overlap.
         """
-        chunks = []
+        chunks: List[str] = []
 
-        # Split on TOPIC: section markers present in your files
-        import re
-        sections = re.split(r"\n(?=TOPIC:|[A-Z][A-Z\s/&]+:)", text)
+        # Split on lines that look like section headers
+        # Matches: "TOPIC: ...", "BORROWING AND LOANS:", "STUDY ROOMS:", etc.
+        sections = re.split(
+            r"\n(?=(?:TOPIC:|[A-Z][A-Z\s,/&()\-]+:)\s)",
+            text
+        )
 
         for section in sections:
             section = section.strip()
-            if not section:
+            if not section or len(section) < 30:  # skip empty/trivial sections
                 continue
 
-            if len(section) <= CHUNK_SIZE * 4:  # ~1600 chars fits in one chunk
+            if len(section) <= CHUNK_MAX_CHARS:
                 chunks.append(section)
             else:
-                # Hard split with overlap for large sections
-                chars = section
+                # Slide window with overlap for large sections
                 start = 0
-                while start < len(chars):
-                    end = start + CHUNK_SIZE * 4
-                    chunk = chars[start:end].strip()
+                while start < len(section):
+                    end = start + CHUNK_MAX_CHARS
+                    chunk = section[start:end].strip()
                     if chunk:
                         chunks.append(chunk)
-                    start += (CHUNK_SIZE * 4) - (CHUNK_OVERLAP * 4)
+                    if end >= len(section):
+                        break
+                    start += CHUNK_MAX_CHARS - CHUNK_OVERLAP_CHARS
 
         return chunks
 
@@ -150,7 +164,7 @@ class RAGService:
         self,
         query: str,
         top_k: int = 5,
-        alpha: float = 0.6   # 0=pure BM25, 1=pure FAISS
+        alpha: float = 0.6,   # 0 = pure BM25, 1 = pure FAISS dense
     ) -> List[dict]:
         """
         Reciprocal Rank Fusion of FAISS (dense) + BM25 (sparse).
@@ -161,44 +175,43 @@ class RAGService:
             return []
 
         try:
-            # Dense retrieval
+            fetch_k = top_k * 3  # over-fetch then re-rank
+
+            # ── Dense retrieval (FAISS) ────────────────────────────────────
             dense_results = await self.vectorstore.asimilarity_search_with_score(
-                query, k=top_k * 3
+                query, k=fetch_k
             )
 
-            # Sparse retrieval
+            # ── Sparse retrieval (BM25) ───────────────────────────────────
             tokenized_query = query.lower().split()
             bm25_scores = self.bm25.get_scores(tokenized_query)
-            top_bm25_idx = np.argsort(bm25_scores)[::-1][:top_k * 3]
+            top_bm25_idx = np.argsort(bm25_scores)[::-1][:fetch_k].tolist()
 
-            # Reciprocal Rank Fusion
-            rrf_scores: dict[int, float] = {}
-            K = 60  # RRF constant
+            # ── Reciprocal Rank Fusion ────────────────────────────────────
+            K = 60  # RRF constant — standard value
+            rrf: dict[int, float] = {}
 
-            # Score dense results
-            for rank, (doc, _) in enumerate(dense_results):
-                # Find corpus index by content match
+            for rank, (doc, _score) in enumerate(dense_results):
                 idx = self._find_corpus_index(doc.page_content)
                 if idx >= 0:
-                    rrf_scores[idx] = rrf_scores.get(idx, 0) + alpha * (1 / (K + rank + 1))
+                    rrf[idx] = rrf.get(idx, 0.0) + alpha * (1.0 / (K + rank + 1))
 
-            # Score BM25 results
             for rank, idx in enumerate(top_bm25_idx):
-                rrf_scores[int(idx)] = rrf_scores.get(int(idx), 0) + (1 - alpha) * (1 / (K + rank + 1))
+                rrf[idx] = rrf.get(idx, 0.0) + (1 - alpha) * (1.0 / (K + rank + 1))
 
-            # Sort by RRF score and return top_k
-            sorted_idx = sorted(rrf_scores, key=lambda i: rrf_scores[i], reverse=True)[:top_k]
+            # ── Return top_k sorted by RRF score ─────────────────────────
+            sorted_idx = sorted(rrf, key=lambda i: rrf[i], reverse=True)[:top_k]
 
             results = []
             for idx in sorted_idx:
                 if idx < len(self.bm25_corpus):
-                    meta = self.bm25_docs[idx]
+                    meta = self.bm25_meta[idx]
                     results.append({
                         "content": self.bm25_corpus[idx],
                         "source": meta.get("source", ""),
                         "title": meta.get("title", ""),
                         "filename": meta.get("filename", ""),
-                        "score": rrf_scores[idx],
+                        "score": round(rrf[idx], 6),
                     })
 
             return results
@@ -208,7 +221,7 @@ class RAGService:
             return []
 
     def _find_corpus_index(self, content: str) -> int:
-        """Find index in bm25_corpus by content match."""
+        """Find index in bm25_corpus by exact content match."""
         content_strip = content.strip()
         for i, c in enumerate(self.bm25_corpus):
             if c.strip() == content_strip:
@@ -219,6 +232,7 @@ class RAGService:
         return {
             "ready": self._ready,
             "chunk_count": len(self.bm25_corpus),
+            "knowledge_dir": str(KNOWLEDGE_DIR),
             "vectorstore_loaded": self.vectorstore is not None,
             "bm25_loaded": self.bm25 is not None,
-        }
+        }