Update app.py

app.py

@@ -1,9 +1,13 @@
 import os
+import time
 import gradio as gr
 from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
 from llama_index.embeddings.mixedbreadai import MixedbreadAIEmbedding
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.llms.groq import Groq
 from llama_parse import LlamaParse
+import mixedbread_ai
+from mixedbread_ai.core.api_error import ApiError
 
 # API keys
 llama_cloud_key = os.environ.get("LLAMA_CLOUD_API_KEY")
@@ -14,7 +18,18 @@ if not (llama_cloud_key and groq_key and mxbai_key):
 
 # Model names
 llm_model_name = "llama-3.1-70b-versatile"
-embed_model_name = "mixedbread-ai/mxbai-embed-large-v1"
+embed_model_name = "mxbai-embed-large-v1"  # Mixedbread AI model
+fallback_embed_model = "sentence-transformers/all-MiniLM-L6-v2"  # Fallback model
+
+# Configure Mixedbread AI SDK
+mixedbread_config = mixedbread_ai.Configuration(
+    api_key=mxbai_key,
+    retry_on=[503],  # Retry on 503 Service Unavailable
+    max_retries=3,
+    retry_delay=2.0,  # Seconds between retries
+    timeout=30.0,  # Request timeout
+)
+mixedbread_client = mixedbread_ai.Client(configuration=mixedbread_config)
 
 # Initialize the parser
 parser = LlamaParse(api_key=llama_cloud_key, result_type="markdown")
@@ -37,8 +52,20 @@ file_extractor = {
 }
 
 # Initialize models with error handling
+def initialize_embed_model():
+    try:
+        return MixedbreadAIEmbedding(
+            api_key=mxbai_key,
+            model_name=embed_model_name,
+            mxbai_client=mixedbread_client,  # Use configured SDK client
+        )
+    except Exception as e:
+        print(f"Failed to initialize Mixedbread AI embedding: {str(e)}")
+        print("Falling back to local HuggingFace embedding model.")
+        return HuggingFaceEmbedding(model_name=fallback_embed_model)
+
 try:
-    embed_model = MixedbreadAIEmbedding(api_key=mxbai_key, model_name=embed_model_name)
+    embed_model = initialize_embed_model()
     llm = Groq(model=llm_model_name, api_key=groq_key)
 except Exception as e:
     raise RuntimeError(f"Failed to initialize models: {str(e)}")
@@ -61,14 +88,20 @@ def load_files(file_path: str):
             input_files=[file_path], 
             file_extractor=file_extractor
         ).load_data()
-        vector_index = VectorStoreIndex.from_documents(
-            document, 
-            embed_model=embed_model
-        )
-        filename = os.path.basename(file_path)
-        return f"Ready to provide responses based on: {filename}"
+        
+        try:
+            vector_index = VectorStoreIndex.from_documents(
+                document, 
+                embed_model=embed_model
+            )
+            filename = os.path.basename(file_path)
+            return f"Ready to provide responses based on: {filename}"
+        except ApiError as e:
+            return f"Error processing file with Mixedbread AI API: {str(e)}. Status code: {e.status_code}"
+        except Exception as e:
+            return f"Unexpected error processing file: {str(e)}"
     except Exception as e:
-        return f"Error processing file: {str(e)}"
+        return f"Error loading file: {str(e)}"
 
 # Respond function
 def respond(message, history):
@@ -115,7 +148,7 @@ with gr.Blocks(
         with gr.Column(scale=3):
             chatbot = gr.ChatInterface(
                 fn=respond,
-                chatbot=gr.Chatbot(height=300, type="messages"),  # Fixed deprecated warning
+                chatbot=gr.Chatbot(height=300, type="messages"),
                 theme="soft",
                 show_progress="full",
                 textbox=gr.Textbox(