Upload folder using huggingface_hub

app.py

@@ -5,8 +5,9 @@ import os
 import requests
 from pypdf import PdfReader
 import gradio as gr
-import chromadb
 import numpy as np
+import pickle
+import os
 
 load_dotenv(override=True)
 
@@ -105,21 +106,34 @@ class Me:
         self.openai = OpenAI()
         self.name = "Alexandre Saadoun"
         
-        # Initialize Chroma connection
-        self.chroma_client = chromadb.PersistentClient(path="./chroma_db")
+        # Initialize simple vector store
+        self.vector_store_path = "./vector_store.pkl"
+        self.knowledge_base = {"documents": [], "embeddings": [], "metadata": []}
         
         # Initialize RAG system - this will auto-load all files in me/
-        self._setup_chroma_collection()
+        self._setup_vector_store()
         self._populate_initial_data()
 
-    def _setup_chroma_collection(self):
-        """Setup Chroma collection for RAG"""
+    def _setup_vector_store(self):
+        """Setup simple vector store for RAG"""
+        try:
+            if os.path.exists(self.vector_store_path):
+                with open(self.vector_store_path, 'rb') as f:
+                    self.knowledge_base = pickle.load(f)
+                print("✅ Loaded existing knowledge base")
+            else:
+                print("✅ Created new knowledge base")
+        except Exception as e:
+            print(f"Error loading knowledge base: {e}")
+            self.knowledge_base = {"documents": [], "embeddings": [], "metadata": []}
+
+    def _save_vector_store(self):
+        """Save vector store to disk"""
         try:
-            self.collection = self.chroma_client.get_collection(name="knowledge_base")
-            print("✅ Loaded existing knowledge base")
-        except:
-            self.collection = self.chroma_client.create_collection(name="knowledge_base")
-            print("✅ Created new knowledge base")
+            with open(self.vector_store_path, 'wb') as f:
+                pickle.dump(self.knowledge_base, f)
+        except Exception as e:
+            print(f"Error saving knowledge base: {e}")
 
     def _get_embedding(self, text):
         """Get embedding for text using OpenAI"""
@@ -130,9 +144,9 @@ class Me:
         return response.data[0].embedding
 
     def _populate_initial_data(self):
-        """Store initial knowledge in Chroma"""
+        """Store initial knowledge in vector store"""
         # Check if data already exists
-        count = self.collection.count()
+        count = len(self.knowledge_base["documents"])
         
         if count == 0:  # Only populate if empty
             print("Auto-loading all files from me/ directory...")
@@ -192,14 +206,20 @@ class Me:
         
         # Clear existing me/ content
         try:
-            # Get all documents from me/ files
-            results = self.collection.get(include=["metadatas"])
-            me_ids = [results["ids"][i] for i, metadata in enumerate(results["metadatas"]) 
-                     if metadata.get("source", "").startswith("me_")]
+            indices_to_remove = []
+            for i, metadata in enumerate(self.knowledge_base["metadata"]):
+                if metadata.get("source", "").startswith("me_"):
+                    indices_to_remove.append(i)
             
-            if me_ids:
-                self.collection.delete(ids=me_ids)
-                print(f"Cleared {len(me_ids)} existing files from me/")
+            # Remove in reverse order to maintain indices
+            for i in reversed(indices_to_remove):
+                del self.knowledge_base["documents"][i]
+                del self.knowledge_base["embeddings"][i]
+                del self.knowledge_base["metadata"][i]
+                
+            if indices_to_remove:
+                print(f"Cleared {len(indices_to_remove)} existing files from me/")
+                self._save_vector_store()
         except Exception as e:
             print(f"Error clearing existing data: {e}")
         
@@ -210,20 +230,40 @@ class Me:
     def _search_knowledge(self, query, limit=3):
         """Search for relevant knowledge using vector similarity"""
         try:
-            results = self.collection.query(
-                query_texts=[query],
-                n_results=limit,
-                include=["documents", "metadatas", "distances"]
-            )
+            if not self.knowledge_base["documents"]:
+                return []
+                
+            # Get query embedding
+            query_embedding = self._get_embedding(query)
+            query_vector = np.array(query_embedding)
+            
+            # Calculate cosine similarities
+            similarities = []
+            for i, doc_embedding in enumerate(self.knowledge_base["embeddings"]):
+                doc_vector = np.array(doc_embedding)
+                
+                # Cosine similarity
+                dot_product = np.dot(query_vector, doc_vector)
+                norm_query = np.linalg.norm(query_vector)
+                norm_doc = np.linalg.norm(doc_vector)
+                
+                if norm_query > 0 and norm_doc > 0:
+                    similarity = dot_product / (norm_query * norm_doc)
+                else:
+                    similarity = 0.0
+                    
+                similarities.append((similarity, i))
+            
+            # Sort by similarity and get top results
+            similarities.sort(reverse=True)
             
             search_results = []
