Rollback: stable local RAG

config.py

@@ -35,4 +35,6 @@ LLM_MODEL = "google/flan-t5-small"
 # Text splitting
 # -----------------------------
 CHUNK_SIZE = 500
-CHUNK_OVERLAP = 50
+CHUNK_OVERLAP = 50
+
+KB_DIR = "./kb"

rag.py

@@ -3,115 +3,83 @@
 import os
 from typing import List, Tuple
 
-from huggingface_hub import hf_hub_download, list_repo_files
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.schema import Document
 from transformers import pipeline
 
 from config import (
-    HF_DATASET_REPO,
+    KB_DIR,
+    CHROMA_DIR,
     EMBEDDING_MODEL,
     LLM_MODEL,
-    CHROMA_DIR,
-    CHUNK_SIZE,
-    CHUNK_OVERLAP,
 )
 
 # -----------------------------
-# Load PDFs from HF Dataset repo
+# Load documents
 # -----------------------------
-def load_documents():
+def load_documents() -> List[Document]:
     docs = []
 
-    try:
-        files = list_repo_files(
-            repo_id=HF_DATASET_REPO,
-            repo_type="dataset"
-        )
-    except Exception as e:
-        print("❌ Could not access dataset:", e)
-        return []
-
-    pdf_files = [f for f in files if f.lower().endswith(".pdf")]
-
-    if not pdf_files:
-        print("⚠️ No PDFs found in dataset")
-        return []
-
-    os.makedirs("kb", exist_ok=True)
+    if not os.path.exists(KB_DIR):
+        print(f"⚠️ KB_DIR not found: {KB_DIR}")
+        return docs
 
-    for pdf in pdf_files:
-        local_path = hf_hub_download(
-            repo_id=HF_DATASET_REPO,
-            filename=pdf,
-            repo_type="dataset"
-        )
-
-        loader = PyPDFLoader(local_path)
-        docs.extend(loader.load())
+    for file in os.listdir(KB_DIR):
+        if file.lower().endswith(".pdf"):
+            loader = PyPDFLoader(os.path.join(KB_DIR, file))
+            docs.extend(loader.load())
 
     return docs
 
 
 # -----------------------------
-# Build vector DB (safe)
+# Build vector DB (once)
 # -----------------------------
-def build_vectorstore():
-    documents = load_documents()
-
-    if not documents:
-        print("⚠️ No documents loaded, vector DB will be empty")
-        return None
-
-    splitter = RecursiveCharacterTextSplitter(
-        chunk_size=CHUNK_SIZE,
-        chunk_overlap=CHUNK_OVERLAP,
-    )
-
-    splits = splitter.split_documents(documents)
+documents = load_documents()
 
-    embeddings = HuggingFaceEmbeddings(
-        model_name=EMBEDDING_MODEL
-    )
+splitter = RecursiveCharacterTextSplitter(
+    chunk_size=800,
+    chunk_overlap=100
+)
 
-    vectordb = Chroma.from_documents(
-        documents=splits,
-        embedding=embeddings,
-        persist_directory=CHROMA_DIR
-    )
+chunks = splitter.split_documents(documents)
 
-    return vectordb
+embeddings = HuggingFaceEmbeddings(
+    model_name=EMBEDDING_MODEL
+)
 
+vectordb = Chroma.from_documents(
+    documents=chunks,
+    embedding=embeddings,
+    persist_directory=CHROMA_DIR
+)
 
-# Build once at startup
-VECTOR_DB = build_vectorstore()
+retriever = vectordb.as_retriever(search_kwargs={"k": 3})
 
 # -----------------------------
-# LLM (CPU-safe)
+# LLM (CORRECT task)
 # -----------------------------
-qa_pipeline = pipeline(
+llm = pipeline(
     "text2text-generation",
     model=LLM_MODEL,
-    max_new_tokens=256
+    device=-1
 )
 
-
 # -----------------------------
-# Public API
+# RAG call
 # -----------------------------
-def ask_rag_with_status(question: str) -> Tuple[str, List[str]]:
+def ask_rag_with_status(question: str) -> Tuple[str, list]:
     status = []
 
-    if VECTOR_DB is None:
-        return "No documents available.", ["Vector DB not initialized"]
+    if vectordb._collection.count() == 0:
+        return "Knowledge base is empty.", ["No documents indexed"]
 
-    retriever = VECTOR_DB.as_retriever(search_kwargs={"k": 3})
     docs = retriever.get_relevant_documents(question)
 
-    if not docs:
-        return "No relevant information found.", ["No matching chunks"]
+    status.append(f"Retrieved {len(docs)} chunks")
 
     context = "\n\n".join(d.page_content for d in docs)
 
@@ -123,11 +91,10 @@ Context:
 
 Question:
 {question}
-"""
 
-    result = qa_pipeline(prompt)[0]["generated_text"]
+Answer:
+"""
 
-    status.append(f"Retrieved {len(docs)} chunks")
-    status.append("Answer generated")
+    result = llm(prompt, max_new_tokens=256)[0]["generated_text"]
 
     return result.strip(), status