Update vector.py

vector.py

@@ -162,7 +162,7 @@ class VectorDatabase:
         
         return chunks
     
-    def store_session_document(self, text: str, filename: str, user_id: str):
+    def store_session_document(self, text: str, filename: str, user_id: str, chat_id: str):
         """Store extracted file content with user session isolation"""
         if not text or len(text) < 10 or not user_id:
             logger.warning(f"Invalid input for {filename}")
@@ -211,6 +211,7 @@ class VectorDatabase:
                 "subtype": chunk.get("type", "general"),
                 "name": chunk.get("name", "unknown"),
                 "user_id": user_id,
+                "chat_id": chat_id,
                 "timestamp": time.time(),
                 "chunk_index": len(final_texts)
             })
@@ -250,7 +251,7 @@ class VectorDatabase:
         if user_vectors < expected_count:
             logger.warning(f"⚠️ Storage mismatch for user {user_id[:8]}")
     
-    def store_chat_context(self, messages: list, user_id: str):
+    def store_chat_context(self, messages: list, user_id: str, chat_id: str):
         """Store chat history as session memory"""
         if not messages or not user_id:
             return False
@@ -282,6 +283,7 @@ class VectorDatabase:
                 "source": "chat_history",
                 "type": "history",
                 "user_id": user_id,
+                "chat_id": chat_id,
                 "timestamp": time.time(),
                 "chunk_index": i
             })
@@ -303,7 +305,7 @@ class VectorDatabase:
             logger.error(f"Failed to store chat history: {e}")
             return False
     
-    def retrieve_session_context(self, query: str, user_id: str, filter_type: str = None, top_k=100, final_k=5):
+    def retrieve_session_context(self, query: str, user_id: str, chat_id: str, filter_type: str = None, top_k=100, final_k=5):
         """Retrieve context only from the user's session"""
         if self.index.ntotal == 0 or not user_id:
             logger.warning(f"Empty index or missing user_id. Index size: {self.index.ntotal}")
@@ -340,7 +342,8 @@ class VectorDatabase:
             # STRICT user filtering
             if item.get("user_id") != user_id:
                 continue
-            
+            if item.get("chat_id") != chat_id:
+                continue
             # Optional type filtering
             if filter_type and item.get("type") != filter_type:
                 continue