Update knowledge_engine.py

knowledge_engine.py

@@ -1,315 +1,112 @@
 import os
-import tempfile
-import shutil
-from typing import Dict, List
-from datetime import datetime
-from concurrent.futures import ThreadPoolExecutor
-
-from langchain_core.documents import Document
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import FAISS
-from langchain.retrievers import BM25Retriever
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from pathlib import Path
+from typing import List, Optional
 
+import faiss
+import numpy as np
+from sentence_transformers import SentenceTransformer
 
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM, pipeline
 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
 
-class CPULLMProvider:
-    """CPU-based LLM provider using HuggingFace models"""
-
-    def __init__(self):
-        self.name = "CPU-LLM"
-        self.is_available = False
-        self.current_model = None
-
-        # CPU-friendly models
-        self.cpu_models = [
-            "google/flan-t5-small",          # Encoder-decoder model
-            "distilbert/distilgpt2"          # Decoder-only (GPT-style)
-        ]
-
-    def initialize(self) -> bool:
-        """Initialize the CPU LLM with the best available model"""
-        for model_id in self.cpu_models:
-            try:
-                print(f"[i] Trying to load {model_id}...")
-
-                tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-                # Detect model type based on name
-                if "flan" in model_id or "t5" in model_id:
-                    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-                    task = "text2text-generation"
-                else:
-                    model = AutoModelForCausalLM.from_pretrained(model_id)
-                    task = "text-generation"
-
-                pipe = pipeline(
-                    task,
-                    model=model,
-                    tokenizer=tokenizer,
-                    max_new_tokens=256,
-                    temperature=0.3,
-                    top_p=0.95,
-                    device="cpu"
-                )
-
-                self.llm = HuggingFacePipeline(pipeline=pipe)
-                self.current_model = model_id
-                self.is_available = True
-
-                # Test model
-                test_response = self.invoke("Hello, who are you?")
-                if test_response and len(test_response) > 0:
-                    print(f"[✓] Successfully loaded {model_id}")
-                    return True
-
-            except Exception as e:
-                print(f"[!] Failed to load {model_id}: {str(e)[:200]}...")
-                continue
-
-        print("[!] All CPU models failed to load")
-        return False
-
-    def invoke(self, prompt: str) -> str:
-        """Invoke the CPU model with prompt"""
-        if not self.llm:
-            raise Exception("CPU LLM not initialized")
-
-        try:
-            # Optionally modify prompt for specific models if needed
-            formatted_prompt = prompt
-            response = self.llm.invoke(formatted_prompt)
-            return response.strip()
-        except Exception as e:
-            print(f"[!] CPU model error: {e}")
-            raise
-
-
+import torch
+from transformers import pipeline
 
 class KnowledgeManager:
-    def __init__(self):
-        self.temp_dir = tempfile.mkdtemp()
-        self.setup_dirs()
-        self.embeddings = self.init_embeddings()
-        self.vector_db = None
-        self.bm25_retriever = None
-        self.llm_provider = CPULLMProvider()
-        self.knowledge_texts = []
-
-        self.init_system()
+    def __init__(self, knowledge_dir="knowledge_base"):
+        self.knowledge_dir = Path(knowledge_dir)
+        self.knowledge_dir.mkdir(exist_ok=True, parents=True)
+
+        self.documents = []
+        self.texts = []
+        self.vectorstore = None
+        self.retriever = None
+        self.qa_chain = None
+        self.llm = 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
+        files = list(self.knowledge_dir.glob("*.txt"))
+        self.documents = []
+        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)
+
+    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 setup_dirs(self):
-        """Setup temporary directories"""
-        self.knowledge_dir = os.path.join(self.temp_dir, "knowledge")
-        os.makedirs(self.knowledge_dir, exist_ok=True)
+    def init_llm(self):
+        # Initialize HuggingFace pipeline + LangChain wrapper LLM
 
-    def init_embeddings(self):
-        """Initialize CPU-friendly embeddings"""
+        # Try flan-t5-small first
         try:
