Update main.py

main.py

@@ -1,7 +1,6 @@
 import os
 import zipfile
 import tempfile
-
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 
@@ -25,27 +24,37 @@ class QueryRequest(BaseModel):
     question: str
 
 
-def _unpack_faiss(src_path: str, extract_to: str) -> str:
+def _unpack_faiss(src_path: str) -> str:
     """
-    If src_path is a valid .zip archive, unzip it into extract_to and
-    return the subdirectory that contains the .faiss index.
-    If src_path is already a directory, return it directly.
+    If src_path is a ZIP, unzip it into a temp dir and return the folder
+    containing the .faiss files; if it’s already a folder, return it.
     """
-    # 1) True ZIP file?
     if zipfile.is_zipfile(src_path):
+        tmp = tempfile.TemporaryDirectory()
         with zipfile.ZipFile(src_path, "r") as zf:
-            zf.extractall(extract_to)
-        # scan until we find any .faiss file
-        for root, _, files in os.walk(extract_to):
+            zf.extractall(tmp.name)
+        for root, _, files in os.walk(tmp.name):
             if any(f.endswith(".faiss") for f in files):
                 return root
         raise RuntimeError(f"No .faiss index found inside ZIP: {src_path}")
-
-    # 2) Already a folder?
-    if os.path.isdir(src_path):
+    elif os.path.isdir(src_path):
         return src_path
+    else:
+        raise RuntimeError(f"Path is neither a valid ZIP nor a directory: {src_path}")
+
+
+def load_and_merge_faiss(path1: str, path2: str, embeddings: HuggingFaceEmbeddings) -> FAISS:
+    """
+    Load two FAISS indexes (either zip files or folders), merge them,
+    and return the combined FAISS vectorstore.
+    """
+    dir1 = _unpack_faiss(path1)
+    dir2 = _unpack_faiss(path2)
 
-    raise RuntimeError(f"Path is neither a valid ZIP nor a directory: {src_path}")
+    vs1 = FAISS.load_local(dir1, embeddings, allow_dangerous_deserialization=True)
+    vs2 = FAISS.load_local(dir2, embeddings, allow_dangerous_deserialization=True)
+    vs1.merge_from(vs2)
+    return vs1
 
 
 @app.on_event("startup")
@@ -57,7 +66,7 @@ def load_components():
         model="meta-llama/llama-4-scout-17b-16e-instruct",
         temperature=0,
         max_tokens=1024,
-        api_key=os.getenv("api_key"),
+        api_key=os.getenv("API_KEY"),
     )
     embeddings = HuggingFaceEmbeddings(
         model_name="intfloat/multilingual-e5-large",
@@ -66,29 +75,17 @@ def load_components():
     )
 
     # --- 2) Load & merge two FAISS indexes ---
-    # (these can be either real .zip files or existing folders)
-    src1 = "faiss_index.zip"          # or "faiss_index" if it's already a folder
-    src2 = "faiss_index_extra.zip"    # or "faiss_index_extra"
-
-    tmp1 = tempfile.TemporaryDirectory()
-    tmp2 = tempfile.TemporaryDirectory()
-
-    dir1 = _unpack_faiss(src1, tmp1.name)
-    dir2 = _unpack_faiss(src2, tmp2.name)
-
-    vs1 = FAISS.load_local(dir1, embeddings, allow_dangerous_deserialization=True)
-    vs2 = FAISS.load_local(dir2, embeddings, allow_dangerous_deserialization=True)
-
-    vs1.merge_from(vs2)
-    vectorstore = vs1
+    src1 = os.getenv("FAISS_INDEX_PATH_1", "faiss_index.zip")
+    src2 = os.getenv("FAISS_INDEX_PATH_2", "faiss_index_extra.zip")
+    vectorstore = load_and_merge_faiss(src1, src2, embeddings)
 
     # --- 3) Build retriever & QA chain ---
     retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
     prompt = PromptTemplate(
         template="""
-You are an expert assistant on Islamic knowledge.  
-Use **only** the information in the “Retrieved context” to answer the user’s question.  
-Do **not** add any outside information, personal opinions, or conjecture—if the answer is not contained in the context, reply with “لا أعلم”.  
+You are an expert assistant on Islamic knowledge.
+Use **only** the information in the “Retrieved context” to answer the user’s question.
+Do **not** add any outside information, personal opinions, or conjecture—if the answer is not contained in the context, reply with “لا أعلم”.
 Be concise, accurate, and directly address the user’s question.
 
 Retrieved context:
@@ -114,7 +111,7 @@ Your response:
 
 @app.get("/")
 def root():
-    return {"message": "Arabic Hadith Finder API is up..."}
+    return {"message": "Arabic Hadith Finder API is up and running!"}
 
 
 @app.post("/query")