Update app.py

app.py

@@ -14,6 +14,8 @@ import shutil
 import logging
 import chromadb
 import tempfile
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
+import requests
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -26,8 +28,8 @@ if os.environ["HUGGINGFACEHUB_API_TOKEN"] == "default-token":
 
 # Model and embedding options
 LLM_MODELS = {
-    "Lightweight (Gemma-2B)": "google/gemma-2b-it",
     "Balanced (Mixtral-8x7B)": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "Lightweight (Gemma-2B)": "google/gemma-2b-it",
     "High Accuracy (Llama-3-8B)": "meta-llama/Llama-3-8b-hf"
 }
 
@@ -133,13 +135,12 @@ def process_documents(files, chunk_size, chunk_overlap, embedding_model):
 
     # Create vector store
     try:
-        # Use a unique collection name to avoid conflicts
+        # Use in-memory Chroma client to avoid filesystem issues
         collection_name = f"doctalk_collection_{int(time.time())}"
-        client = chromadb.PersistentClient(path=PERSIST_DIRECTORY)
+        client = chromadb.Client()
         vector_store = Chroma.from_documents(
             documents=doc_splits,
             embedding=embeddings,
-            persist_directory=PERSIST_DIRECTORY,
             collection_name=collection_name
         )
         return f"Processed {len(documents)} documents into {len(doc_splits)} chunks.", None
@@ -147,7 +148,12 @@ def process_documents(files, chunk_size, chunk_overlap, embedding_model):
         logger.error(f"Error creating vector store: {str(e)}")
         return f"Error creating vector store: {str(e)}", None
 
-# Function to initialize QA chain
+# Function to initialize QA chain with retry logic
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((requests.exceptions.HTTPError, requests.exceptions.ConnectionError))
+)
 def initialize_qa_chain(llm_model, temperature):
     global qa_chain
     if not vector_store:
@@ -159,20 +165,37 @@ def initialize_qa_chain(llm_model, temperature):
             task="text-generation",
             temperature=float(temperature),
             max_new_tokens=512,
-            huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"]
+            huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"],
+            timeout=30
         )
+        # Dynamically set k based on vector store size
+        collection = vector_store._collection
+        doc_count = collection.count()
+        k = min(3, doc_count) if doc_count > 0 else 1
         qa_chain = ConversationalRetrievalChain.from_llm(
             llm=llm,
-            retriever=vector_store.as_retriever(search_kwargs={"k": 3}),
+            retriever=vector_store.as_retriever(search_kwargs={"k": k}),
             memory=memory
         )
-        logger.info(f"Initialized QA chain with {llm_model}.")
+        logger.info(f"Initialized QA chain with {llm_model} and k={k}.")
         return "QA Doctor: QA chain initialized successfully.", None
+    except requests.exceptions.HTTPError as e:
+        logger.error(f"HTTP error initializing QA chain for {llm_model}: {str(e)}")
+        if "503" in str(e):
+            return f"Error: Hugging Face API temporarily unavailable for {llm_model}. Try 'Balanced (Mixtral-8x7B)' or wait and retry.", None
+        elif "403" in str(e):
+            return f"Error: Access denied for {llm_model}. Ensure your HF token has access.", None
+        return f"Error initializing QA chain: {str(e)}.", None
     except Exception as e:
         logger.error(f"Error initializing QA chain for {llm_model}: {str(e)}")
         return f"Error initializing QA chain: {str(e)}. Ensure your HF token has access to {llm_model}.", None
 
-# Function to handle user query
+# Function to handle user query with retry logic
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((requests.exceptions.HTTPError, requests.exceptions.ConnectionError))
+)
 def answer_question(question, llm_model, embedding_model, temperature, chunk_size, chunk_overlap):
     global chat_history
     if not vector_store:
@@ -183,11 +206,18 @@ def answer_question(question, llm_model, embedding_model, temperature, chunk_siz
         return "Please enter a valid question.", chat_history
 
     try:
-        response = qa_chain({"question": question})["answer"]
+        response = qa_chain.invoke({"question": question})["answer"]
         chat_history.append({"role": "user", "content": question})
         chat_history.append({"role": "assistant", "content": response})
         logger.info(f"Answered question: {question}")
         return response, chat_history
+    except requests.exceptions.HTTPError as e:
+        logger.error(f"HTTP error answering question: {str(e)}")
+        if "503" in str(e):
+            return f"Error: Hugging Face API temporarily unavailable for {llm_model}. Try 'Balanced (Mixtral-8x7B)' or wait and retry.", chat_history
+        elif "403" in str(e):
+            return f"Error: Access denied for {llm_model}. Ensure your HF token has access.", chat_history
+        return f"Error answering question: {str(e)}", chat_history
     except Exception as e:
         logger.error(f"Error answering question: {str(e)}")
         return f"Error answering question: {str(e)}", chat_history
@@ -242,7 +272,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DocTalk: Document Q&A Chatbot") as
             status = gr.Textbox(label="Status", interactive=False)
 
         with gr.Column(scale=1):
-            llm_model = gr.Dropdown(choices=list(LLM_MODELS.keys()), label="Select LLM Model", value="Lightweight (Gemma-2B)")
+            llm_model = gr.Dropdown(choices=list(LLM_MODELS.keys()), label="Select LLM Model", value="Balanced (Mixtral-8x7B)")
             embedding_model = gr.Dropdown(choices=list(EMBEDDING_MODELS.keys()), label="Select Embedding Model", value="Lightweight (MiniLM-L6)")
             temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.7, label="Temperature")
             chunk_size = gr.Slider(minimum=500, maximum=2000, step=100, value=1000, label="Chunk Size")