Update app.py

app.py

@@ -1,50 +1,39 @@
 import os, io, json, zipfile, hashlib
-from typing import List, Dict, Any, Optional, Union
+from typing import List, Dict, Any, Optional
 import gradio as gr
 from pydantic import BaseModel
-from tenacity import retry, stop_after_attempt, wait_exponential
+from tenacity import retry, stop_after_attempt, wait_exponential, RetryError
 
-# --- Optional .env support ---
+# .env support (optional)
 try:
     from dotenv import load_dotenv
     load_dotenv()
 except Exception:
     pass
 
-# --- Model SDKs ---
-# OpenAI (GPT-5)
+# SDKs
 try:
     from openai import OpenAI
 except Exception:
     OpenAI = None
 
-# Anthropic (Claude Sonnet 4.5 / latest)
 try:
     import anthropic
+    from anthropic import NotFoundError as AnthropicNotFound
 except Exception:
     anthropic = None
+    AnthropicNotFound = Exception  # fallback type
 
-# Firecrawl SDK
 from firecrawl import Firecrawl  # v2.x
 
-# --------------------------
-# Utilities
-# --------------------------
+# -------------------- utils --------------------
 def _to_dict(obj: Any) -> Any:
-    """
-    Recursively convert Firecrawl/Pydantic models (SearchData, ScrapeData, CrawlData, etc.)
-    or other objects into plain Python dict/list primitives.
-    """
-    # Pydantic v2 models
     if isinstance(obj, BaseModel):
         return obj.model_dump()
-    # Mapping-like
     if isinstance(obj, dict):
         return {k: _to_dict(v) for k, v in obj.items()}
-    # List/Tuple
     if isinstance(obj, (list, tuple)):
         return [_to_dict(v) for v in obj]
-    # Objects with __dict__ (fallback)
     if hasattr(obj, "__dict__") and not isinstance(obj, (str, bytes)):
         try:
             return {k: _to_dict(v) for k, v in vars(obj).items()}
@@ -59,120 +48,154 @@ def _pretty_json(data: Any, limit: int = 300_000) -> str:
     except Exception as e:
         return f"<!> Could not serialize to JSON: {e}"
 
-# --------------------------
-# Session keys
-# --------------------------
+def _listify(x) -> List[Any]:
+    if x is None:
+        return []
+    if isinstance(x, list):
+        return x
+    return [x]
+
+# -------------------- keys --------------------
 class Keys(BaseModel):
     openai: Optional[str] = None
     anthropic: Optional[str] = None
     firecrawl: Optional[str] = None
 
-def resolve_keys(session: Keys) -> Keys:
+def resolve_keys(s: Keys) -> Keys:
     return Keys(
-        openai=session.openai or os.getenv("OPENAI_API_KEY"),
-        anthropic=session.anthropic or os.getenv("ANTHROPIC_API_KEY"),
-        firecrawl=session.firecrawl or os.getenv("FIRECRAWL_API_KEY"),
+        openai=s.openai or os.getenv("OPENAI_API_KEY"),
+        anthropic=s.anthropic or os.getenv("ANTHROPIC_API_KEY"),
+        firecrawl=s.firecrawl or os.getenv("FIRECRAWL_API_KEY"),
     )
 
-# --------------------------
-# Firecrawl helpers
-# --------------------------
-def fc_client(session: Keys) -> Firecrawl:
-    keys = resolve_keys(session)
-    if not keys.firecrawl:
+# -------------------- firecrawl --------------------
+def fc_client(s: Keys) -> Firecrawl:
+    k = resolve_keys(s)
+    if not k.firecrawl:
         raise gr.Error("Missing FIRECRAWL_API_KEY. Enter it in Keys → Save.")
-    return Firecrawl(api_key=keys.firecrawl)
+    return Firecrawl(api_key=k.firecrawl)
 
 @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=8))
-def fc_search(session: Keys, query: str, limit: int = 5, scrape_formats: Optional[List[str]] = None, location: Optional[str] = None) -> Dict[str, Any]:
-    fc = fc_client(session)
+def fc_search(s: Keys, query: str, limit: int = 5, scrape_formats: Optional[List[str]] = None, location: Optional[str] = None) -> Dict[str, Any]:
+    fc = fc_client(s)
     kwargs: Dict[str, Any] = {"query": query, "limit": limit}
     if location:
         kwargs["location"] = location
     if scrape_formats:
         kwargs["scrape_options"] = {"formats": scrape_formats}
