Update app.py

app.py

@@ -1,189 +1,130 @@
-"""
-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.  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 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 in the env."
+# API keys
+llama_cloud_key = os.environ.get("LLAMA_CLOUD_API_KEY")
+groq_key = os.environ.get("GROQ_API_KEY")
+mxbai_key = os.environ.get("MXBAI_API_KEY")
+if not (llama_cloud_key and groq_key and mxbai_key):
+    raise ValueError(
+        "API Keys not found! Ensure they are passed to the Docker container."
     )
 
-# ──────────────────────────────────────────────────────────────────
-# 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")
-
-SUPPORTED_EXTS = (
-    ".pdf", ".docx", ".doc", ".txt", ".csv", ".xlsx",
-    ".pptx", ".html", ".jpg", ".jpeg", ".png", ".webp", ".svg",
-)
-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 file is None:
-        return "⚠️  No file selected."
-
-    if file.suffix.lower() not in SUPPORTED_EXTS:
-        allow = ", ".join(SUPPORTED_EXTS)
-        return f"⚠️  Unsupported file type. Allowed: {allow}"
+# models name
+llm_model_name = "llama-3.1-70b-versatile"
+embed_model_name = "mixedbread-ai/mxbai-embed-large-v1"
 
-    docs = SimpleDirectoryReader(
-        input_files=[str(file)],
-        file_extractor=file_extractor,
-    ).load_data()
+# Initialize the parser
+parser = LlamaParse(api_key=llama_cloud_key, result_type="markdown")
 
-    idx = _safe_build_index(docs)
-    if idx is None:
-        return "🚧 Embedding service busy. Please retry in ~1 minute."
+# Define file extractor with various common extensions
+file_extractor = {
+    ".pdf": parser,
+    ".docx": parser,
+    ".doc": parser,
+    ".txt": parser,
+    ".csv": parser,
+    ".xlsx": parser,
+    ".pptx": parser,
+    ".html": parser,
+    ".jpg": parser,
+    ".jpeg": parser,
+    ".png": parser,
+    ".webp": parser,
+    ".svg": parser,
+}
 
-    vector_index = idx
-    return f"✅ Parsed **{file.name}** — you can start chatting!"
+# Initialize the embedding model
+embed_model = MixedbreadAIEmbedding(api_key=mxbai_key, model_name=embed_model_name)
 
+# Initialize the LLM
 
-def respond(message: str, history: List[List[str]]):
-    """Stream answer chunks to the Chatbot."""
-    if vector_index is None:
-        yield "➡️  Please upload a document first."
-        return
+llm = Groq(model="llama-3.1-70b-versatile", api_key=groq_key)
 
-    query_engine = vector_index.as_query_engine(streaming=True, llm=llm)
-    response = query_engine.query(message)
 
-    partial = ""
-    for token in response.response_gen:
-        partial += token
-        yield partial
-
-
-def clear():
-    """Reset everything (file widget + status + index)."""
+# File processing function
+def load_files(file_path: str):
+    global vector_index
+    if not file_path:
+        return "No file path provided. Please upload a file."
+    
+    valid_extensions = ', '.join(file_extractor.keys())
+    if not any(file_path.endswith(ext) for ext in file_extractor):
+        return f"The parser can only parse the following file types: {valid_extensions}"
+
+    document = SimpleDirectoryReader(input_files=[file_path], file_extractor=file_extractor).load_data()
+    vector_index = VectorStoreIndex.from_documents(document, embed_model=embed_model)
+    print(f"Parsing completed for: {file_path}")
+    filename = os.path.basename(file_path)
+    return f"Ready to provide responses based on: {filename}"
+
+
+# Respond function
+def respond(message, history):
+    try:
+        # Use the preloaded LLM
+        query_engine = vector_index.as_query_engine(streaming=True, llm=llm)
+        streaming_response = query_engine.query(message)
+        partial_text = ""
+        for new_text in streaming_response.response_gen:
+            partial_text += new_text
+            # Yield an empty string to cleanup the message textbox and the updated conversation history
+            yield partial_text
+    except (AttributeError, NameError):
+        print("An error occurred while processing your request.")
+        yield "Please upload the file to begin chat."
+
+
+# Clear function
+def clear_state():
     global vector_index
     vector_index = None
-    return None, "", None   # file_input, status_md, chatbot history
+    return [None, None, None]
 
 
-# ──────────────────────────────────────────────────────────────────
-# 4.  Gradio UI  (5.x syntax)
-# ──────────────────────────────────────────────────────────────────
+# UI Setup
 with gr.Blocks(
     theme=gr.themes.Default(
         primary_hue="green",
         secondary_hue="blue",
         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("# DataCamp Doc Q&A 🤖📃")
     with gr.Row():
         with gr.Column(scale=1):
             file_input = gr.File(
-                label="Upload document",
-                file_count="single",
-                type="filepath",
-                show_label=True,
+                file_count="single", type="filepath", label="Upload Document"
             )
-            status_md  = gr.Markdown()
             with gr.Row():
-                clear_btn = gr.Button("Reset 🔄", variant="secondary")
-
+                btn = gr.Button("Submit", variant="primary")
+                clear = gr.Button("Clear")
+            output = gr.Textbox(label="Status")
         with gr.Column(scale=3):
-            chatbot = gr.Chatbot(height=340)
-            txt_box = gr.Textbox(
-                placeholder="Ask something about the uploaded document…",
-                container=False,
-                scale=7,
+            chatbot = gr.ChatInterface(
+                fn=respond,
+                chatbot=gr.Chatbot(height=300),
+                theme="soft",
+                show_progress="full",
+                textbox=gr.Textbox(
+                    placeholder="Ask questions about the uploaded document!",
+                    container=False,
+                ),
             )
-            send_btn = gr.Button("Send", variant="primary")
 
-    # 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,
+    # Set up Gradio interactions
+    btn.click(fn=load_files, inputs=file_input, outputs=output)
+    clear.click(
+        fn=clear_state,  # Use the clear_state function
+        outputs=[file_input, output],
     )
-    clear_btn.click(
-        fn=clear,
-        outputs=[file_input, status_md, chatbot],
-    )
-
-# optional: disable public OpenAPI schema (old crash guard)
-demo.queue(api_open=False)
 
-# ──────────────────────────────────────────────────────────────────
-# 5.  Launch
-# ──────────────────────────────────────────────────────────────────
+# Launch the demo
 if __name__ == "__main__":
-    demo.launch(share=True, server_name="0.0.0.0", server_port=7860)
+    demo.launch()