-            return HuggingFaceEmbeddings(
-                model_name="sentence-transformers/all-MiniLM-L6-v2",
-                model_kwargs={'device': 'cpu'},
-                encode_kwargs={'normalize_embeddings': True}
+            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 embeddings: {e}")
-            return None
-
-    def init_system(self):
-        """Initialize the RAG system"""
-        print("[i] Initializing CPU LLM...")
-        if self.llm_provider.initialize():
-            print(f"[✓] Using model: {self.llm_provider.current_model}")
-        else:
-            print("[!] Continuing without LLM (retrieval only)")
-
-        # Load default knowledge
-        self._load_default_knowledge()
-
-        # Build retrievers
-        self.build_retrievers()
-
-    def _load_default_knowledge(self):
-        """Load default knowledge base"""
-        default_content = """Sirraya xBrain - CPU-based AI Platform
-
-Features:
-- Uses efficient CPU-based language models like Phi-2
-- Implements RAG (Retrieval-Augmented Generation)
-- Combines vector search and keyword retrieval
-- Optimized for CPU-only environments
-
-Technical Details:
-- Embeddings: all-MiniLM-L6-v2
-- Vector Store: FAISS
-- Keyword Retrieval: BM25
-- LLM: Microsoft Phi-2 or similar CPU-friendly models"""
-
-        self.knowledge_texts = [{
-            "filename": "default_knowledge.txt",
-            "content": default_content
-        }]
-
-        # Save to file
-        with open(os.path.join(self.knowledge_dir, "default_knowledge.txt"), "w") as f:
-            f.write(default_content)
-
-    def build_retrievers(self):
-        """Build the retrieval components"""
-        if not self.embeddings:
-            print("[!] No embeddings available")
-            return
-
-        try:
-            # Create documents
-            documents = [
-                Document(
-                    page_content=text["content"],
-                    metadata={"source": text["filename"]}
-                )
-                for text in self.knowledge_texts
-            ]
-
-            # Split documents
-            splitter = RecursiveCharacterTextSplitter(
-                chunk_size=512,
-                chunk_overlap=128,
-                separators=["\n\n", "\n", ". ", "! ", "? ", "; ", " ", ""]
+            print(f"Failed to load flan-t5-small: {e}")
+            self.llm = None
+
+        # Fallback: if no LLM, set to None and warn
+        if self.llm is None:
+            print("No LLM available, will fallback to retrieval-only.")
+
+    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"
             )
-            chunks = splitter.split_documents(documents)
-
-            # Create vector store
-            self.vector_db = FAISS.from_documents(
-                chunks,
-                self.embeddings,
-                distance_strategy="COSINE"
-            )
-
-            # Create BM25 retriever
-            self.bm25_retriever = BM25Retriever.from_documents(chunks)
-            self.bm25_retriever.k = 3
-
-            print(f"[✓] Built retrievers with {len(chunks)} chunks")
-
-        except Exception as e:
-            print(f"[!] Error building retrievers: {e}")
-
-    def retrieve_documents(self, query: str) -> List[Document]:
-        """Retrieve relevant documents using both methods"""
-        if not self.vector_db or not self.bm25_retriever:
-            return []
-
-        def vector_search():
-            try:
-                return self.vector_db.similarity_search(query, k=2)
-            except:
-                return []
-
-        def bm25_search():
-            try:
-                return self.bm25_retriever.invoke(query)
-            except:
-                return []
-
-        with ThreadPoolExecutor(max_workers=2) as executor:
-            vector_future = executor.submit(vector_search)
-            bm25_future = executor.submit(bm25_search)
-            vector_results = vector_future.result()
-            bm25_results = bm25_future.result()
-
-        # Combine and deduplicate
-        combined = vector_results + bm25_results
-        unique_docs = []
-        seen = set()
-
-        for doc in combined:
-            content_hash = hash(doc.page_content)
-            if content_hash not in seen:
-                seen.add(content_hash)
-                unique_docs.append(doc)
-
-        return unique_docs[:3]  # Return top 3 unique docs
-
-    def query(self, query: str) -> Dict[str, any]:
-        """Process a query with RAG"""
-        start_time = datetime.now()
-
-        # Retrieve relevant documents
-        docs = self.retrieve_documents(query)
-
-        if not docs:
-            return {
-                "answer": "No relevant information found.",
-                "sources": [],
-                "model": "none",
-                "time_ms": 0
-            }
-
-        # Prepare context
-        context = "\n\n".join([doc.page_content for doc in docs])
-
-        # Generate answer if LLM is available
-        if self.llm_provider.is_available:
-            try:
-                prompt = f"""Use the following context to answer the question:
-
-Context:
-{context}
-
-Question: {query}
-
-Answer:"""
-
-                answer = self.llm_provider.invoke(prompt)
-
-                return {
-                    "answer": answer,
-                    "sources": [doc.metadata.get("source", "") for doc in docs],
-                    "model": self.llm_provider.current_model,
-                    "time_ms": (datetime.now() - start_time).total_seconds() * 1000
-                }
-            except Exception as e:
-                print(f"[!] LLM error: {e}")
-                # Fall through to retrieval mode
-
-        # Fallback: return best matching document
-        best_doc = docs[0].page_content[:500] + "..." if len(docs[0].page_content) > 500 else docs[0].page_content
-        return {
-            "answer": f"Relevant information:\n\n{best_doc}",
-            "sources": [doc.metadata.get("source", "") for doc in docs],
-            "model": "retrieval-only",
-            "time_ms": (datetime.now() - start_time).total_seconds() * 1000
-        }
-
-    def add_document(self, filename: str, content: str) -> bool:
-        """Add a document to the knowledge base"""
-        try:
-            self.knowledge_texts.append({
-                "filename": filename,
-                "content": content
-            })
-
-            # Save to file
-            with open(os.path.join(self.knowledge_dir, filename), "w") as f:
-                f.write(content)
-
-            # Rebuild retrievers
-            self.build_retrievers()
-            return True
-
-        except Exception as e:
-            print(f"[!] Error adding document: {e}")
-            return False
-
-    def cleanup(self):
-        """Clean up temporary files"""
-        try:
-            shutil.rmtree(self.temp_dir)
-        except:
-            pass
-
-    def __del__(self):
-        self.cleanup()
+        else:
+            self.qa_chain = None
+
+    def get_knowledge_summary(self) -> str:
+        count = len(self.texts) if self.texts else 0
+        return f"{count} document chunks loaded."
+
+    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.", []