-    result = fc.search(**kwargs)   # returns a Pydantic model
-    return _to_dict(result)
+    res = fc.search(**kwargs)
+    return _to_dict(res)
 
-@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=8))
-def fc_scrape(session: Keys, url: str, formats: Optional[List[str]] = None, timeout_ms: Optional[int] = None, mobile: bool = False) -> Dict[str, Any]:
-    fc = fc_client(session)
+@retry(stop=stop_after_attempt(2), wait=wait_exponential(multiplier=1, min=1, max=10))
+def fc_scrape(s: Keys, url: str, formats: Optional[List[str]] = None, timeout_ms: Optional[int] = None, mobile: bool = False) -> Dict[str, Any]:
+    fc = fc_client(s)
     kwargs: Dict[str, Any] = {"url": url}
     if formats:
         kwargs["formats"] = formats
+    # give slow pages more time; cap at 40s
     if timeout_ms:
-        kwargs["timeout"] = timeout_ms
+        kwargs["timeout"] = min(int(timeout_ms), 40000)
     if mobile:
         kwargs["mobile"] = True
-    result = fc.scrape(**kwargs)   # Pydantic model
-    return _to_dict(result)
+    res = fc.scrape(**kwargs)
+    return _to_dict(res)
 
-@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=8))
-def fc_crawl(session: Keys, url: str, max_pages: int = 25, formats: Optional[List[str]] = None) -> Dict[str, Any]:
-    fc = fc_client(session)
+@retry(stop=stop_after_attempt(2), wait=wait_exponential(multiplier=1, min=1, max=10))
+def fc_crawl(s: Keys, url: str, max_pages: int = 25, formats: Optional[List[str]] = None) -> Dict[str, Any]:
+    fc = fc_client(s)
     kwargs: Dict[str, Any] = {"url": url, "limit": max_pages}
     if formats:
         kwargs["scrape_options"] = {"formats": formats}
-    result = fc.crawl(**kwargs)    # Pydantic model
-    return _to_dict(result)
+    res = fc.crawl(**kwargs)
+    return _to_dict(res)
 
-# --------------------------
-# LLM helpers
-# --------------------------
+# -------------------- LLMs --------------------
 SYSTEM_STEER = (
     "You are ZEN's VibeCoder: extract web insights, generate clean scaffolds, "
     "and produce production-ready artifacts. Prefer structured outlines, code blocks, and checklists. "
     "When asked to clone or refactor, output file trees and exact text."
 )
 
-def use_openai(session: Keys):
-    keys = resolve_keys(session)
-    if not keys.openai:
+def use_openai(s: Keys):
+    k = resolve_keys(s)
+    if not k.openai:
         raise gr.Error("Missing OPENAI_API_KEY.")
     if OpenAI is None:
         raise gr.Error("OpenAI SDK not installed.")
-    return OpenAI(api_key=keys.openai)
+    return OpenAI(api_key=k.openai)
 
-def use_anthropic(session: Keys):
-    keys = resolve_keys(session)
-    if not keys.anthropic:
+def use_anthropic(s: Keys):
+    k = resolve_keys(s)
+    if not k.anthropic:
         raise gr.Error("Missing ANTHROPIC_API_KEY.")
     if anthropic is None:
         raise gr.Error("Anthropic SDK not installed.")
-    return anthropic.Anthropic(api_key=keys.anthropic)
-
-def llm_summarize(session: Keys, provider: str, model_name: str, prompt: str, context_md: str, temp: float=0.4) -> str:
+    return anthropic.Anthropic(api_key=k.anthropic)
+
+ANTHROPIC_FALLBACKS = [
+    # try a few known-good Sonnet identifiers
+    "claude-3-7-sonnet-2025-06-13",  # example new tag
+    "claude-3-7-sonnet",
+    "claude-3-5-sonnet-20241022",
+    "claude-3-5-sonnet-20240620",
+]
+
+OPENAI_FALLBACKS = [
+    "gpt-5",           # user-preferred
+    "gpt-4.1",         # safe fallback
+    "gpt-4o",          # vision-capable fallback
+    "gpt-4o-mini",     # economical fallback
+]
+
+def llm_summarize(s: Keys, provider: str, model_name: str, prompt: str, ctx_md: str, temp: float=0.4) -> str:
+    ctx = (ctx_md or "")[:150000]
     if provider == "openai":
