Update app.py

app.py

@@ -8,7 +8,6 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain.chains import RetrievalQA
-# Use existing imports since langchain_aws is not available
 from langchain_community.chat_models import BedrockChat
 from langchain_openai import ChatOpenAI
 from langchain_community.llms import Ollama
@@ -246,12 +245,15 @@ def create_interface():
             """)
         return demo
 
-    # Initialize agents - changed to initialize lazily to avoid startup errors
+    # Initialize agents dictionary - will be initialized on demand
     audit_agents = {}
     
     with gr.Blocks(theme=gr.themes.Base()) as demo:
         gr.Markdown("# 🔍 Amy - Your Audit Copilot")
         
+        # Status indicator for initialization and operations
+        status_message = gr.Textbox(label="Status", value="Ready")
+        
         with gr.Row():
             with gr.Column(scale=1):
                 file_upload = gr.File(
@@ -293,9 +295,6 @@ def create_interface():
                 )
                 query_button = gr.Button("Query")
                 query_output = gr.Markdown(label="Response")
-                
-        # Status indicator for initialization and operations
-        status_message = gr.Textbox(label="Status", value="Ready")
         
         # Track the selected model
         selected_model = gr.State("claude-3-sonnet")
@@ -309,54 +308,119 @@ def create_interface():
             
         model_tabs.select(update_selected_model, outputs=[selected_model])
         
-        # Lazy initialization of models when first used
+        # COMPLETELY REVISED: Initialize an agent and return both agent and status message
         def get_or_initialize_agent(model_name):
-            if model_name not in audit_agents:
-                try:
-                    status_message.update(value=f"Initializing {model_name}...")
-                    config = llm_configs[model_name]
-                    audit_agents[model_name] = AuditAgent(config["name"], config["provider"])
-                    status_message.update(value=f"{model_name} initialized successfully")
-                except Exception as e:
-                    status_message.update(value=f"Error initializing {model_name}: {str(e)}")
-                    return None
-            return audit_agents[model_name]
+            """Initialize an agent if not already initialized and return a status message"""
+            init_message = f"Initializing {model_name}..."
+            
+            # If agent already exists, return it with a status message
+            if model_name in audit_agents:
+                return audit_agents[model_name], f"{model_name} ready"
+            
+            # Try to initialize the agent
+            try:
+                config = llm_configs[model_name]
+                logging.info(init_message)
+                agent = AuditAgent(config["name"], config["provider"])
+                audit_agents[model_name] = agent
+                success_message = f"{model_name} initialized successfully"
+                logging.info(success_message)
+                return agent, success_message
+            except Exception as e:
+                error_message = f"Error initializing {model_name}: {str(e)}"
+                logging.error(error_message)
+                return None, error_message
 
+        # Handle chat separately
         def handle_chat(query, model_name):
-            agent = get_or_initialize_agent(model_name)
-            if not agent:
-                return f"Could not initialize {model_name}. Please check logs for details."
-            return agent.process_query(query)
+            # First update status message
+            status = f"Processing query with {model_name}..."
+            
+            # Get or initialize agent
+            agent, init_status = get_or_initialize_agent(model_name)
+            
+            # If initialization failed
+            if agent is None:
+                return f"Could not initialize {model_name}. {init_status}", init_status
+            
+            # Process the query
+            try:
+                result = agent.process_query(query)
+                return result, f"Query processed with {model_name}"
+            except Exception as e:
+                error_msg = f"Error processing query: {str(e)}"
+                return error_msg, error_msg
 
+        # Handle numerical problem
         def handle_problem(problem, model_name):
-            agent = get_or_initialize_agent(model_name)
-            if not agent:
-                return f"Could not initialize {model_name}. Please check logs for details."
-            return agent.process_query(problem)
+            status = f"Solving problem with {model_name}..."
+            
+            # Get or initialize agent
+            agent, init_status = get_or_initialize_agent(model_name)
+            
+            # If initialization failed
+            if agent is None:
+                return f"Could not initialize {model_name}. {init_status}", init_status
+            
+            # Process the problem
+            try:
+                result = agent.process_query(problem)
+                return result, f"Problem solved with {model_name}"
+            except Exception as e:
+                error_msg = f"Error solving problem: {str(e)}"
+                return error_msg, error_msg
 
+        # Handle file upload
         def handle_file_upload(file, model_name):
-            agent = get_or_initialize_agent(model_name)
-            if not agent:
-                return f"Could not initialize {model_name}. Please check logs for details."
-            return agent.process_documents(file)
+            if file is None:
+                return "No file uploaded. Please upload a file."
+                
+            status = f"Processing document with {model_name}..."
+            
+            # Get or initialize agent
+            agent, init_status = get_or_initialize_agent(model_name)
+            
+            # If initialization failed
+            if agent is None:
+                return init_status
+            
+            # Process the document
+            try:
+                result = agent.process_documents(file)
+                return result
+            except Exception as e:
+                return f"Error processing document: {str(e)}"
 
+        # Handle document query
         def handle_query(query, model_name):
-            agent = get_or_initialize_agent(model_name)
-            if not agent:
-                return f"Could not initialize {model_name}. Please check logs for details."
-            return agent.query_documents(query)
+            status = f"Querying documents with {model_name}..."
+            
+            # Get or initialize agent
+            agent, init_status = get_or_initialize_agent(model_name)
+            
+            # If initialization failed
+            if agent is None:
+                return f"Could not initialize {model_name}. {init_status}", init_status
+            
+            # Query the documents
+            try:
+                result = agent.query_documents(query)
+                return result, f"Documents queried with {model_name}"
+            except Exception as e:
+                error_msg = f"Error querying documents: {str(e)}"
+                return error_msg, error_msg
 
-        # Set up event handlers
+        # Set up event handlers - UPDATED to include status_message updates
         chat_button.click(
             handle_chat,
             inputs=[chat_input, selected_model],
-            outputs=[chat_output]
+            outputs=[chat_output, status_message]
         )
         
         solve_button.click(
             handle_problem,
             inputs=[problem_input, selected_model],
-            outputs=[solution_output]
+            outputs=[solution_output, status_message]
         )
         
         file_upload.upload(
@@ -368,11 +432,11 @@ def create_interface():
         query_button.click(
             handle_query,
             inputs=[query_input, selected_model],
-            outputs=[query_output]
+            outputs=[query_output, status_message]
         )
 
     return demo
 
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch(share=True)  
+    demo.launch(share=True)