Update app/main.py

app/main.py

@@ -10,7 +10,10 @@ Run:  uvicorn app.main:app --host 0.0.0.0 --port 8000
 
 from __future__ import annotations
 
+import io
 import json
+import os
+import zipfile
 from datetime import date
 from typing import List
 
@@ -19,6 +22,7 @@ from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
+from app.config import settings
 from app.graph.builder import get_graph
 
 app = FastAPI(title="Sattva AI API")
@@ -178,5 +182,39 @@ async def chat_stream_endpoint(request: ChatRequest):
     return StreamingResponse(event_generator(), media_type="text/event-stream")
 
 
+@app.get("/download_store")
+def download_store():
+    """Zip and download the built Chroma vector store (chroma_db_unified/).
+
+    Use this once the build is 100% done (after the `Created unified vector store
+    with N chunks` log) to pull the store off the Space and reuse it locally /
+    commit it — so you never have to rebuild from scratch.
+
+        curl -o chroma_db_unified.zip https://<space>.hf.space/download_store
+        unzip chroma_db_unified.zip
+
+    Returns 404 if the store hasn't been built yet (empty/missing directory).
+    """
+    base = os.path.join(os.getcwd(), settings.RAG_PERSIST_DIR)  # chroma_db_unified
+    if not os.path.isdir(base) or not os.listdir(base):
+        raise HTTPException(status_code=404, detail="Vector store not built yet.")
+
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
+        for root, _, files in os.walk(base):
+            for fname in files:
+                fpath = os.path.join(root, fname)
+                # arcname keeps the top-level `chroma_db_unified/` folder in the zip
+                arcname = os.path.relpath(fpath, os.path.dirname(base))
+                zf.write(fpath, arcname)
+    buf.seek(0)
+
+    return StreamingResponse(
+        buf,
+        media_type="application/zip",
+        headers={"Content-Disposition": f"attachment; filename={settings.RAG_PERSIST_DIR}.zip"},
+    )
+
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)