Update agent.py

agent.py

@@ -365,11 +365,13 @@ class MyVectorStore:
 # Prepare Documents
 # -----------------------------
 # Define the URL where the JSON file is hosted
-
 from typing import TypedDict, Annotated, List
 import gradio as gr
 from langchain.schema import Document
 import json
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+import faiss
 
 # 1. Type-Checked State for Gradio
 class ChatState(TypedDict):
@@ -416,29 +418,45 @@ def create_documents(data_source: str, data: List[dict]) -> List[Document]:
     
     return docs
 
-texts = [doc.page_content for doc in docs]
-
-# Initialize the embedding model
-embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-
-docs = [Document(page_content=texts, metadata={"task_id": item["task_id"]}) for item in data]
+# 4. Vector Store Integration
+class MyVectorStore:
+    def __init__(self, index: faiss.Index):
+        self.index = index
 
-# Create the FAISS index
-#vector_store = FAISS.from_documents(docs, embedding_model)
-vector_store = FAISS.from_texts(texts = text_chunks, embedding = embeddings)
+    def save_local(self, path: str):
+        faiss.write_index(self.index, path)
 
-# Save the FAISS index
-vector_store = MyVectorStore(index)
+    @classmethod
+    def load_local(cls, path: str):
+        index = faiss.read_index(path)
+        return cls(index)
 
-# Assuming you already have a FAISS index
-index = faiss.IndexFlatL2(768)  # Or whatever dimensionality your embeddings have
-vector_store = MyVectorStore(index)
+# Usage Example
+if __name__ == "__main__":
+    # Process JSON data
+    with open("questions.json", "r", encoding="utf-8") as f:
+        json_data = json.load(f)
+    
+    # Create documents from JSON
+    docs = create_documents("json", json_data)
+    texts = [doc.page_content for doc in docs]
 
-# Save
-vector_store.save_local("/home/wendy/Downloads/faiss_index.index")
+    # Initialize embedding model
+    embedding_model = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2"
+    )
 
-# Load
-loaded_store = MyVectorStore.load_local("/home/wendy/Downloads/faiss_index.index")
+    # Create FAISS index
+    vector_store = FAISS.from_documents(
+        documents=docs,
+        embedding=embedding_model
+    )
+        
+    # Save
+    vector_store.save_local("/home/wendy/Downloads/faiss_index.index")
+    
+    # Load
+    loaded_store = MyVectorStore.load_local("/home/wendy/Downloads/faiss_index.index")