Enhance error handling and logging in various modules

- Improved exception handling in app.py to provide more specific error messages.
- Added timeout and error handling for API requests in extract_results.py.
- Implemented thread safety for Solr client access in ui.py.
- Enhanced error messages in connections.py and data_processing.py for better debugging.

app.py

@@ -29,8 +29,11 @@ def main():
     demo = create_ui(llm_model, solr_client)
     try:
         demo.queue().launch(debug=True, share=True, allowed_paths=['/tmp/plots'])
-    except Exception as e:
+    except (IOError, OSError) as e:
         print(f"An error occurred while launching the Gradio app: {e}")
+        print("Please check if the port is already in use or if you have the necessary permissions.")
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
     finally:
         # Ensure the SSH tunnel is closed when the app is shut down
         print("\nClosing SSH tunnel...")

connections.py

@@ -23,6 +23,7 @@ def initialize_connections():
     ssh_tunnel_server = None
     try:
         # 1. Configure and start the SSH Tunnel
+        print("Attempting to start SSH tunnel...")
         ssh_tunnel_server = SSHTunnelForwarder(
             (config.SSH_HOST, config.SSH_PORT),
             ssh_username=config.SSH_USER,
@@ -36,19 +37,36 @@ def initialize_connections():
         # 2. Initialize the pysolr client
         solr_url = f'http://127.0.0.1:{ssh_tunnel_server.local_bind_port}/solr/{config.SOLR_CORE_NAME}'
         solr_client = pysolr.Solr(solr_url, auth=(config.SOLR_USER, config.SOLR_PASS), always_commit=True)
-        solr_client.ping()
+        solr_client.ping() # Verify connection
         print(f"✅ Solr connection successful on core '{config.SOLR_CORE_NAME}'.")
 
         # 3. Initialize the LLM
+        if not config.GEMINI_API_KEY:
+            print("❌ CRITICAL: GEMINI_API_KEY is not set. LLM will not be available.")
+            raise ValueError("GEMINI_API_KEY is missing.")
+        
         genai.configure(api_key=config.GEMINI_API_KEY)
-        llm_model = genai.GenerativeModel('gemini-2.5-flash', generation_config=genai.types.GenerationConfig(temperature=0))
+        llm_model = genai.GenerativeModel('gemini-1.5-flash', generation_config=genai.types.GenerationConfig(temperature=0))
         print(f"✅ LLM Model '{llm_model.model_name}' initialized.")
 
         print("✅ System Initialized Successfully.")
         return ssh_tunnel_server, solr_client, llm_model
 
+    except pysolr.SolrError as e:
+        print(f"\n❌ Solr Error: {e}")
+        print("Please check if the Solr core '{config.SOLR_CORE_NAME}' exists and is running.")
+        if ssh_tunnel_server and ssh_tunnel_server.is_active:
+            ssh_tunnel_server.stop()
+        return None, None, None
+    
+    except ValueError as e:
+        print(f"\n❌ Configuration Error: {e}")
+        if ssh_tunnel_server and ssh_tunnel_server.is_active:
+            ssh_tunnel_server.stop()
+        return None, None, None
+
     except Exception as e:
-        print(f"\n❌ An error occurred during setup: {e}")
+        print(f"\n❌ An unexpected error occurred during setup: {e}")
         if ssh_tunnel_server and ssh_tunnel_server.is_active:
             ssh_tunnel_server.stop()
         return None, None, None

data_processing.py

@@ -1,3 +1,4 @@
+
 """
 Core data processing and analysis logic for the PharmaCircle AI Data Analyst.
 
@@ -9,7 +10,6 @@ This module orchestrates the main analysis workflow:
 5.  Synthesizes the findings into a comprehensive, user-facing report.
 """
 
-
 import json
 import re
 import datetime
@@ -21,6 +21,7 @@ import concurrent.futures
 import copy
 import google.generativeai as genai
 import urllib
+import pysolr
 
 from llm_prompts import (
     get_analysis_plan_prompt,
@@ -29,7 +30,6 @@ from llm_prompts import (
 )
 from extract_results import get_search_list_params
 
-
 def parse_suggestions_from_report(report_text):
     """Extracts numbered suggestions from the report's markdown text."""
     suggestions_match = re.search(r"### (?:Deeper Dive: Suggested Follow-up Analyses|Suggestions for Further Exploration)\s*\n(.*?)$", report_text, re.DOTALL | re.IGNORECASE)
@@ -38,58 +38,52 @@ def parse_suggestions_from_report(report_text):
     suggestions = re.findall(r"^\s*\d+\.\s*(.*)", suggestions_text, re.MULTILINE)
     return [s.strip() for s in suggestions]
 
-
 def llm_generate_analysis_plan_with_history(llm_model, natural_language_query, chat_history):
-        """
-        Generates a complete analysis plan from a user query, considering chat history
-        and dynamic field suggestions from an external API.
-        """
-        search_fields, search_name, field_mappings = [], "", {}
-        try:
-            # Call the external API to get dynamic fields, core name, and mappings
-            search_fields, search_name, field_mappings = get_search_list_params(natural_language_query)
-            print(f"API returned core: '{search_name}' with {len(search_fields)} fields and {len(field_mappings)} mappings.")
-        except Exception as e:
-            print(f"Warning: Could not retrieve dynamic search fields. Proceeding without them. Error: {e}")
-
-        # Determine the core name, default to 'news' if not provided by the API
-        core_name = search_name if search_name else 'news'
-
-        # Apply the field mappings to the suggestions before sending them to the LLM
-        mapped_search_fields = []
-        if search_fields and field_mappings:
-            for field in search_fields:
-                original_name = field.get('field_name')
-                # Create a new dict to avoid modifying the original
-                mapped_field = field.copy()
-                if original_name in field_mappings:
-                    mapped_field['field_name'] = field_mappings[original_name]
-                    print(f"Mapped field '{original_name}' to '{mapped_field['field_name']}'")
-                mapped_search_fields.append(mapped_field)
-        else:
-            mapped_search_fields = search_fields
-
+    """
+    Generates a complete analysis plan from a user query, considering chat history
+    and dynamic field suggestions from an external API.
+    """
+    search_fields, search_name, field_mappings = [], "", {}
+    try:
+        search_fields, search_name, field_mappings = get_search_list_params(natural_language_query)
+        print(f"API returned core: '{search_name}' with {len(search_fields)} fields and {len(field_mappings)} mappings.")
+    except Exception as e:
+        print(f"Warning: Could not retrieve dynamic search fields. Proceeding without them. Error: {e}")
 
-        # Generate the prompt, passing the mapped fields and the dynamic core name
-        prompt = get_analysis_plan_prompt(natural_language_query, chat_history, mapped_search_fields, core_name)
+    core_name = search_name if search_name else 'news'
 
-        try:
-            response = llm_model.generate_content(prompt)
-            cleaned_text = re.sub(r'```json\s*|\s*```', '', response.text, flags=re.MULTILINE | re.DOTALL).strip()
-            plan = json.loads(cleaned_text)
-            # Return the plan, the mapped fields for UI display, and the core name
-            return plan, mapped_search_fields, core_name
-        except Exception as e:
-            raw_response_text = response.text if 'response' in locals() else 'N/A'
-            print(f"Error in llm_generate_analysis_plan_with_history: {e}\nRaw Response:\n{raw_response_text}")
-            # Return None for the plan but still return other data for debugging
-            return None, mapped_search_fields, core_name
+    mapped_search_fields = []
+    if search_fields and field_mappings:
+        for field in search_fields:
+            original_name = field.get('field_name')
+            mapped_field = field.copy()
+            if original_name in field_mappings:
+                mapped_field['field_name'] = field_mappings[original_name]
+                print(f"Mapped field '{original_name}' to '{mapped_field['field_name']}'")
+            mapped_search_fields.append(mapped_field)
+    else:
+        mapped_search_fields = search_fields
 
+    prompt = get_analysis_plan_prompt(natural_language_query, chat_history, mapped_search_fields, core_name)
 
+    try:
+        response = llm_model.generate_content(prompt)
+        cleaned_text = re.sub(r'```json\s*|\s*```', '', response.text, flags=re.MULTILINE | re.DOTALL).strip()
+        plan = json.loads(cleaned_text)
+        return plan, mapped_search_fields, core_name
+    except json.JSONDecodeError as e:
+        raw_response_text = response.text if 'response' in locals() else 'N/A'
+        print(f"Error decoding JSON from LLM response: {e}\nRaw Response:\n{raw_response_text}")
+        return None, mapped_search_fields, core_name
+    except Exception as e:
+        raw_response_text = response.text if 'response' in locals() else 'N/A'
+        print(f"Error in llm_generate_analysis_plan_with_history: {e}\nRaw Response:\n{raw_response_text}")
+        return None, mapped_search_fields, core_name
 
 def execute_quantitative_query(solr_client, plan):
     """Executes the facet query to get aggregate data."""
     if not plan or 'quantitative_request' not in plan or 'json.facet' not in plan.get('quantitative_request', {}):
+        print("Skipping quantitative query due to incomplete plan.")
         return None, None
     try:
         params = {
@@ -97,38 +91,43 @@ def execute_quantitative_query(solr_client, plan):
             "rows": 0,
             "json.facet": json.dumps(plan['quantitative_request']['json.facet'])
         }
-        # Build the full Solr URL manually (for logging) from the client's current URL
         base_url = f"{solr_client.url}/select"
         query_string = urllib.parse.urlencode(params)
         full_url = f"{base_url}?{query_string}"
         print(f"[DEBUG] Solr QUANTITATIVE query URL: {full_url}")
         results = solr_client.search(**params)
         return results.raw_response.get("facets", {}), full_url
+    except pysolr.SolrError as e:
+        print(f"Solr Error in quantitative query on core {solr_client.url}: {e}")
+        return None, None
     except Exception as e:
-        print(f"Error in quantitative query on core specified in client ({solr_client.url}): {e}")
+        print(f"Unexpected error in quantitative query: {e}")
         return None, None
 
 def execute_qualitative_query(solr_client, plan):
     """Executes the grouping query to get the best example docs."""
     if not plan or 'qualitative_request' not in plan:
+        print("Skipping qualitative query due to incomplete plan.")
         return None, None
     try:
         qual_request = copy.deepcopy(plan['qualitative_request'])
         params = {
             "q": plan.get('query_filter', '*_*'),
-            "rows": 5,  # Get a few examples per group
+            "rows": 5,
             "fl": "*,score",
             **qual_request
         }
-        # Build the full Solr URL manually (for logging) from the client's current URL
         base_url = f"{solr_client.url}/select"
         query_string = urllib.parse.urlencode(params)
         full_url = f"{base_url}?{query_string}"
         print(f"[DEBUG] Solr QUALITATIVE query URL: {full_url}")
         results = solr_client.search(**params)
         return results.grouped, full_url
+    except pysolr.SolrError as e:
+        print(f"Solr Error in qualitative query on core {solr_client.url}: {e}")
+        return None, None
     except Exception as e:
-        print(f"Error in qualitative query on core specified in client ({solr_client.url}): {e}")
+        print(f"Unexpected error in qualitative query: {e}")
         return None, None
 
 def llm_synthesize_enriched_report_stream(llm_model, query, quantitative_data, qualitative_data, plan):
@@ -143,7 +142,7 @@ def llm_synthesize_enriched_report_stream(llm_model, query, quantitative_data, q
             yield chunk.text
     except Exception as e:
         print(f"Error in llm_synthesize_enriched_report_stream: {e}")
-        yield "Sorry, I was unable to generate a report for this data."
+        yield "Sorry, an error occurred while generating the report. Please check the logs for details."
 
 def llm_generate_visualization_code(llm_model, query_context, facet_data):
     """Generates Python code for visualization based on query and data."""
@@ -154,23 +153,45 @@ def llm_generate_visualization_code(llm_model, query_context, facet_data):
         code = re.sub(r'^```python\s*|```$', '', response.text, flags=re.MULTILINE)
         return code
     except Exception as e:
-        print(f"Error in llm_generate_visualization_code: {e}\nRaw response: {response.text}")
+        raw_response_text = response.text if 'response' in locals() else 'N/A'
+        print(f"Error in llm_generate_visualization_code: {e}\nRaw response: {raw_response_text}")
         return None
 
 def execute_viz_code_and_get_path(viz_code, facet_data):
     """Executes visualization code and returns the path to the saved plot image."""
-    if not viz_code: return None
+    if not viz_code:
+        return None
+
+    # --- SECURITY WARNING ---
+    # The following code executes code generated by an LLM. This is a security
+    # risk and should be handled with extreme care in a production environment.
+    # Ideally, this code should be run in a sandboxed environment.
+    print("\n--- WARNING: Executing LLM-generated code. ---")
+
     try:
-        if not os.path.exists('/tmp/plots'): os.makedirs('/tmp/plots')
+        if not os.path.exists('/tmp/plots'):
+            os.makedirs('/tmp/plots')
         plot_path = f"/tmp/plots/plot_{datetime.datetime.now().timestamp()}.png"
+
+        # Create a restricted global environment for execution
         exec_globals = {'facet_data': facet_data, 'plt': plt, 'sns': sns, 'pd': pd}
+
+
         exec(viz_code, exec_globals)
+
         fig = exec_globals.get('fig')
         if fig:
             fig.savefig(plot_path, bbox_inches='tight')
             plt.close(fig)
+            print("--- LLM-generated code executed successfully. ---")
             return plot_path
-        return None
+        else:
+            print("--- LLM-generated code did not produce a 'fig' object. ---")
+            return None
+
     except Exception as e:
-        print(f"ERROR executing visualization code: {e}\n---Code---\n{viz_code}")
-        return None
+        print(f"\n--- ERROR executing visualization code: ---")
+        print(f"Error: {e}")
+        print(f"--- Code---\n{viz_code}")
+        print("-----------------------------------------")
+        return None

extract_results.py

@@ -26,42 +26,50 @@ def get_search_list_params(query, k=20):
     """
     url = "https://aitest.ebalina.com/stream"
     
-    response = requests.post(
-        url,
-        headers={'Content-Type': 'application/json'},
-        json={"query": query, "k": k},
-        stream=True
-    )
-    
-    search_fields = []
-    search_name = ""
-    field_mappings_str = ""
+    try:
+        response = requests.post(
+            url,
+            headers={'Content-Type': 'application/json'},
+            json={"query": query, "k": k},
+            stream=True,
+            timeout=30  # Add a 30-second timeout
+        )
+        response.raise_for_status()  # Raise an exception for bad status codes
+        
+        search_fields = []
+        search_name = ""
+        field_mappings_str = ""
 
-    for line in response.iter_lines():
-        if line and line.startswith(b'data: '):
-            try:
-                line_str = line.decode('utf-8')[6:]
-                if not line_str or line_str.isspace():
-                    continue
-                
-                data = json.loads(line_str)
-                log_title = data.get('log_title')
+        for line in response.iter_lines():
+            if line and line.startswith(b'data: '):
+                try:
+                    line_str = line.decode('utf-8')[6:]
+                    if not line_str or line_str.isspace():
+                        continue
+                    
+                    data = json.loads(line_str)
+                    log_title = data.get('log_title')
 
-                if log_title == 'Search List Result':
-                    content = data.get('content', '')
-                    if content:
-                        yaml_data = yaml.safe_load(content)
-                        print("DEBUG:", yaml_data)
-                        # This is the dynamic core name
-                        search_name = yaml_data.get('search_name', '')
-                        search_fields = yaml_data.get('search_fields', [])
+                    if log_title == 'Search List Result':
+                        content = data.get('content', '')
+                        if content:
+                            yaml_data = yaml.safe_load(content)
+                            print("DEBUG:", yaml_data)
+                            # This is the dynamic core name
+                            search_name = yaml_data.get('search_name', '')
+                            search_fields = yaml_data.get('search_fields', [])
 
-                elif log_title == 'Field Mapping Outputs':
-                    field_mappings_str = data.get('content', '')
+                    elif log_title == 'Field Mapping Outputs':
+                        field_mappings_str = data.get('content', '')
 
-            except (json.JSONDecodeError, yaml.YAMLError, AttributeError):
-                continue
-    
-    field_mappings = _parse_mappings(field_mappings_str)
-    
-    return search_fields, search_name, field_mappings
+                except (json.JSONDecodeError, yaml.YAMLError, AttributeError) as e:
+                    print(f"Error parsing stream line: {e}\nLine: {line_str}")
+                    continue
+        
+        field_mappings = _parse_mappings(field_mappings_str)
+        
+        return search_fields, search_name, field_mappings
+
+    except requests.exceptions.RequestException as e:
+        print(f"Error connecting to the external API: {e}")
+        return [], "", {}

ui.py

@@ -11,6 +11,7 @@ by calling functions from the data_processing module.
 import gradio as gr
 import json
 import concurrent.futures
+import threading
 from data_processing import (
     llm_generate_analysis_plan_with_history,
     execute_quantitative_query,
@@ -21,6 +22,9 @@ from data_processing import (
     parse_suggestions_from_report
 )
 
+# Create a lock to protect the Solr client from concurrent access
+solr_lock = threading.Lock()
+
 
 def create_ui(llm_model, solr_client):
     """
@@ -121,28 +125,27 @@ def create_ui(llm_model, solr_client):
             yield (history, state, None, None, gr.update(value=formatted_plan, visible=True), None, None, None, None, suggestions_display_update)
 
             # --- DYNAMIC CORE SWITCH (Thread-safe) ---
-            original_solr_url = solr_client.url
-            # Correctly construct the new URL by replacing the last component (the core name)
-            base_url = original_solr_url.rsplit('/', 1)[0]
-            new_url = f"{base_url}/{core_name}"
-            solr_client.url = new_url
-            print(f"[INFO] Switched Solr client to core: {core_name} at URL: {solr_client.url}")
-            # ---
-
-            # Execute queries in parallel
-            aggregate_data, quantitative_url = None, None
-            example_data, qualitative_url = None, None
-            try:
-                with concurrent.futures.ThreadPoolExecutor() as executor:
-                    future_agg = executor.submit(execute_quantitative_query, solr_client, analysis_plan)
-                    future_ex = executor.submit(execute_qualitative_query, solr_client, analysis_plan)
-                    aggregate_data, quantitative_url = future_agg.result()
-                    example_data, qualitative_url = future_ex.result()
-            finally:
-                # --- IMPORTANT: Reset client to default URL ---
-                solr_client.url = original_solr_url
-                print(f"[INFO] Reset Solr client to default URL: {original_solr_url}")
-                # ---
+            with solr_lock:
+                original_solr_url = solr_client.url
+                # Correctly construct the new URL by replacing the last component (the core name)
+                base_url = original_solr_url.rsplit('/', 1)[0]
+                new_url = f"{base_url}/{core_name}"
+                solr_client.url = new_url
+                print(f"[INFO] Switched Solr client to core: {core_name} at URL: {solr_client.url}")
+
+                # Execute queries in parallel
+                aggregate_data, quantitative_url = None, None
+                example_data, qualitative_url = None, None
+                try:
+                    with concurrent.futures.ThreadPoolExecutor() as executor:
+                        future_agg = executor.submit(execute_quantitative_query, solr_client, analysis_plan)
+                        future_ex = executor.submit(execute_qualitative_query, solr_client, analysis_plan)
+                        aggregate_data, quantitative_url = future_agg.result()
+                        example_data, qualitative_url = future_ex.result()
+                finally:
+                    # --- IMPORTANT: Reset client to default URL ---
+                    solr_client.url = original_solr_url
+                    print(f"[INFO] Reset Solr client to default URL: {original_solr_url}")
 
             if not aggregate_data or aggregate_data.get('count', 0) == 0:
                 history.append((None, f"No data was found for your query in the '{core_name}' core. Please try a different question."))