Update knowledge_engine.py

knowledge_engine.py

@@ -6,26 +6,40 @@ from concurrent.futures import ThreadPoolExecutor
 
 from config import Config
 
+# Setup Hugging Face token securely (Make sure to set this in your environment securely)
+# os.environ["HUGGINGFACEHUB_API_TOKEN"] = "your_token_here"
+
 # Core ML/AI libraries
 from langchain_community.document_loaders import TextLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import OllamaEmbeddings
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
-from langchain_community.llms import Ollama
-from langchain_community.retrievers import BM25Retriever
+from langchain.retrievers import BM25Retriever
 
+# Only use Hugging Face embeddings and LLM (no Ollama fallback)
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.llms import HuggingFaceHub
 
-class KnowledgeManager:
-    """Main knowledge management class handling document processing and Q&A with CoT & MoE routing"""
 
+class KnowledgeManager:
     def __init__(self):
         Config.setup_dirs()
-        self.embeddings = OllamaEmbeddings(model="mxbai-embed-large")
+        self.embeddings = self._init_embeddings()
         self.vector_db, self.bm25_retriever = self._init_retrievers()
         self.qa_chain = self._create_moe_qa_chain()
 
+    def _init_embeddings(self):
+        print("[i] Using Hugging Face embeddings")
+        return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+    def _init_llm(self):
+        print("[i] Using Hugging Face LLM")
+        return HuggingFaceHub(
+            repo_id="tiiuae/falcon-7b-instruct",
+            model_kwargs={"temperature": 0.1, "max_new_tokens": 512}
+        )
+
     def _init_retrievers(self):
         faiss_index_path = Config.VECTOR_STORE_PATH / "index.faiss"
         faiss_pkl_path = Config.VECTOR_STORE_PATH / "index.pkl"
@@ -42,7 +56,7 @@ class KnowledgeManager:
                         bm25_retriever = pickle.load(f)
                     return vector_db, bm25_retriever
             except Exception as e:
-                print(f"[!] Error loading existing vector store: {e}. Rebuilding...")
+                print(f"[!] Error loading vector store: {e}. Rebuilding...")
 
         return self._build_retrievers_from_documents()
 
@@ -77,18 +91,15 @@ class KnowledgeManager:
         return vector_db, bm25_retriever
 
     def _create_default_knowledge(self):
-        default_text = """Sirraya xBrain - Advanced AI Platform\n\nCreated by Amir Hameed.\n\nFeatures:\n- Hybrid Retrieval (Vector + BM25)\n- LISA Assistant\n- FAISS, Ollama, BM25 Integration"""
+        default_text = """Sirraya xBrain - Advanced AI Platform\n\nCreated by Amir Hameed.\n\nFeatures:\n- Hybrid Retrieval (Vector + BM25)\n- LISA Assistant\n- FAISS, BM25 Integration"""
         with open(Config.KNOWLEDGE_DIR / "sirraya_xbrain.txt", "w", encoding="utf-8") as f:
             f.write(default_text)
 
     def _parallel_retrieve(self, question: str):
-        """Parallel retrieval execution: simulates Mixture of Experts routing"""
-
         def retrieve_with_bm25():
             return self.bm25_retriever.get_relevant_documents(question)
 
         def retrieve_with_vector():
-            # Lowered threshold to 0.3 for better doc retrieval (adjust as needed)
             retriever = self.vector_db.as_retriever(
                 search_type="similarity_score_threshold",
                 search_kwargs={"k": Config.MAX_CONTEXT_CHUNKS, "score_threshold": 0.83}
@@ -101,7 +112,6 @@ class KnowledgeManager:
             bm25_results = bm25_future.result()
             vector_results = vector_future.result()
 
-        # Combine results; duplicates are possible, consider deduplication if needed
         return vector_results + bm25_results
 
     def _create_moe_qa_chain(self):
@@ -123,9 +133,9 @@ Instructions:
 Answer:"""
 
         return RetrievalQA.from_chain_type(
-            llm=Ollama(model="phi", temperature=0.1),
+            llm=self._init_llm(),
             chain_type="stuff",
-            retriever=self.vector_db.as_retriever(search_kwargs={"k": 1}),  # Dummy retriever to satisfy LangChain
+            retriever=self.vector_db.as_retriever(search_kwargs={"k": 1}),
             chain_type_kwargs={
                 "prompt": PromptTemplate(
                     template=prompt_template,
@@ -136,7 +146,6 @@ Answer:"""
         )
 
     def query(self, question: str) -> Dict[str, Any]:
-        """Query system using CoT + MoE logic"""
         if not self.qa_chain:
             return {
                 "answer": "Knowledge system not initialized. Please reload.",
@@ -148,14 +157,11 @@ Answer:"""
             start_time = datetime.now()
             docs = self._parallel_retrieve(question)
 
-            # If no docs found, fallback to retriever without threshold for testing
             if not docs:
                 retriever = self.vector_db.as_retriever(search_kwargs={"k": Config.MAX_CONTEXT_CHUNKS})
                 docs = retriever.get_relevant_documents(question)
 
-            # Use invoke() for chains with multiple outputs
             result = self.qa_chain.invoke({"input_documents": docs, "query": question})
-
             processing_time = (datetime.now() - start_time).total_seconds() * 1000
 
             return {