Update app.py

app.py

@@ -1,125 +1,189 @@
+"""
+Doc-Q&A app (Gradio 5.x + Llama-Index 0.12.x, June 2025)
+
+Key upgrades
+------------
+▪ Gradio 5.34—new event system (`upload`, `clear` etc.)  
+▪ Llama-Index 0.12.42—`VectorStoreIndex.from_documents` signature unchanged  
+▪ MixedbreadAIEmbedding 0.3.0  ➜ supports `batch_size`, `timeout`  
+▪ Tenacity for exponential-back-off when MXBAI returns 5xx / rate limits
+"""
+
+from __future__ import annotations
+
 import os
+from pathlib import Path
+from typing import List
+
 import gradio as gr
+from tenacity import retry, wait_exponential, stop_after_attempt
+from mixedbread_ai.core.api_error import ApiError
+
 from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
 from llama_index.embeddings.mixedbreadai import MixedbreadAIEmbedding
 from llama_index.llms.groq import Groq
 from llama_parse import LlamaParse
 
-# ────────────────────────────────
-# 1.  Check environment variables
-# ────────────────────────────────
-llama_cloud_key = os.getenv("LLAMA_CLOUD_API_KEY")
-groq_key        = os.getenv("GROQ_API_KEY")
-mxbai_key       = os.getenv("MXBAI_API_KEY")
+# ──────────────────────────────────────────────────────────────────
+# 1.  Environment variables (fail-fast if missing)
+# ──────────────────────────────────────────────────────────────────
+LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
+GROQ_API_KEY        = os.getenv("GROQ_API_KEY")
+MXBAI_API_KEY       = os.getenv("MXBAI_API_KEY")
 
-if not (llama_cloud_key and groq_key and mxbai_key):
+if not all([LLAMA_CLOUD_API_KEY, GROQ_API_KEY, MXBAI_API_KEY]):
     raise EnvironmentError(
-        "LLAMA_CLOUD_API_KEY, GROQ_API_KEY and MXBAI_API_KEY must be set."
+        "LLAMA_CLOUD_API_KEY, GROQ_API_KEY and MXBAI_API_KEY must be set in the env."
     )
 
-# ────────────────────────────────
-# 2.  Model / parser setup
-# ────────────────────────────────
-LLM_MODEL   = "llama-3.1-70b-versatile"
-EMBED_MODEL = "mixedbread-ai/mxbai-embed-large-v1"
+# ──────────────────────────────────────────────────────────────────
+# 2.  Models & parsers  (latest defaults - June 2025)
+# ──────────────────────────────────────────────────────────────────
+LLM_MODEL   = "llama-3.1-70b-versatile"                 # Groq’s best for Q&A
+EMBED_MODEL = "mixedbread-ai/mxbai-embed-large-v1"      # 1024-dim
+
+parser = LlamaParse(api_key=LLAMA_CLOUD_API_KEY, result_type="markdown")
 
-parser       = LlamaParse(api_key=llama_cloud_key, result_type="markdown")
-file_extractor = {ext: parser for ext in (
+SUPPORTED_EXTS = (
     ".pdf", ".docx", ".doc", ".txt", ".csv", ".xlsx",
     ".pptx", ".html", ".jpg", ".jpeg", ".png", ".webp", ".svg",
-)}
-embed_model  = MixedbreadAIEmbedding(api_key=mxbai_key, model_name=EMBED_MODEL)
-llm          = Groq(model=LLM_MODEL, api_key=groq_key)
-
-# Global cache for the current document
-vector_index = None
-
-
-# ────────────────────────────────
-# 3.  Helper functions
-# ────────────────────────────────
-def load_files(file_path: str) -> str:
-    """Parse the uploaded document and build a vector index."""
+)
+file_extractor = {ext: parser for ext in SUPPORTED_EXTS}
+
+embed_model = MixedbreadAIEmbedding(
+    api_key     = MXBAI_API_KEY,
+    model_name  = EMBED_MODEL,
+    batch_size  = 8,          # keep requests < 100 KB
+    timeout     = 60,         # generous server-side processing window
+)
+
+llm = Groq(model=LLM_MODEL, api_key=GROQ_API_KEY)
+
+# A simple global cache (could be swapped for Redis, etc.)
+vector_index: VectorStoreIndex | None = None
+
+
+# ──────────────────────────────────────────────────────────────────
+# 3.  Helper wrappers
+# ──────────────────────────────────────────────────────────────────
+@retry(
+    wait=wait_exponential(multiplier=2, min=4, max=32),
+    stop=stop_after_attempt(4),
+    retry_error_callback=lambda retry_state: None,   # bubble up as None
+    reraise=False,
+)
+def _safe_build_index(docs) -> VectorStoreIndex | None:
+    """Retry MXBAI 503 / 429 transparently."""
+    try:
+        return VectorStoreIndex.from_documents(docs, embed_model=embed_model)
+    except ApiError as e:
+        # Tenacity will catch and retry unless non-5xx / non-429
+        if e.status_code not in (429, 500, 502, 503, 504):
+            raise
+        raise   # trigger retry
+
+
+def load_files(file: Path | None) -> str:
+    """Parse uploaded file and build vector index (with retries)."""
     global vector_index
