Update knowledge_engine.py

knowledge_engine.py

@@ -1,112 +1,61 @@
 import os
 from pathlib import Path
-from typing import List, Optional
-
-import faiss
-import numpy as np
-from sentence_transformers import SentenceTransformer
-
-from langchain.llms import HuggingFacePipeline
-from langchain.chains import RetrievalQA
-from langchain.vectorstores.faiss import FAISS
-from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.document_loaders import TextLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-
-import torch
-from transformers import pipeline
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.chains import RetrievalQA
+from langchain.llms import HuggingFaceHub
 
 class KnowledgeManager:
-    def __init__(self, knowledge_dir="knowledge_base"):
+    def __init__(self, knowledge_dir="."):  # root dir by default
         self.knowledge_dir = Path(knowledge_dir)
-        self.knowledge_dir.mkdir(exist_ok=True, parents=True)
-
         self.documents = []
-        self.texts = []
+        self.embeddings = None
         self.vectorstore = None
         self.retriever = None
-        self.qa_chain = None
         self.llm = None
+        self.qa_chain = None
 
-        self.device = "cpu"  # For HF Spaces, CPU only
-
-        # Initialize embeddings
-        self.embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-
-        # Load and prepare knowledge
-        self.load_documents()
-        self.create_vectorstore()
-        self.init_llm()
-        self.init_qa_chain()
-
-    def load_documents(self):
-        # Load text files and split into chunks
+        self._load_documents()
+        if self.documents:
+            self._initialize_embeddings()
+            self._initialize_vectorstore()
+            self._initialize_llm()
+            self._initialize_qa_chain()
+
+    def _load_documents(self):
+        if not self.knowledge_dir.exists():
+            raise FileNotFoundError(f"Directory {self.knowledge_dir} does not exist.")
+        
         files = list(self.knowledge_dir.glob("*.txt"))
-        self.documents = []
+        if not files:
+            raise FileNotFoundError(f"No .txt files found in {self.knowledge_dir}. Please upload your knowledge base files in root.")
+        
         for file in files:
-            loader = TextLoader(str(file), encoding="utf-8")
-            docs = loader.load()
-            self.documents.extend(docs)
-
-        # Split into smaller chunks (to improve retrieval granularity)
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-        self.texts = text_splitter.split_documents(self.documents)
+            loader = TextLoader(str(file))
+            self.documents.extend(loader.load())
+        
+        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+        self.documents = splitter.split_documents(self.documents)
 
-    def create_vectorstore(self):
-        if not self.texts:
-            self.vectorstore = None
-            return
-        self.vectorstore = FAISS.from_documents(self.texts, self.embeddings)
-        self.retriever = self.vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
-
-    def init_llm(self):
-        # Initialize HuggingFace pipeline + LangChain wrapper LLM
+    def _initialize_embeddings(self):
+        self.embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
-        # Try flan-t5-small first
-        try:
-            pipe = pipeline(
-                "text2text-generation",
-                model="google/flan-t5-small",
-                device=-1,  # CPU only
-                max_length=256,
-                do_sample=False,
-            )
-            self.llm = HuggingFacePipeline(pipeline=pipe)
-        except Exception as e:
-            print(f"Failed to load flan-t5-small: {e}")
-            self.llm = None
+    def _initialize_vectorstore(self):
+        self.vectorstore = FAISS.from_documents(self.documents, self.embeddings)
+        self.retriever = self.vectorstore.as_retriever()
 
-        # Fallback: if no LLM, set to None and warn
-        if self.llm is None:
-            print("No LLM available, will fallback to retrieval-only.")
+    def _initialize_llm(self):
+        self.llm = HuggingFaceHub(repo_id="google/flan-t5-small", model_kwargs={"temperature":0, "max_length":256})
 
-    def init_qa_chain(self):
-        if self.llm and self.retriever:
-            self.qa_chain = RetrievalQA.from_chain_type(
-                llm=self.llm,
-                retriever=self.retriever,
-                return_source_documents=True,
-                chain_type="stuff",  # Stuff all docs in prompt, or "map_reduce"
-            )
-        else:
-            self.qa_chain = None
+    def _initialize_qa_chain(self):
+        self.qa_chain = RetrievalQA.from_chain_type(llm=self.llm, chain_type="stuff", retriever=self.retriever)
 
-    def get_knowledge_summary(self) -> str:
-        count = len(self.texts) if self.texts else 0
-        return f"{count} document chunks loaded."
+    def ask(self, query):
+        if not self.qa_chain:
+            return "Knowledge base not initialized properly."
+        return self.qa_chain.run(query)
 
-    def query(self, question: str):
-        if self.qa_chain:
-            # Use LLM + retrieval
-            result = self.qa_chain({"query": question})
-            answer = result.get("result", "No answer found.")
-            sources = result.get("source_documents", [])
-            source_texts = [doc.page_content for doc in sources]
-            return answer, source_texts
-        elif self.retriever:
-            # Retrieval only fallback
-            docs = self.retriever.get_relevant_documents(question)
-            answers = [doc.page_content for doc in docs]
-            return "\n\n".join(answers), []
-        else:
-            return "Knowledge base not initialized.", []
+    def get_knowledge_summary(self):
+        return f"Loaded {len(self.documents)} document chunks from {self.knowledge_dir}"