Update chatbot_handler.py

chatbot_handler.py

@@ -1,84 +1,141 @@
 # chatbot_handler.py
 import logging
 import json
-from google import genai
+from google import genai # Assuming this is the correct SDK
 import os
+import asyncio # Added for asyncio.to_thread
 
 # Gemini API key configuration
 GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', '')
 
 client = None
-model_name = "gemini-2.0-flash"
+# model_name = "gemini-1.0-pro" # Or your preferred model like "gemini-2.0-flash"
+model_name = "gemini-1.5-flash-latest" # Using a more recent Flash model
 safety_settings = []
 
-generation_config = genai.types.GenerationConfig(
-    temperature=0.7,
-    top_k=1,
-    top_p=1,
-    max_output_tokens=2048,
-)
 
-try:
-    if GEMINI_API_KEY:
-        client = genai.Client(api_key=GEMINI_API_KEY)
+generation_config = { # New SDK style
+    "temperature": 0.7,
+    "top_p": 1,
+    "top_k": 1,
+    "max_output_tokens": 2048,
+}
 
-        # Optional: safety settings
-        safety_settings = [
-            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-            {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-            {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-        ]
+# Define safety settings list to be used by both client types
+common_safety_settings = [
+    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+]
 
-        logging.info(f"Gemini client initialized with model '{model_name}'")
+try:
+    if GEMINI_API_KEY:
+        if hasattr(genai, 'Client'): # Check for older SDK structure
+             client = genai.Client(api_key=GEMINI_API_KEY)
+             logging.info(f"Gemini client (genai.Client) initialized with model '{model_name}' for older SDK structure.")
+        else: # Fallback to current recommended practice (genai.GenerativeModel)
+             genai.configure(api_key=GEMINI_API_KEY)
+             client = genai.GenerativeModel(
+                 model_name=model_name,
+                 safety_settings=common_safety_settings,
+                 generation_config=generation_config
+             )
+             logging.info(f"Gemini client (genai.GenerativeModel) initialized with model '{model_name}'")
     else:
         logging.error("Gemini API Key is not set.")
 except Exception as e:
-    logging.error(f"Failed to initialize Gemini client: {e}", exc_info=True)
+    logging.error(f"Failed to initialize Gemini client/model: {e}", exc_info=True)
 
 
 def format_history_for_gemini(gradio_chat_history: list) -> list:
+    """Converts Gradio chat history to Gemini content format."""
     gemini_contents = []
     for msg in gradio_chat_history:
-        role = "user" if msg["role"] == "user" else "model"
+        role = "user" if msg.get("role") == "user" else "model"
         content = msg.get("content")
         if isinstance(content, str):
             gemini_contents.append({"role": role, "parts": [{"text": content}]})
+        elif isinstance(content, list) and len(content) > 0 and isinstance(content[0], dict) and "type" in content[0]:
+            parts = []
+            for part_item in content:
+                if part_item.get("type") == "text":
+                    parts.append({"text": part_item.get("text", "")})
+            if parts:
+                 gemini_contents.append({"role": role, "parts": parts})
+            else:
+                logging.warning(f"Skipping complex but empty content part in chat history: {content}")
         else:
-            logging.warning(f"Skipping non-string content in chat history: {content}")
+            logging.warning(f"Skipping non-string/non-standard content in chat history: {content}")
     return gemini_contents
 
 
 async def generate_llm_response(user_message: str, plot_id: str, plot_label: str, chat_history_for_plot: list, plot_data_summary: str = None):
     if not client:
-        logging.error("Gemini client not initialized.")
+        logging.error("Gemini client/model not initialized.")
         return "The AI model is not available. Configuration error."
 
     gemini_formatted_history = format_history_for_gemini(chat_history_for_plot)
 
     if not gemini_formatted_history:
-        logging.error("Empty or invalid chat history.")
-        return "There was an issue processing the conversation history."
+        if not any(part.get("text", "").strip() for message in gemini_formatted_history for part in message.get("parts",[])):
+             logging.error("Formatted history for Gemini is empty or contains no text.")
+             return "There was an issue processing the conversation history for the AI model (empty text)."
 
     try:
-        response = await client.models.generate_content_async(
-            model=model_name,
-            contents=gemini_formatted_history,
-            generation_config=generation_config,
-            safety_settings=safety_settings,
-        )
-
-        if response.prompt_feedback and response.prompt_feedback.block_reason:
-            reason = response.prompt_feedback.block_reason.name
-            logging.warning(f"Blocked by prompt feedback: {reason}")
-            return f"Blocked due to content policy: {reason}."
-
-        if response.candidates and response.candidates[0].content.parts:
-            return "".join(part.text for part in response.candidates[0].content.parts)
-
-        finish_reason = response.candidates[0].finish_reason.name if response.candidates and response.candidates[0].finish_reason else "UNKNOWN"
-        return f"Unexpected response. Finish reason: {finish_reason}."
-
+        response = None
+        if isinstance(client, genai.GenerativeModel):
+            logging.debug("Using genai.GenerativeModel.generate_content_async")
+            response = await client.generate_content_async(
+                contents=gemini_formatted_history
+            )
+        elif hasattr(client, 'models') and hasattr(client.models, 'generate_content'): # Check for the synchronous method
+             logging.debug("Using genai.Client.models.generate_content (synchronous via asyncio.to_thread)")
+             qualified_model_name = model_name if model_name.startswith("models/") else f"models/{model_name}"
+             
+             # Ensure safety_settings and generation_config are passed correctly
+             # to the synchronous method if it's part of this older client structure.
+             # The `client.models.generate_content` might take these as direct args.
+             response = await asyncio.to_thread(
+                 client.models.generate_content, # The synchronous function
+                 model=qualified_model_name,
+                 contents=gemini_formatted_history,
+                 generation_config=generation_config, # Pass the dict directly
+                 safety_settings=common_safety_settings # Pass the list of dicts
+             )
+        else:
+            logging.error(f"Gemini client is not a recognized type for generating content. Type: {type(client)}")
+            return "AI model interaction error (client type)."
+
+        if hasattr(response, 'prompt_feedback') and response.prompt_feedback and response.prompt_feedback.block_reason:
+            reason = response.prompt_feedback.block_reason
+            reason_name = getattr(reason, 'name', str(reason))
+            logging.warning(f"Blocked by prompt feedback: {reason_name}")
+            return f"Blocked due to content policy: {reason_name}."
+
+        if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
+            return "".join(part.text for part in response.candidates[0].content.parts if hasattr(part, 'text'))
+        
+        finish_reason = "UNKNOWN"
+        if response.candidates and response.candidates[0].finish_reason:
+            finish_reason_val = response.candidates[0].finish_reason
+            finish_reason = getattr(finish_reason_val, 'name', str(finish_reason_val))
+            
+        if not (response.candidates and response.candidates[0].content and response.candidates[0].content.parts):
+            logging.warning(f"No content parts in response. Finish reason: {finish_reason}")
+            if finish_reason == "SAFETY":
+                 return f"Response generation stopped due to safety reasons. Finish reason: {finish_reason}."
+            return f"The AI model returned an empty response. Finish reason: {finish_reason}."
+
+        return f"Unexpected response structure from AI model. Finish reason: {finish_reason}."
+
+    except AttributeError as ae:
+        logging.error(f"AttributeError during Gemini call for plot '{plot_label}': {ae}", exc_info=True)
+        if "generate_content_async" in str(ae) or "generate_content" in str(ae):
+            return f"AI model error: SDK method not found or mismatch. Details: {ae}"
+        return f"AI model error (Attribute): {type(ae).__name__} - {ae}."
     except Exception as e:
         logging.error(f"Error generating response for plot '{plot_label}': {e}", exc_info=True)
-        return f"Unexpected error occurred: {type(e).__name__}."
+        if "API key not valid" in str(e):
+            return "AI model error: API key is not valid. Please check configuration."
+        return f"An unexpected error occurred while contacting the AI model: {type(e).__name__}."