-    if not file_path:
+    if file is None:
         return "⚠️  No file selected."
 
-    if not any(file_path.endswith(ext) for ext in file_extractor):
-        return ("⚠️  Unsupported file type. "
-                f"Allowed: {', '.join(file_extractor.keys())}")
+    if file.suffix.lower() not in SUPPORTED_EXTS:
+        allow = ", ".join(SUPPORTED_EXTS)
+        return f"⚠️  Unsupported file type. Allowed: {allow}"
 
     docs = SimpleDirectoryReader(
-        input_files=[file_path], file_extractor=file_extractor
+        input_files=[str(file)],
+        file_extractor=file_extractor,
     ).load_data()
 
-    vector_index = VectorStoreIndex.from_documents(docs, embed_model=embed_model)
-    return f"✅ Parsed **{os.path.basename(file_path)}**. Ask away!"
+    idx = _safe_build_index(docs)
+    if idx is None:
+        return "🚧 Embedding service busy. Please retry in ~1 minute."
+
+    vector_index = idx
+    return f"✅ Parsed **{file.name}** — you can start chatting!"
 
 
-def respond(message: str, history: list) -> str:
-    """Chat handler. Streams partial tokens back to the UI."""
+def respond(message: str, history: List[List[str]]):
+    """Stream answer chunks to the Chatbot."""
     if vector_index is None:
-        return "➡️  Please upload a document first."
+        yield "➡️  Please upload a document first."
+        return
 
-    query_engine   = vector_index.as_query_engine(streaming=True, llm=llm)
-    streaming_resp = query_engine.query(message)
+    query_engine = vector_index.as_query_engine(streaming=True, llm=llm)
+    response = query_engine.query(message)
 
     partial = ""
-    for chunk in streaming_resp.response_gen:
-        partial += chunk
-        yield partial     # <─ streaming to the frontend
+    for token in response.response_gen:
+        partial += token
+        yield partial
 
 
-def clear_state():
-    """Reset everything."""
+def clear():
+    """Reset everything (file widget + status + index)."""
     global vector_index
     vector_index = None
-    return [None, ""]
+    return None, "", None   # file_input, status_md, chatbot history
 
 
-# ────────────────────────────────
-# 4.  Gradio UI
-# ────────────────────────────────
+# ──────────────────────────────────────────────────────────────────
+# 4.  Gradio UI  (5.x syntax)
+# ──────────────────────────────────────────────────────────────────
 with gr.Blocks(
     theme=gr.themes.Default(
         primary_hue="green",
         secondary_hue="blue",
-        font=[gr.themes.GoogleFont("Poppins")]
+        font=[gr.themes.GoogleFont("Poppins")],
     ),
-    css="footer {visibility: hidden}"
+    css="footer {visibility:hidden}",
 ) as demo:
 
-    gr.Markdown("<h1 style='text-align:center'>DataCamp Doc Q&A 🤖📃</h1>")
+    gr.Markdown("<h1 style='text-align:center'>DataCamp Doc Q&A 🤖📑</h1>")
+
     with gr.Row():
         with gr.Column(scale=1):
-            file_input = gr.File(file_count="single",
-                                 type="filepath",
-                                 label="Upload document")
+            file_input = gr.File(
+                label="Upload document",
+                file_count="single",
+                type="filepath",
+                show_label=True,
+            )
+            status_md  = gr.Markdown()
             with gr.Row():
-                submit_btn = gr.Button("Submit", variant="primary")
-                clear_btn  = gr.Button("Clear")
-            status_box = gr.Markdown()
+                clear_btn = gr.Button("Reset 🔄", variant="secondary")
 
         with gr.Column(scale=3):
-            chat = gr.ChatInterface(
-                fn=respond,
-                chatbot=gr.Chatbot(height=300),
-                show_progress="full",   # keep the nice progress bar
-                textbox=gr.Textbox(
-                    placeholder="Ask a question about the uploaded document…",
-                    container=False,
-                ),
+            chatbot = gr.Chatbot(height=340)
+            txt_box = gr.Textbox(
+                placeholder="Ask something about the uploaded document…",
+                container=False,
+                scale=7,
             )
+            send_btn = gr.Button("Send", variant="primary")
 
-    submit_btn.click(load_files, inputs=file_input, outputs=status_box)
-    clear_btn.click(clear_state, outputs=[file_input, status_box])
+    # events (v5 style)
+    file_input.upload(
+        fn=load_files,
+        inputs=file_input,
+        outputs=status_md,
+    )
+    send_btn.click(
+        fn=respond,
+        inputs=[txt_box, chatbot],
+        outputs=chatbot,
+    )
+    clear_btn.click(
+        fn=clear,
+        outputs=[file_input, status_md, chatbot],
+    )
 
-# Disable OpenAPI generation (avoids the bool/‘const’ bug) …
+# optional: disable public OpenAPI schema (old crash guard)
 demo.queue(api_open=False)
 
-# ────────────────────────────────
+# ──────────────────────────────────────────────────────────────────
 # 5.  Launch
-# ────────────────────────────────
+# ──────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
-    # …and make a public share link so the container doesn’t choke on localhost
     demo.launch(share=True, server_name="0.0.0.0", server_port=7860)