Update chatbot_handler.py

chatbot_handler.py

@@ -2,165 +2,83 @@
 import logging
 import json
 from google import genai
-import os # For potential API key loading if Canvas injection fails for the library
-
-# --- Gemini Configuration ---
-# Option 1: Rely on Canvas to make this work with an empty key for the library.
-# This is the preferred approach as per Canvas guidelines for 'fetch'.
-GEMINI_API_KEY = "" 
-GEMINI_API_KEY = os.getenv('GEMINI_API_KEY')
-
-# Option 2: Fallback to environment variable if direct empty key doesn't work with the library via Canvas.
-# if not GEMINI_API_KEY: # This check would be if we explicitly want to load from env
-#     GEMINI_API_KEY = os.getenv('GEMINI_API_KEY_ENV_VAR_NAME') # Replace with your actual env var name if you use one
-#     if not GEMINI_API_KEY:
-#         logging.warning("GEMINI_API_KEY not found via direct assignment or environment variable.")
-        # If you have a default key for local testing (NOT FOR PRODUCTION/CANVAS)
-        # GEMINI_API_KEY = "YOUR_LOCAL_DEV_API_KEY" 
-
-model = None
-gen_config = None
+import os
+
+# Gemini API key configuration
+GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', '')
+
+client = None
+model_name = "gemini-2.0-flash"
 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 is not None: # Check if it's set (even if empty string for Canvas)
-        genai.configure(api_key=GEMINI_API_KEY)
-        
-        # As per general instructions, use gemini-2.0-flash if not told otherwise.
-        MODEL_NAME = "gemini-2.0-flash" 
-        model = genai.GenerativeModel(MODEL_NAME)
-        
-        gen_config = genai.types.GenerationConfig(
-            temperature=0.7,
-            top_k=1, # Per user's original config
-            top_p=1, # Per user's original config
-            max_output_tokens=2048, # Per user's original config
-        )
-        
-        # Standard safety settings
+    if GEMINI_API_KEY:
+        client = genai.Client(api_key=GEMINI_API_KEY)
+
+        # 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"},
         ]
-        logging.info(f"Gemini model '{MODEL_NAME}' configured successfully.")
-    else:
-        logging.error("Gemini API Key is None. Model not configured.")
 
+        logging.info(f"Gemini client initialized with model '{model_name}'")
+    else:
+        logging.error("Gemini API Key is not set.")
 except Exception as e:
-    logging.error(f"Failed to configure Gemini or instantiate model: {e}", exc_info=True)
-    model = None # Ensure model is None if setup fails
+    logging.error(f"Failed to initialize Gemini client: {e}", exc_info=True)
+
 
 def format_history_for_gemini(gradio_chat_history: list) -> list:
-    """
-    Converts Gradio chat history (list of dicts with 'role' and 'content')
-    to Gemini API's 'contents' format (list of dicts with 'role' and 'parts').
-    Gemini SDK expects roles 'user' and 'model'.
-    """
     gemini_contents = []
     for msg in gradio_chat_history:
-        # Map Gradio 'assistant' role to Gemini 'model' role
         role = "user" if msg["role"] == "user" else "model"
-        
         content = msg.get("content")
         if isinstance(content, str):
             gemini_contents.append({"role": role, "parts": [{"text": content}]})
         else:
-            logging.warning(f"Skipping non-string content in chat history for Gemini: {content}")
+            logging.warning(f"Skipping non-string 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):
-    """
-    Generates a response from the LLM using Gemini API via the Python SDK.
-    Args:
-        user_message (str): The latest message from the user.
-        plot_id (str): The ID of the plot being discussed.
-        plot_label (str): The label of the plot being discussed.
-        chat_history_for_plot (list): The current conversation history for this plot.
-                                      This list ALREADY includes the initial assistant message (with data summary)
-                                      and the latest user_message.
-        plot_data_summary (str, optional): The textual summary of the plot data. 
-                                           While it's in the history, passing it explicitly might be useful
-                                           for future system prompt enhancements if needed.
-    Returns:
-        str: The LLM's response text.
-    """
-    logging.info(f"Generating LLM response for plot_id: {plot_id} ('{plot_label}'). User message: '{user_message}'")
-    # Log the provided data summary for debugging
-    # logging.debug(f"Data summary for '{plot_label}':\n{plot_data_summary}")
-
-
-    if not model:
-        logging.error("Gemini model not configured. Cannot generate LLM response.")
-        return "I'm sorry, the AI model is not available at the moment. (Configuration Error)"
-
-    # The chat_history_for_plot already contains the initial assistant message with the summary,
-    # and the latest user message which triggered this call.
+    if not client:
+        logging.error("Gemini client 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("Cannot generate LLM response: Formatted history is empty or invalid.")
-        return "I'm sorry, there was an issue processing the conversation history."
-
-    # Optional: Construct a system instruction if desired, though the initial message in history helps.
-    # system_instruction_text = (
-    #     f"You are an expert in Employer Branding and LinkedIn social media strategy. "
-    #     f"You are discussing the graph: '{plot_label}' (ID: '{plot_id}'). "
-    #     f"A data summary for this graph was provided in your initial message: \n---\n{plot_data_summary}\n---\n"
-    #     f"Refer to this summary and the conversation history to answer questions. "
-    #     f"If specific data is not in the summary, clearly state that the provided snapshot doesn't contain that detail."
-    # )
-    # contents_for_api = [{"role": "system", "parts": [{"text": system_instruction_text}]}] + gemini_formatted_history
-    # For now, relying on the summary being in the `gemini_formatted_history` via the first assistant message.
+        logging.error("Empty or invalid chat history.")
+        return "There was an issue processing the conversation history."
 
     try:
-        logging.debug(f"Sending to Gemini API. History: {json.dumps(gemini_formatted_history, indent=2)}")
-        
-        response = await model.generate_content_async(
-            contents=gemini_formatted_history, # History already includes user's latest message
-            generation_config=gen_config,
-            safety_settings=safety_settings
+        response = await client.models.generate_content_async(
+            model=model_name,
+            contents=gemini_formatted_history,
+            generation_config=generation_config,
+            safety_settings=safety_settings,
         )
-        
-        # logging.debug(f"LLM API Raw Response object for '{plot_label}': {response}")
 
-        # Check for blocking based on prompt_feedback first (as per SDK examples)
         if response.prompt_feedback and response.prompt_feedback.block_reason:
-            reason = response.prompt_feedback.block_reason.name # e.g., 'SAFETY'
-            # safety_ratings_info = [f"{rating.category.name}: {rating.probability.name}" for rating in response.prompt_feedback.safety_ratings]
-            # details = f" Safety Ratings: {', '.join(safety_ratings_info)}" if safety_ratings_info else ""
-            logging.warning(f"Content blocked by API (prompt_feedback) for '{plot_label}'. Reason: {reason}.")
-            return f"I'm sorry, your request was blocked due to content policy: {reason}."
-
-        # Accessing response text (handle multi-part if any)
-        if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
-            response_text = "".join(part.text for part in response.candidates[0].content.parts)
-            logging.info(f"LLM generated response for '{plot_label}': {response_text[:150]}...")
-            return response_text
-        else:
-            # This case might occur if the response was empty but not blocked by prompt_feedback
-            # (e.g. finish_reason other than SAFETY, or no candidates)
-            finish_reason_str = "UNKNOWN"
-            if response.candidates and response.candidates[0].finish_reason:
-                 finish_reason_str = response.candidates[0].finish_reason.name # e.g. 'STOP', 'MAX_TOKENS', 'SAFETY', 'RECITATION', 'OTHER'
-            
-            if finish_reason_str == 'SAFETY': # Content blocked at candidate level
-                 logging.warning(f"Content blocked by API (candidate safety) for '{plot_label}'. Finish Reason: {finish_reason_str}.")
-                 return f"I'm sorry, I can't provide a response due to safety filters regarding: {finish_reason_str}."
-
-            logging.error(f"Unexpected LLM API response structure or empty content for '{plot_label}'. Finish Reason: {finish_reason_str}. Full response: {response}")
-            return f"Sorry, I received an unexpected or empty response from the AI model (Finish Reason: {finish_reason_str})."
-
-    except google.api_core.exceptions.PermissionDenied as e:
-        logging.error(f"LLM API Permission Denied (Status 403) for '{plot_label}': {e}", exc_info=True)
-        return "Sorry, there's an issue with API permissions. Please ensure the API key is correct and the service is enabled. (Error 403)"
-    except google.api_core.exceptions.InvalidArgument as e:
-        logging.error(f"LLM API Invalid Argument (Status 400) for '{plot_label}': {e}. History: {json.dumps(gemini_formatted_history, indent=2)}", exc_info=True)
-        return "Sorry, there was an issue with the request sent to the AI model (e.g. malformed history). (Error 400)"
-    except google.api_core.exceptions.GoogleAPIError as e: # Catch other Google API errors
-        logging.error(f"Google API Error during LLM call for '{plot_label}': {e}", exc_info=True)
-        return f"An API error occurred while trying to get an AI response: {type(e).__name__}."
+            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}."
+
     except Exception as e:
-        logging.error(f"Generic error during LLM call for '{plot_label}': {e}", exc_info=True)
-        return f"An unexpected error occurred while trying to get an AI response: {type(e).__name__}."
+        logging.error(f"Error generating response for plot '{plot_label}': {e}", exc_info=True)
+        return f"Unexpected error occurred: {type(e).__name__}."