Update main.py

main.py

@@ -25,36 +25,41 @@ class QueryRequest(BaseModel):
     question: str
 
 
-def _unpack_faiss(src: str, dest_dir: str) -> str:
+def _unpack_faiss(src_path: str, extract_to: str) -> str:
     """
-    If src is a .zip, unzip it into dest_dir and return
-    the path to the extracted FAISS folder. Otherwise
-    assume src is already a folder and return it.
+    If src_path is a .zip, unzip to extract_to and return the directory
+    containing the .faiss file.  If it's already a folder, just return it.
     """
-    if zipfile.is_zipfile(src):
-        with zipfile.ZipFile(src, "r") as zf:
-            zf.extractall(dest_dir)
-        # if there’s exactly one subfolder, use it
-        items = os.listdir(dest_dir)
-        if len(items) == 1 and os.path.isdir(os.path.join(dest_dir, items[0])):
-            return os.path.join(dest_dir, items[0])
-        return dest_dir
-    else:
-        # src is already a directory
-        return src
+    # 1) ZIP case
+    if src_path.lower().endswith(".zip"):
+        if not os.path.isfile(src_path):
+            raise FileNotFoundError(f"Could not find zip file: {src_path}")
+        with zipfile.ZipFile(src_path, "r") as zf:
+            zf.extractall(extract_to)
+
+        # walk until we find any .faiss file
+        for root, _, files in os.walk(extract_to):
+            if any(fn.endswith(".faiss") for fn in files):
+                return root
+        raise RuntimeError(f"No .faiss index found inside {src_path}")
+
+    # 2) directory case
+    if os.path.isdir(src_path):
+        return src_path
+
+    raise RuntimeError(f"Path is neither a .zip nor a directory: {src_path}")
 
 
 @app.on_event("startup")
 def load_components():
     global llm, embeddings, vectorstore, retriever, chain
 
-    # --- 1) Initialize LLM & Embeddings ---
-    api_key = os.getenv("api_key")
+    # --- 1) Init LLM & Embeddings ---
     llm = ChatGroq(
         model="meta-llama/llama-4-scout-17b-16e-instruct",
         temperature=0,
         max_tokens=1024,
-        api_key=api_key,
+        api_key=os.getenv("api_key"),
     )
 
     embeddings = HuggingFaceEmbeddings(
@@ -64,25 +69,23 @@ def load_components():
     )
 
     # --- 2) Load & merge two FAISS indexes ---
-    # Paths to your two vectorstores (could be .zip or folders)
     src1 = "faiss_index.zip"
     src2 = "faiss_index_extra.zip"
 
-    # Temporary dirs for extraction
-    tmp1 = tempfile.mkdtemp()
-    tmp2 = tempfile.mkdtemp()
+    # Use TemporaryDirectory objects so they stick around until program exit
+    tmp1 = tempfile.TemporaryDirectory()
+    tmp2 = tempfile.TemporaryDirectory()
 
-    # Unpack and load each
-    path1 = _unpack_faiss(src1, tmp1)
-    vs1 = FAISS.load_local(path1, embeddings, allow_dangerous_deserialization=True)
+    # Unpack & locate
+    dir1 = _unpack_faiss(src1, tmp1.name)
+    dir2 = _unpack_faiss(src2, tmp2.name)
 
-    path2 = _unpack_faiss(src2, tmp2)
-    vs2 = FAISS.load_local(path2, embeddings, allow_dangerous_deserialization=True)
+    # Load them
+    vs1 = FAISS.load_local(dir1, embeddings, allow_dangerous_deserialization=True)
+    vs2 = FAISS.load_local(dir2, embeddings, allow_dangerous_deserialization=True)
 
     # Merge vs2 into vs1
     vs1.merge_from(vs2)
-
-    # Assign the merged store to our global
     vectorstore = vs1
 
     # --- 3) Build retriever & QA chain ---
@@ -114,7 +117,7 @@ Your response:
         chain_type_kwargs={"prompt": prompt},
     )
 
-    print("✅ Loaded and merged both FAISS indexes, QA chain is ready.")
+    print("✅ Loaded & merged both FAISS indexes, QA chain ready.")
 
 
 @app.get("/")