Update agent.py

agent.py

@@ -401,6 +401,16 @@ def create_documents(data_source: str, data: List[dict]) -> List[Document]:
     return docs
 
 # 4. Vector Store Integration
+import os
+import json
+import faiss
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.tools import create_retriever_tool
+from langchain.llms import HuggingFaceEndpoint
+from your_module import create_documents  # Replace with your actual import
+
+# Custom FAISS wrapper (optional, if you still want it)
 class MyVector_Store:
     def __init__(self, index: faiss.Index):
         self.index = index
@@ -413,40 +423,41 @@ class MyVector_Store:
         index = faiss.read_index(path)
         return cls(index)
 
+# -----------------------------
+# Process JSON data and create documents
+# -----------------------------
+with open("questions.json", "r", encoding="utf-8") as f:
+    json_data = json.load(f)
 
-    # 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]
-
-    # Initialize embedding model
-    embedding_model = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2"
-    )
+docs = create_documents("json", json_data)
+texts = [doc.page_content for doc in docs]
 
-    # 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 = vector_store.load_local("/home/wendy/Downloads/faiss_index.index")
+# -----------------------------
+# Initialize embedding model
+# -----------------------------
+embedding_model = HuggingFaceEmbeddings(
+    model_name="sentence-transformers/all-MiniLM-L6-v2"
+)
 
+# -----------------------------
+# Create FAISS index and save it
+# -----------------------------
+vector_store = FAISS.from_documents(
+    documents=docs,
+    embedding=embedding_model
+)
 
+vector_store.save_local("/home/wendy/Downloads/faiss_index")
 
+# -----------------------------
+# Load FAISS index properly
+# -----------------------------
+loaded_store = FAISS.load_local("/home/wendy/Downloads/faiss_index", embedding_model)
 
 # -----------------------------
 # Create LangChain Retriever Tool
 # -----------------------------
-
-retriever = FAISS.loaded_store("/home/wendy/Downloads/faiss_index.index", embedding_model).as_retriever()
+retriever = loaded_store.as_retriever()
 
 question_retriever_tool = create_retriever_tool(
     retriever=retriever,
@@ -454,17 +465,9 @@ question_retriever_tool = create_retriever_tool(
     description="A tool to retrieve documents related to a user's question."
 )
 
-
-
-# Define the LLM before using it
-#llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo") # or "gpt-3.5-turbo" "gpt-4"
-#llm = ChatMistralAI(model="mistral-7b-instruct-v0.1")
-
-
-# Get the Hugging Face API token from the environment variable
-#hf_token = os.getenv("HF_TOKEN")
-
-
+# -----------------------------
+# Load HuggingFace LLM
+# -----------------------------
 llm = HuggingFaceEndpoint(
     repo_id="HuggingFaceH4/zephyr-7b-beta",
     task="text-generation",
@@ -473,7 +476,7 @@ llm = HuggingFaceEndpoint(
     max_new_tokens=512
 )
 
-
+        
 
 # -------------------------------
 # Step 8: Use the Planner, Classifier, and Decision Logic