Update chatbot_handler.py

chatbot_handler.py

@@ -1,112 +1,166 @@
 # chatbot_handler.py
 import logging
 import json
-import aiohttp # Using asynchronous aiohttp
-# import os # No longer needed for API key from environment
+import google.generativeai as genai
+import os # For potential API key loading if Canvas injection fails for the library
 
-# Ensure GEMINI_API_KEY is set in your environment variables
-# api_key = os.getenv('GEMINI_API_KEY') # Removed: API key will be empty string, provided by Canvas
+# --- 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
+safety_settings = []
+
+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
+        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.")
+
+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
 
 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.
-    Gemini expects roles 'user' and 'model'.
-    It also filters out system messages if any, as Gemini handles system prompts differently.
+    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:
-        role = "user" if msg["role"] == "user" else "model" # Gemini uses 'model' for assistant
-        # Ensure content is a string, skip if not (e.g. if a gr.Plot was accidentally in history)
-        if isinstance(msg.get("content"), str):
-            gemini_contents.append({"role": role, "parts": [{"text": msg["content"]}]})
+        # 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: {msg.get('content')}")
+            logging.warning(f"Skipping non-string content in chat history for Gemini: {content}")
     return gemini_contents
 
-async def generate_llm_response(user_message: str, plot_id: str, plot_label: str, chat_history_for_plot: list):
+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.
+    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 latest user_message.
+                                      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.
     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.")
+        logging.error("Cannot generate LLM response: Formatted history is empty or invalid.")
         return "I'm sorry, there was an issue processing the conversation history."
 
-    payload = {
-        "contents": gemini_formatted_history,
-        "generationConfig": { 
-            "temperature": 0.7,
-            "topK": 1,
-            "topP": 1,
-            "maxOutputTokens": 2048,
-        }
-    }
-    
-    # API key is an empty string. Canvas will automatically provide it in runtime.
-    apiKey = "" 
-    # Using gemini-2.0-flash as per general instructions
-    apiUrl = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={apiKey}"
-
-    async with aiohttp.ClientSession() as session:
-        try:
-            async with session.post(apiUrl, headers={'Content-Type': 'application/json'}, json=payload, timeout=45) as resp:
-                response_json = await resp.json()
-                logging.debug(f"LLM API Raw Response for '{plot_label}': {json.dumps(response_json, indent=2)}")
-
-                if resp.status != 200:
-                    error_detail = response_json.get('error', {}).get('message', 'Unknown API error')
-                    if response_json.get("promptFeedback") and response_json["promptFeedback"].get("blockReason"):
-                         reason = response_json["promptFeedback"]["blockReason"]
-                         safety_ratings_info = [f"{rating['category']}: {rating['probability']}" for rating in response_json['promptFeedback'].get('safetyRatings', [])]
-                         details = f" Safety Ratings: {', '.join(safety_ratings_info)}" if safety_ratings_info else ""
-                         logging.warning(f"Content blocked by API (Status {resp.status}) for '{plot_label}'. Reason: {reason}.{details}")
-                         return f"I'm sorry, I can't respond to that due to content policy: {reason}."
-                    logging.error(f"LLM API Error (Status {resp.status}) for '{plot_label}': {error_detail}")
-                    return f"Sorry, the AI model returned an error (Status {resp.status}). Please try again later."
-
-                if response_json.get("candidates") and \
-                   response_json["candidates"][0].get("content") and \
-                   response_json["candidates"][0]["content"].get("parts") and \
-                   response_json["candidates"][0]["content"]["parts"][0].get("text"):
-                    response_text = response_json["candidates"][0]["content"]["parts"][0]["text"]
-                    logging.info(f"LLM generated response for '{plot_label}': {response_text[:100]}...")
-                    return response_text
-                elif response_json.get("promptFeedback") and response_json["promptFeedback"].get("blockReason"):
-                    reason = response_json["promptFeedback"]["blockReason"]
-                    safety_ratings_info = [f"{rating['category']}: {rating['probability']}" for rating in response_json['promptFeedback'].get('safetyRatings', [])]
-                    details = f" Safety Ratings: {', '.join(safety_ratings_info)}" if safety_ratings_info else ""
-                    logging.warning(f"Content blocked by API (Status 200 but no candidate) for '{plot_label}'. Reason: {reason}.{details}")
-                    return f"I'm sorry, your request was processed but could not be answered due to content policy: {reason}."
-                else:
-                    logging.error(f"Unexpected LLM API response structure for '{plot_label}': {response_json}")
-                    return "Sorry, I received an unexpected or empty response from the AI model."
-
-        except aiohttp.ClientTimeout:
-            logging.error(f"LLM API call timed out for '{plot_label}'.", exc_info=True)
-            return "Sorry, the request to the AI model timed out. Please try again."
-        except aiohttp.ClientError as e:
-            logging.error(f"Error calling LLM API (aiohttp) for '{plot_label}': {e}", exc_info=True)
-            return f"Sorry, I couldn't connect to the AI model at the moment. Network Error: {type(e).__name__}."
-        except json.JSONDecodeError as e:
-            logging.error(f"Error decoding LLM API response for '{plot_label}': {e}", exc_info=True)
-            try:
-                raw_text_response = await resp.text()
-                logging.error(f"LLM API Raw Text Response (on JSONDecodeError) for '{plot_label}': {raw_text_response}")
-            except Exception as read_err:
-                logging.error(f"Could not read raw text response: {read_err}")
-            return "Sorry, I received an unreadable response from the AI model."
-        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__}."
+    # 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.
+
+    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
+        )
+        
+        # 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__}."
+    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__}."