-            if results["documents"] and results["documents"][0]:
-                for i, doc in enumerate(results["documents"][0]):
-                    search_results.append({
-                        "content": doc,
-                        "type": results["metadatas"][0][i].get("type", "unknown") if results["metadatas"] else "unknown",
-                        "score": 1 - results["distances"][0][i] if results["distances"] else 1.0
-                    })
+            for similarity, idx in similarities[:limit]:
+                search_results.append({
+                    "content": self.knowledge_base["documents"][idx],
+                    "type": self.knowledge_base["metadata"][idx].get("type", "unknown"),
+                    "score": similarity
+                })
             
             return search_results
         except Exception as e:
@@ -231,18 +271,19 @@ class Me:
             return []
 
     def _store_new_knowledge(self, information, context=""):
-        """Store new information in Chroma"""
+        """Store new information in vector store"""
         try:
-            doc_id = f"conv_{len(self.collection.get()['ids'])}"
-            self.collection.add(
-                documents=[information],
-                metadatas=[{
-                    "type": "conversation",
-                    "context": context,
-                    "timestamp": str(np.datetime64('now'))
-                }],
-                ids=[doc_id]
-            )
+            embedding = self._get_embedding(information)
+            
+            self.knowledge_base["documents"].append(information)
+            self.knowledge_base["embeddings"].append(embedding)
+            self.knowledge_base["metadata"].append({
+                "type": "conversation",
+                "context": context,
+                "timestamp": str(np.datetime64('now'))
+            })
+            
+            self._save_vector_store()
         except Exception as e:
             print(f"Error storing knowledge: {e}")
 
@@ -268,25 +309,19 @@ class Me:
         
         # Store each chunk
         try:
-            documents = []
-            metadatas = []
-            ids = []
-            
             for i, chunk in enumerate(chunks):
-                documents.append(chunk)
-                metadatas.append({
+                embedding = self._get_embedding(chunk)
+                
+                self.knowledge_base["documents"].append(chunk)
+                self.knowledge_base["embeddings"].append(embedding)
+                self.knowledge_base["metadata"].append({
                     "type": "text_content",
                     "source": source_name,
                     "chunk_index": i,
                     "timestamp": str(np.datetime64('now'))
                 })
-                ids.append(f"{source_name}_chunk_{i}")
             
-            self.collection.add(
-                documents=documents,
-                metadatas=metadatas,
-                ids=ids
-            )
+            self._save_vector_store()
         except Exception as e:
             print(f"Error storing chunks: {e}")
         
@@ -340,22 +375,31 @@ class Me:
         """
         try:
             if knowledge_type:
-                # Get documents of specific type
-                results = self.collection.get(include=["metadatas"])
-                type_ids = [results["ids"][i] for i, metadata in enumerate(results["metadatas"]) 
-                           if metadata.get("type") == knowledge_type]
+                # Remove documents of specific type
+                indices_to_remove = []
+                for i, metadata in enumerate(self.knowledge_base["metadata"]):
+                    if metadata.get("type") == knowledge_type:
+                        indices_to_remove.append(i)
+                
+                # Remove in reverse order to maintain indices
+                for i in reversed(indices_to_remove):
+                    del self.knowledge_base["documents"][i]
+                    del self.knowledge_base["embeddings"][i]
+                    del self.knowledge_base["metadata"][i]
                 
-                if type_ids:
-                    self.collection.delete(ids=type_ids)
-                    print(f"Deleted {len(type_ids)} {knowledge_type} documents")
+                if indices_to_remove:
+                    print(f"Deleted {len(indices_to_remove)} {knowledge_type} documents")
+                    self._save_vector_store()
                 else:
                     print(f"No {knowledge_type} documents found")
             else:
-                # Clear entire collection
-                all_ids = self.collection.get()["ids"]
-                if all_ids:
-                    self.collection.delete(ids=all_ids)
-                    print(f"Deleted {len(all_ids)} documents")
+                # Clear entire knowledge base
+                count = len(self.knowledge_base["documents"])
+                self.knowledge_base = {"documents": [], "embeddings": [], "metadata": []}
+                
+                if count > 0:
+                    print(f"Deleted {count} documents")
+                    self._save_vector_store()
                 else:
                     print("No documents to delete")
                     
@@ -365,12 +409,10 @@ class Me:
     def get_knowledge_stats(self):
         """Get statistics about the knowledge base"""
         try:
-            results = self.collection.get(include=["metadatas"])
-            
             stats = {}
-            total = len(results["ids"])
+            total = len(self.knowledge_base["documents"])
             
-            for metadata in results["metadatas"]:
+            for metadata in self.knowledge_base["metadata"]:
                 doc_type = metadata.get("type", "unknown")
                 stats[doc_type] = stats.get(doc_type, 0) + 1