Update app.py

app.py

@@ -1,12 +1,12 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-import os, logging, traceback
-from pathlib import Path
 from typing import Optional, Dict
+import os, shutil, logging, traceback
+from pathlib import Path
 
-# Writable base (HF Spaces persistent volume)
-DATA_DIR = os.getenv("DATA_DIR", "/data")
+# ===== Persistent storage on HF =====
+DATA_DIR   = os.getenv("DATA_DIR", "/data")
 INDEX_ROOT = os.getenv("INDEX_DIR", os.path.join(DATA_DIR, "index"))
 Path(INDEX_ROOT).mkdir(parents=True, exist_ok=True)
 
@@ -15,15 +15,14 @@ from agent.tools import FetchTools
 from ingest.sec import fetch_recent_filings_by_cik
 
 log = logging.getLogger("uvicorn.error")
-
 app = FastAPI(title="DeepDive IR Agent")
 
-# CORS: allow localhost + all vercel previews + your prod app
+# ===== CORS: localhost + vercel previews + your prod app =====
 app.add_middleware(
     CORSMiddleware,
     allow_origins=[
         "http://localhost:3000",
-        "https://deepdive-ir-agent.vercel.app",
+        "https://deepdive-ir-agent.vercel.app",  # <-- change if your prod URL is different
     ],
     allow_origin_regex=r"https://.*\.vercel\.app$",
     allow_credentials=True,
@@ -31,21 +30,21 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ---- Models ----
+# ===== Models =====
 class IngestRequest(BaseModel):
     cik: str
     ir_url: Optional[str] = None
 
 class AskRequest(BaseModel):
     question: str
-    cik: Optional[str] = None  # optional: if omitted we use the last ingested CIK
+    cik: Optional[str] = None  # optional; if omitted we use the last ingested CIK
 
-# ---- State: cache graphs per CIK ----
-graphs: Dict[str, AgentGraph] = {}
+# ===== State =====
+graphs: Dict[str, AgentGraph] = {}   # CIK -> AgentGraph
 last_cik: Optional[str] = None
 tools = FetchTools()
 
-# ---- Helpers ----
+# ===== Helpers =====
 def norm_cik(raw: str) -> str:
     s = raw.strip()
     if not s.isdigit():
@@ -54,11 +53,12 @@ def norm_cik(raw: str) -> str:
         raise HTTPException(400, "CIK too long; use 10 digits.")
     return s.zfill(10)
 
-def company_index_dir(cik: str) -> str:
+def idx_dir_for(cik: str) -> str:
     d = os.path.join(INDEX_ROOT, cik)
     Path(d).mkdir(parents=True, exist_ok=True)
     return d
 
+# ===== Routes =====
 @app.get("/")
 def root():
     return {"ok": True, "msg": "DeepDive IR Agent API"}
@@ -69,10 +69,14 @@ def healthz():
 
 @app.post("/ingest")
 async def ingest(req: IngestRequest):
+    """
+    Build a fresh index for this CIK under /data/index/<CIK>.
+    """
     global last_cik
     try:
         cik = norm_cik(req.cik)
 
+        # Fetch recent filings + optional IR site
         filings = await fetch_recent_filings_by_cik(cik)
         docs = []
         for form, url, title in filings:
@@ -95,14 +99,17 @@ async def ingest(req: IngestRequest):
         if not docs:
             raise HTTPException(400, "No documents fetched.")
 
-        idx_dir = company_index_dir(cik)
+        # Fresh per-CIK folder
+        idx_dir = idx_dir_for(cik)
+        shutil.rmtree(idx_dir, ignore_errors=True)
+        Path(idx_dir).mkdir(parents=True, exist_ok=True)
 
-        # Some libraries write relative paths like "index/vecs.npy".
-        # Make those land under /data/index/<CIK>/ by chdir-ing temporarily.
+        # Some libs may write relative paths like "index/vecs.npy".
+        # Build from the CIK directory so relative paths resolve to /data/index/<CIK>/...
         prev = os.getcwd()
         os.chdir(idx_dir)
         try:
-            g = AgentGraph(index_dir=idx_dir)   # absolute, per-CIK dir
+            g = AgentGraph(index_dir=idx_dir)  # absolute per-CIK dir
             g.build_index(docs)
         finally:
             os.chdir(prev)