-        client = use_openai(session)
-        resp = client.chat.completions.create(
-            model=model_name or "gpt-5",
-            temperature=temp,
-            messages=[
-                {"role": "system", "content": SYSTEM_STEER},
-                {"role": "user", "content": f"{prompt}\n\n=== SOURCE (markdown) ===\n{(context_md or '')[:150000]}"},
-            ],
-        )
-        return (resp.choices[0].message.content or "").strip()
+        client = use_openai(s)
+        candidates = [model_name] + OPENAI_FALLBACKS if model_name else OPENAI_FALLBACKS
+        last_err = None
+        for m in candidates:
+            try:
+                resp = client.chat.completions.create(
+                    model=m,
+                    temperature=temp,
+                    messages=[
+                        {"role": "system", "content": SYSTEM_STEER},
+                        {"role": "user", "content": f"{prompt}\n\n=== SOURCE (markdown) ===\n{ctx}"},
+                    ],
+                )
+                return (resp.choices[0].message.content or "").strip()
+            except Exception as e:
+                last_err = e
+                continue
+        raise gr.Error(f"OpenAI failed across fallbacks: {last_err}")
     else:
-        client = use_anthropic(session)
-        resp = client.messages.create(
-            model=model_name or "claude-3-5-sonnet-latest",
-            max_tokens=4000,
-            temperature=temp,
-            system=SYSTEM_STEER,
-            messages=[{"role": "user", "content": f"{prompt}\n\n=== SOURCE (markdown) ===\n{(context_md or '')[:150000]}"}],
-        )
-        chunks = []
-        for blk in resp.content:
-            # anthropic returns TextBlock objects
-            text = getattr(blk, "text", None)
-            if text:
-                chunks.append(text)
-        return "".join(chunks).strip()
-
-# --------------------------
-# ZIP export
-# --------------------------
+        client = use_anthropic(s)
+        candidates = [model_name] + ANTHROPIC_FALLBACKS if model_name else ANTHROPIC_FALLBACKS
+        last_err = None
+        for m in candidates:
+            try:
+                resp = client.messages.create(
+                    model=m,
+                    max_tokens=4000,
+                    temperature=temp,
+                    system=SYSTEM_STEER,
+                    messages=[{"role": "user", "content": f"{prompt}\n\n=== SOURCE (markdown) ===\n{ctx}"}],
+                )
+                chunks = []
+                for blk in resp.content:
+                    t = getattr(blk, "text", None)
+                    if t:
+                        chunks.append(t)
+                return "".join(chunks).strip()
+            except AnthropicNotFound as e:
+                last_err = e
+                continue
+            except Exception as e:
+                last_err = e
+                continue
+        raise gr.Error(f"Anthropic failed across fallbacks: {last_err}")
+
+# -------------------- ZIP export --------------------
 def pack_zip(pages: List[Dict[str, Any]]) -> bytes:
     mem = io.BytesIO()
     with zipfile.ZipFile(mem, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
@@ -183,18 +206,14 @@ def pack_zip(pages: List[Dict[str, Any]]) -> bytes:
             md = p.get("markdown") or p.get("data", {}).get("markdown") or p.get("content") or ""
             html = p.get("html") or p.get("data", {}).get("html") or ""
             links = p.get("links") or p.get("data", {}).get("links") or []
-            if md:
-                zf.writestr(f"{i:03d}_{slug}.md", md)
-            if html:
-                zf.writestr(f"{i:03d}_{slug}.html", html)
+            if md:   zf.writestr(f"{i:03d}_{slug}.md", md)
+            if html: zf.writestr(f"{i:03d}_{slug}.html", html)
             manifest.append({"url": url, "title": p.get("title") or p.get("metadata", {}).get("title"), "links": links})
         zf.writestr("manifest.json", json.dumps(manifest, indent=2))
     mem.seek(0)
     return mem.read()
 
-# --------------------------
-# Gradio actions
-# --------------------------
+# -------------------- actions --------------------
 def save_keys(openai_key, anthropic_key, firecrawl_key):
     return Keys(
         openai=(openai_key or "").strip() or None,
@@ -202,67 +221,73 @@ def save_keys(openai_key, anthropic_key, firecrawl_key):
         firecrawl=(firecrawl_key or "").strip() or None,
     ), gr.Info("Keys saved to this session. (Env vars still apply if set.)")
 
-def action_search(session: Keys, query: str, limit: int, scrape_content: bool, location: str):
+def action_search(sess: Keys, query: str, limit: int, scrape_content: bool, location: str):
     if not query.strip():
         raise gr.Error("Enter a search query.")
     formats = ["markdown", "links"] if scrape_content else None
-    result = fc_search(session, query=query.strip(), limit=limit, scrape_formats=formats, location=(location or None))
-    # Firecrawl v2 search returns {'data': {'web': [...], 'news': [...], ...}} or {'data': [...]}
-    data = result.get("data", result)  # tolerate both shapes
+    res = fc_search(sess, query=query.strip(), limit=limit, scrape_formats=formats, location=(location or None))
+    data = res.get("data", res)
     items: List[Any] = []
     if isinstance(data, dict):
         for bucket in ("web", "news", "images", "videos", "discussion"):
-            items.extend(_to_dict(data.get(bucket, [])))
+            b = data.get(bucket)
+            if b:
+                items.extend(_listify(_to_dict(b)))
     elif isinstance(data, list):
         items = _to_dict(data)
     else:
-        items = [_to_dict(data)]
+        items = _listify(_to_dict(data))
+    if not items:
+        return _pretty_json(res)  # show raw result if buckets are empty
     return json.dumps(items, indent=2)
 
-def action_scrape(session: Keys, url: str, mobile: bool, formats_sel: List[str], timeout_ms: int):
+def action_scrape(sess: Keys, url: str, mobile: bool, formats_sel: List[str], timeout_ms: int):
     if not url.strip():
         raise gr.Error("Enter a URL.")
-    formats = formats_sel or ["markdown", "html", "links"]
-    out = fc_scrape(session, url.strip(), formats=formats, timeout_ms=(timeout_ms or None), mobile=mobile)
-    pretty = _pretty_json(out)
-    md = (
-        out.get("markdown")
-        or out.get("data", {}).get("markdown")
-        or out.get("content")
-        or ""
-    )
-    return pretty, md
+    formats = formats_sel or ["markdown", "links"]
+    try:
+        out = fc_scrape(sess, url.strip(), formats=formats, timeout_ms=(timeout_ms or 15000), mobile=mobile)
+        pretty = _pretty_json(out)
+        md = out.get("markdown") or out.get("data", {}).get("markdown") or out.get("content") or ""
+        return pretty, md
+    except RetryError as e:
+        return f"<!> Scrape timed out after retries. Try increasing timeout, unchecking 'mobile', or limiting formats.\n\n{e}", ""
+    except Exception as e:
+        return f"<!> Scrape error: {e}", ""
 
-def action_crawl(session: Keys, base_url: str, max_pages: int, formats_sel: List[str]):
+def action_crawl(sess: Keys, base_url: str, max_pages: int, formats_sel: List[str]):
     if not base_url.strip():
         raise gr.Error("Enter a base URL to crawl.")
     formats = formats_sel or ["markdown", "links"]
-    out = fc_crawl(session, base_url.strip(), max_pages=max_pages, formats=formats)
-    pages = out.get("data")
-    if not isinstance(pages, list) or not pages:
-        raise gr.Error("Crawl returned no pages.")
-    zip_bytes = pack_zip(pages)
-    return gr.File.update(value=io.BytesIO(zip_bytes), visible=True, filename="site_clone.zip"), f"Crawled {len(pages)} pages. ZIP is ready."
-
-def action_generate(session: Keys, provider: str, model_name: str, sys_prompt: str, user_prompt: str, context_md: str, temp: float):
+    try:
+        out = fc_crawl(sess, base_url.strip(), max_pages=max_pages, formats=formats)
+        pages = out.get("data")
+        if not isinstance(pages, list) or not pages:
+            raise gr.Error("Crawl returned no pages.")
+        zip_bytes = pack_zip(pages)
+        return gr.File.update(value=io.BytesIO(zip_bytes), visible=True, filename="site_clone.zip"), f"Crawled {len(pages)} pages. ZIP is ready."
+    except RetryError as e:
+        return gr.File.update(visible=False), f"<!> Crawl timed out after retries. Reduce Max Pages or try again.\n\n{e}"
+    except Exception as e:
+        return gr.File.update(visible=False), f"<!> Crawl error: {e}"
+
+def action_generate(sess: Keys, provider: str, model_name: str, sys_prompt: str, user_prompt: str, context_md: str, temp: float):
     if not user_prompt.strip():
         raise gr.Error("Enter a prompt or click a starter tile.")
-    model = (model_name or "").strip() or ("gpt-5" if provider == "openai" else "claude-3-5-sonnet-latest")
+    model = (model_name or "").strip()
     steer = (sys_prompt or "").strip()
     prompt = (("SYSTEM:\n" + steer + "\n\n") if steer else "") + user_prompt.strip()
-    out = llm_summarize(session, provider, model, prompt, context_md or "", temp=temp)
+    out = llm_summarize(sess, provider, model, prompt, context_md or "", temp=temp)
     return out
 
-# --------------------------
-# UI (Blocks)
-# --------------------------
+# -------------------- UI --------------------
 with gr.Blocks(css="#keys .wrap.svelte-1ipelgc { filter: none !important; }") as demo:
     gr.Markdown("## ZEN VibeCoder — Web Clone & Research Foundry")
     session_state = gr.State(Keys())
 
     with gr.Accordion("🔐 Keys (session)", open=True):
         with gr.Row():
-            openai_key = gr.Textbox(label="OPENAI_API_KEY (GPT-5)", type="password", placeholder="sk-...", value=os.getenv("OPENAI_API_KEY") or "")
+            openai_key = gr.Textbox(label="OPENAI_API_KEY (GPT-5 / fallbacks)", type="password", placeholder="sk-...", value=os.getenv("OPENAI_API_KEY") or "")
             anthropic_key = gr.Textbox(label="ANTHROPIC_API_KEY (Claude Sonnet)", type="password", placeholder="anthropic-key...", value=os.getenv("ANTHROPIC_API_KEY") or "")
             firecrawl_key = gr.Textbox(label="FIRECRAWL_API_KEY", type="password", placeholder="fc-...", value=os.getenv("FIRECRAWL_API_KEY") or "")
         save_btn = gr.Button("Save keys", variant="primary")
@@ -271,7 +296,7 @@ with gr.Blocks(css="#keys .wrap.svelte-1ipelgc { filter: none !important; }") as
 
     with gr.Tabs():
         with gr.Tab("🔎 Search"):
-            query = gr.Textbox(label="Query", placeholder='ex: "best open-source vector databases in 2025 site:docs"')
+            query = gr.Textbox(label="Query", placeholder='ex: site:docs "vector database" 2025')
             with gr.Row():
                 limit = gr.Slider(1, 20, value=6, step=1, label="Limit")
                 scrape_content = gr.Checkbox(label="Also scrape results (markdown + links)", value=True)
@@ -283,7 +308,7 @@ with gr.Blocks(css="#keys .wrap.svelte-1ipelgc { filter: none !important; }") as
         with gr.Tab("🕸️ Scrape • Crawl • Clone"):
             with gr.Row():
                 target_url = gr.Textbox(label="URL to Scrape", placeholder="https://example.com")
-                timeout_ms = gr.Number(label="Timeout (ms)", value=15000)
+                timeout_ms = gr.Number(label="Timeout (ms, max 40000)", value=15000)
             with gr.Row():
                 formats_sel = gr.CheckboxGroup(choices=["markdown","html","links","screenshot"], value=["markdown","links"], label="Formats")
                 mobile = gr.Checkbox(label="Emulate mobile", value=False)
@@ -306,7 +331,7 @@ with gr.Blocks(css="#keys .wrap.svelte-1ipelgc { filter: none !important; }") as
         with gr.Tab("✨ Vibe Code (Synthesis)"):
             with gr.Row():
                 provider = gr.Radio(choices=["openai","anthropic"], value="openai", label="Provider")
-                model_name = gr.Textbox(label="Model (override)", placeholder="gpt-5  |  claude-3-5-sonnet-latest")
+                model_name = gr.Textbox(label="Model (override)", placeholder="(blank = auto fallback)")
                 temp = gr.Slider(0.0, 1.2, value=0.4, step=0.05, label="Temperature")
 
             sys_prompt = gr.Textbox(label="System Style (optional)",
@@ -347,5 +372,4 @@ with gr.Blocks(css="#keys .wrap.svelte-1ipelgc { filter: none !important; }") as
     gr.Markdown("Built for **ZEN Arena** pipelines. Export ZIPs → ingest → credentialize via ZEN Cards.")
 
 if __name__ == "__main__":
-    # If SSR causes issues in your Space, you can disable it by: demo.launch(ssr_mode=False)
-    demo.launch()
+    demo.launch(ssr_mode=False)