Update src/services/rag_service.py

src/services/rag_service.py

@@ -1,13 +1,224 @@
+import os
+import time
+import logging
+from pathlib import Path
 from typing import List, Optional
+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
+
 
 class RAGService:
     def __init__(self):
         self.vectorstore = None
-    
+        self.bm25 = None
+        self.bm25_corpus: List[str] = []
+        self.bm25_docs: List[dict] = []   # parallel to bm25_corpus
+        self._embeddings = None
+        self._ready = False
+
     def is_ready(self) -> bool:
-        return self.vectorstore is not None
-    
-    async def hybrid_search(self, query: str, top_k: int = 7) -> List:
-        """Placeholder for hybrid search"""
-        # Will implement FAISS + BM25 here
-        return []
+        return self._ready
+
+    # ── Startup ────────────────────────────────────────────────────────────
+
+    async def initialize(self, openai_api_key: str):
+        """Call once from lifespan. Loads files, chunks, vectorizes."""
+        t0 = time.time()
+        logger.info("RAGService: starting initialization...")
+
+        try:
+            from langchain_openai import OpenAIEmbeddings
+            from langchain_community.vectorstores import FAISS
+            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
+            )
+
+            # 1. Load and chunk all .txt files
+            chunks, metadatas = self._load_and_chunk_all()
+            if not chunks:
+                logger.error("RAGService: no chunks loaded — check KNOWLEDGE_DIR path")
+                return
+
+            logger.info(f"RAGService: loaded {len(chunks)} chunks from {KNOWLEDGE_DIR}")
+
+            # 2. Build FAISS index
+            self.vectorstore = await FAISS.afrom_texts(
+                chunks,
+                self._embeddings,
+                metadatas=metadatas
+            )
+
+            # 3. Build BM25 index (parallel corpus)
+            self.bm25_corpus = chunks
+            self.bm25_docs = metadatas
+            tokenized = [c.lower().split() for c in chunks]
+            self.bm25 = BM25Okapi(tokenized)
+
+            self._ready = True
+            logger.info(
+                f"RAGService: ready — {len(chunks)} chunks indexed in {time.time()-t0:.1f}s"
+            )
+
+        except Exception as e:
+            logger.error(f"RAGService initialization failed: {e}", exc_info=True)
+
+    # ── 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 = []
+
+        if not KNOWLEDGE_DIR.exists():
+            logger.error(f"KNOWLEDGE_DIR not found: {KNOWLEDGE_DIR}")
+            return [], []
+
+        txt_files = sorted(KNOWLEDGE_DIR.glob("*.txt"))
+        logger.info(f"RAGService: found {len(txt_files)} .txt files")
+
+        for fpath in txt_files:
+            try:
+                text = fpath.read_text(encoding="utf-8").strip()
+                if not text:
+                    continue
+
+                # Extract SOURCE and TITLE from header if present
+                source_url = ""
+                title = fpath.stem.replace("_", " ").replace("-", " ")
+                for line in text.splitlines()[:5]:
+                    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:
+                    all_chunks.append(chunk)
+                    all_meta.append({
+                        "source": source_url,
+                        "title": title,
+                        "filename": fpath.name,
+                    })
+
+            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]:
+        """
+        Split on TOPIC: headers first (respects your file structure),
+        then hard-split oversized sections by character count.
+        """
+        chunks = []
+
+        # Split on TOPIC: section markers present in your files
+        import re
+        sections = re.split(r"\n(?=TOPIC:|[A-Z][A-Z\s/&]+:)", text)
+
+        for section in sections:
+            section = section.strip()
+            if not section:
+                continue
+
+            if len(section) <= CHUNK_SIZE * 4:  # ~1600 chars fits in one chunk
+                chunks.append(section)
+            else:
+                # Hard split with overlap for large sections
+                chars = section
+                start = 0
+                while start < len(chars):
+                    end = start + CHUNK_SIZE * 4
+                    chunk = chars[start:end].strip()
+                    if chunk:
+                        chunks.append(chunk)
+                    start += (CHUNK_SIZE * 4) - (CHUNK_OVERLAP * 4)
+
+        return chunks
+
+    # ── Retrieval ──────────────────────────────────────────────────────────
+
+    async def hybrid_search(
+        self,
+        query: str,
+        top_k: int = 5,
+        alpha: float = 0.6   # 0=pure BM25, 1=pure FAISS
+    ) -> List[dict]:
+        """
+        Reciprocal Rank Fusion of FAISS (dense) + BM25 (sparse).
+        Returns list of {content, source, title, filename, score}
+        """
+        if not self._ready:
+            logger.warning("hybrid_search called before RAGService is ready")
+            return []
+
+        try:
+            # Dense retrieval
+            dense_results = await self.vectorstore.asimilarity_search_with_score(
+                query, k=top_k * 3
+            )
+
+            # Sparse retrieval
+            tokenized_query = query.lower().split()
+            bm25_scores = self.bm25.get_scores(tokenized_query)
+            top_bm25_idx = np.argsort(bm25_scores)[::-1][:top_k * 3]
+
+            # Reciprocal Rank Fusion
+            rrf_scores: dict[int, float] = {}
+            K = 60  # RRF constant
+
+            # Score dense results
+            for rank, (doc, _) in enumerate(dense_results):
+                # Find corpus index by content match
+                idx = self._find_corpus_index(doc.page_content)
+                if idx >= 0:
+                    rrf_scores[idx] = rrf_scores.get(idx, 0) + alpha * (1 / (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))
+
+            # Sort by RRF score and return top_k
+            sorted_idx = sorted(rrf_scores, key=lambda i: rrf_scores[i], reverse=True)[:top_k]
+
+            results = []
+            for idx in sorted_idx:
+                if idx < len(self.bm25_corpus):
+                    meta = self.bm25_docs[idx]
+                    results.append({
+                        "content": self.bm25_corpus[idx],
+                        "source": meta.get("source", ""),
+                        "title": meta.get("title", ""),
+                        "filename": meta.get("filename", ""),
+                        "score": rrf_scores[idx],
+                    })
+
+            return results
+
+        except Exception as e:
+            logger.error(f"hybrid_search error: {e}", exc_info=True)
+            return []
+
+    def _find_corpus_index(self, content: str) -> int:
+        """Find index in bm25_corpus by content match."""
+        content_strip = content.strip()
+        for i, c in enumerate(self.bm25_corpus):
+            if c.strip() == content_strip:
+                return i
+        return -1
+
+    def get_stats(self) -> dict:
+        return {
+            "ready": self._ready,
+            "chunk_count": len(self.bm25_corpus),
+            "vectorstore_loaded": self.vectorstore is not None,
+            "bm25_loaded": self.bm25 is not None,
+        }