Update app.py

app.py

@@ -5,6 +5,8 @@ import requests
 import pandas as pd
 from langchain_core.messages import HumanMessage
 from agent import build_graph
+from langchain_core.documents import Document
+from langchain_community.vectorstores import Chroma
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -183,6 +185,34 @@ if __name__ == "__main__":
     space_host_startup = os.getenv("SPACE_HOST")
     space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup
 
+    # Load metadata.jsonl
+    import json
+    # Load the metadata.jsonl file
+    with open('metadata.jsonl', 'r') as jsonl_file:
+        json_list = list(jsonl_file)
+    
+    json_QA = []
+    for json_str in json_list:
+        json_data = json.loads(json_str)
+        json_QA.append(json_data)
+
+    # 1. Prepare your documents (same format as before)
+    documents = [
+        Document(
+            page_content=f"Q: {q['Question']}\nA: {q['Final answer']}",
+            metadata={"source": q['task_id']}
+        ) for q in json_QA
+    ]
+    
+    # 2. Initialize ChromaDB
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") #  dim=768
+    vector_store = Chroma.from_documents(
+        collection_name = 'documents',
+        documents=documents,
+        embedding=embeddings,
+        persist_directory="./vector_db"  # Omit for in-memory only
+    )
+
     if space_host_startup:
         print(f"✅ SPACE_HOST found: {space_host_startup}")
         print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")