Update app.py

app.py

@@ -3,13 +3,12 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import os, logging, traceback
 from pathlib import Path
+from typing import Optional, Dict
 
-# ⬅️ we only need DATA_DIR here; we'll derive the absolute index path from it
-from settings import DATA_DIR
-
-# Make sure the HF persistent folder exists: /data/index
-INDEX_ABS = os.path.join(DATA_DIR, "index")
-Path(INDEX_ABS).mkdir(parents=True, exist_ok=True)
+# Writable base (HF Spaces persistent volume)
+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)
 
 from agent.graph import AgentGraph
 from agent.tools import FetchTools
@@ -19,27 +18,46 @@ log = logging.getLogger("uvicorn.error")
 
 app = FastAPI(title="DeepDive IR Agent")
 
-# CORS — tighten later once your Vercel URL is final
+# CORS: allow localhost + all vercel previews + your prod app
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],       # for debugging; later: ["https://<your-vercel>.vercel.app", "http://localhost:3000"]
+    allow_origins=[
+        "http://localhost:3000",
+        "https://deepdive-ir-agent.vercel.app",
+    ],
+    allow_origin_regex=r"https://.*\.vercel\.app$",
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Ensure index path exists (and is persistent on HF: set INDEX_DIR="/data/index" in settings)
-Path(INDEX_DIR).mkdir(parents=True, exist_ok=True)
-
-graph = AgentGraph(index_dir=INDEX_DIR)
-tools = FetchTools()
-
+# ---- Models ----
 class IngestRequest(BaseModel):
     cik: str
-    ir_url: str | None = None
+    ir_url: Optional[str] = None
 
 class AskRequest(BaseModel):
     question: str
+    cik: Optional[str] = None  # optional: if omitted we use the last ingested CIK
+
+# ---- State: cache graphs per CIK ----
+graphs: Dict[str, AgentGraph] = {}
+last_cik: Optional[str] = None
+tools = FetchTools()
+
+# ---- Helpers ----
+def norm_cik(raw: str) -> str:
+    s = raw.strip()
+    if not s.isdigit():
+        raise HTTPException(400, "CIK must be digits only.")
+    if len(s) > 10:
+        raise HTTPException(400, "CIK too long; use 10 digits.")
+    return s.zfill(10)
+
+def company_index_dir(cik: str) -> str:
+    d = os.path.join(INDEX_ROOT, cik)
+    Path(d).mkdir(parents=True, exist_ok=True)
+    return d
 
 @app.get("/")
 def root():
@@ -51,54 +69,76 @@ def healthz():
 
 @app.post("/ingest")
 async def ingest(req: IngestRequest):
+    global last_cik
     try:
-        # SEC wants a *10-digit, zero-padded* CIK string. Keep as string to preserve leading zeros.
-        cik = req.cik.strip()
-        if not cik.isdigit():
-            raise HTTPException(status_code=400, detail="CIK must be digits only.")
-        if len(cik) > 10:
-            raise HTTPException(status_code=400, detail="CIK too long; use 10 digits.")
-        cik = cik.zfill(10)
-
-        filings = await fetch_recent_filings_by_cik(cik)  # must send a real User-Agent inside!
+        cik = norm_cik(req.cik)
+
+        filings = await fetch_recent_filings_by_cik(cik)
         docs = []
         for form, url, title in filings:
-            text = await tools.get_text_from_url(url)     # must send a real User-Agent inside!
-            if not text:
-                log.warning(f"Empty text for {url}")
-                continue
-            docs.append({"title": title, "url": url, "text": text})
+            try:
+                text = await tools.get_text_from_url(url)
+            except Exception as e:
+                log.warning(f"Fetch failed for {url}: {e}")
+                text = ""
+            if text:
+                docs.append({"title": title, "url": url, "text": text})
 
         if req.ir_url:
-            ir_text = await tools.get_text_from_url(req.ir_url)
-            if ir_text:
-                docs.append({"title": "IR site", "url": req.ir_url, "text": ir_text})
+            try:
+                ir_text = await tools.get_text_from_url(req.ir_url)
+                if ir_text:
+                    docs.append({"title": "IR site", "url": req.ir_url, "text": ir_text})
+            except Exception as e:
+                log.warning(f"IR fetch failed for {req.ir_url}: {e}")
 
         if not docs:
-            raise HTTPException(status_code=400, detail="No documents fetched.")
+            raise HTTPException(400, "No documents fetched.")
+
+        idx_dir = company_index_dir(cik)
 
-        graph.build_index(docs)
-        return {"ok": True, "num_docs": len(docs)}
+        # Some libraries write relative paths like "index/vecs.npy".
+        # Make those land under /data/index/<CIK>/ by chdir-ing temporarily.
+        prev = os.getcwd()
+        os.chdir(idx_dir)
+        try:
+            g = AgentGraph(index_dir=idx_dir)   # absolute, per-CIK dir
+            g.build_index(docs)
+        finally:
+            os.chdir(prev)
+
+        graphs[cik] = g
+        last_cik = cik
+        return {"ok": True, "cik": cik, "num_docs": len(docs)}
 
     except HTTPException:
         raise
     except Exception as e:
-        # Log full traceback for debugging in HF "Logs"
         log.error("Ingest failed: %s\n%s", e, traceback.format_exc())
         raise HTTPException(status_code=502, detail=f"Ingest failed: {type(e).__name__}: {e}")
 
 @app.post("/ask")
 def ask(req: AskRequest):
     try:
-        return graph.answer(req.question)
+        cik = norm_cik(req.cik) if req.cik else last_cik
+        if not cik or cik not in graphs:
+            raise HTTPException(400, "No index available. Call /ingest with a CIK first.")
+        return graphs[cik].answer(req.question)
+    except HTTPException:
+        raise
     except Exception as e:
         log.error("Ask failed: %s\n%s", e, traceback.format_exc())
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(500, detail=str(e))
 
 @app.get("/brief")
-def brief():
+def brief(cik: Optional[str] = None):
     try:
-        return graph.brief()
+        c = norm_cik(cik) if cik else last_cik
+        if not c or c not in graphs:
+            raise HTTPException(400, "No index available. Call /ingest with a CIK first.")
+        return graphs[c].brief()
+    except HTTPException:
+        raise
     except Exception as e:
         log.error("Brief failed: %s\n%s", e, traceback.format_exc())
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(500, detail=str(e))