Update knowledge_engine.py

knowledge_engine.py

@@ -1,10 +1,11 @@
 import os
 import pickle
-from typing import Dict, Any
+from typing import Dict, Any, List
 from datetime import datetime
 from concurrent.futures import ThreadPoolExecutor
 
 from config import Config
+from langchain_core.documents import Document
 from langchain_community.document_loaders import TextLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
@@ -12,7 +13,7 @@ from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
 from langchain.retrievers import BM25Retriever
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.llms import HuggingFaceHub
+from langchain_huggingface import HuggingFaceEndpoint
 
 class KnowledgeManager:
     def __init__(self):
@@ -23,18 +24,18 @@ class KnowledgeManager:
 
     def _init_embeddings(self):
         print("[i] Using Hugging Face embeddings")
-        return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        return HuggingFaceEmbeddings(
+            model_name="sentence-transformers/all-mpnet-base-v2",
+            model_kwargs={'device': 'cpu'}
+        )
 
     def _init_llm(self):
-        print("[i] Using HuggingFaceHub with Mistral-7B")
-        return HuggingFaceHub(
+        print("[i] Using HuggingFaceEndpoint with Mistral-7B")
+        return HuggingFaceEndpoint(
             repo_id="mistralai/Mistral-7B-Instruct-v0.1",
-            huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"),
-            model_kwargs={
-                "temperature": 0.1,
-                "max_new_tokens": 512,
-                "do_sample": True
-            }
+            temperature=0.1,
+            max_length=512,
+            token=os.getenv("HUGGINGFACEHUB_API_TOKEN")
         )
 
     def _init_retrievers(self):
@@ -76,7 +77,11 @@ class KnowledgeManager:
         )
         chunks = splitter.split_documents(docs)
 
-        vector_db = FAISS.from_documents(chunks, self.embeddings)
+        vector_db = FAISS.from_documents(
+            chunks, 
+            self.embeddings,
+            distance_strategy="COSINE"  # Ensures scores between 0-1
+        )
         vector_db.save_local(str(Config.VECTOR_STORE_PATH))
 
         bm25_retriever = BM25Retriever.from_documents(chunks)
@@ -92,16 +97,19 @@ class KnowledgeManager:
         with open(Config.KNOWLEDGE_DIR / "sirraya_xbrain.txt", "w", encoding="utf-8") as f:
             f.write(default_text)
 
-    def _parallel_retrieve(self, question: str):
+    def _parallel_retrieve(self, question: str) -> List[Document]:
         def retrieve_with_bm25():
-            return self.bm25_retriever.get_relevant_documents(question)
+            return self.bm25_retriever.invoke(question)  # Updated to use invoke()
 
         def retrieve_with_vector():
             retriever = self.vector_db.as_retriever(
                 search_type="similarity_score_threshold",
-                search_kwargs={"k": Config.MAX_CONTEXT_CHUNKS, "score_threshold": 0.83}
+                search_kwargs={
+                    "k": Config.MAX_CONTEXT_CHUNKS,
+                    "score_threshold": 0.83
+                }
             )
-            return retriever.get_relevant_documents(question)
+            return retriever.invoke(question)  # Updated to use invoke()
 
         with ThreadPoolExecutor(max_workers=2) as executor:
             bm25_future = executor.submit(retrieve_with_bm25)
@@ -132,7 +140,7 @@ Answer:"""
         return RetrievalQA.from_chain_type(
             llm=self._init_llm(),
             chain_type="stuff",
-            retriever=self.vector_db.as_retriever(search_kwargs={"k": 1}),
+            retriever=self.vector_db.as_retriever(search_kwargs={"k": Config.MAX_CONTEXT_CHUNKS}),
             chain_type_kwargs={
                 "prompt": PromptTemplate(
                     template=prompt_template,
@@ -155,10 +163,12 @@ Answer:"""
             docs = self._parallel_retrieve(question)
 
             if not docs:
-                retriever = self.vector_db.as_retriever(search_kwargs={"k": Config.MAX_CONTEXT_CHUNKS})
-                docs = retriever.get_relevant_documents(question)
+                retriever = self.vector_db.as_retriever(
+                    search_kwargs={"k": Config.MAX_CONTEXT_CHUNKS}
+                )
+                docs = retriever.invoke(question)  # Updated to use invoke()
 
-            result = self.qa_chain.invoke({"input_documents": docs, "query": question})
+            result = self.qa_chain.invoke({"query": question, "input_documents": docs})
             processing_time = (datetime.now() - start_time).total_seconds() * 1000
 
             return {
@@ -169,7 +179,7 @@ Answer:"""
         except Exception as e:
             print(f"[!] Query error: {e}")
             return {
-                "answer": f"Error: {e}",
+                "answer": f"Error processing your query: {str(e)}",
                 "processing_time": 0,
                 "source_chunks": []
             }