Update eb_agent_module.py

eb_agent_module.py

@@ -15,6 +15,7 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 try:
     from google import genai
     from google.genai import types # Assuming this provides necessary types like SafetySetting, HarmCategory etc.
+    from google.genai import errors
     # If GenerationConfig or EmbedContentConfig are from a different submodule, adjust imports.
     # For google-generativeai, GenerationConfig is often passed as a dict or genai.types.GenerationConfig
     # and EmbedContentConfig might be implicit or part of task_type.
@@ -424,97 +425,141 @@ class EmployerBrandingAgent:
         If a query is ambiguous or requires data not present, ask for clarification or state the limitation.
         """).strip()
 
-    async def _generate_response(self, current_user_query: str) -> str:
-        """
-        Generates a response from the LLM based on the current query, system prompts,
-        data summaries, RAG context, and the agent's chat history.
-        Assumes self.chat_history is already populated by app.py and includes the current_user_query as the last entry.
-        """
-        if not self.is_ready:
-            return "Agent is not ready. Please initialize."
-        if not client and not self.llm_model_instance:
-            return "Error: AI service is not available. Check API configuration."
-
-        try:
-            system_prompt_text = self._build_system_prompt()
-            data_summary_text = self._get_dataframes_summary()
-            rag_context_text = await self.rag_system.retrieve_relevant_info(current_user_query, top_k=2, min_similarity=0.25) # Fine-tuned RAG params
-
-            # Construct the messages for the LLM API call
-            # The history (self.chat_history) is set by app.py and includes the current user query.
-            llm_messages = []
-
-            # 1. System-level instructions and context (as a first "user" turn)
-            initial_context_prompt = (
-                f"{system_prompt_text}\n\n"
-                f"## Available Data Overview:\n{data_summary_text}\n\n"
-                f"## Relevant Background Information (if any):\n{rag_context_text if rag_context_text else 'No specific background information retrieved for this query.'}\n\n"
-                f"Given this context, please respond to the user queries that follow in the chat history."
-            )
-            llm_messages.append({"role": "user", "parts": [{"text": initial_context_prompt}]})
-            # 2. Priming assistant message
-            llm_messages.append({"role": "model", "parts": [{"text": "Understood. I have reviewed the context and data overview. I am ready to assist with your Employer Branding analysis based on our conversation."}]})
-            
-            # 3. Append the actual conversation history (already includes the current user query)
-            for entry in self.chat_history: # self.chat_history is set by app.py
-                llm_messages.append({"role": entry["role"], "parts": [{"text": entry["content"]}]})
-
-            # Prepare generation config and safety settings for the API
-            gen_config_payload = self.generation_config_dict
-            safety_settings_payload = self.safety_settings_list # Already formatted if types.SafetySetting used
-
-            if GENAI_AVAILABLE and hasattr(types, 'GenerationConfig') and not isinstance(self.generation_config_dict, types.GenerationConfig):
-                 try:
-                    gen_config_payload = types.GenerationConfig(**self.generation_config_dict)
-                 except Exception as e:
-                    logging.warning(f"Could not convert gen_config_dict to types.GenerationConfig: {e}")
-
-
-            # --- Make the API call ---
-            response_text = ""
-            if self.llm_model_instance: # Standard google-generativeai usage
-                logging.debug(f"Using google-generativeai.GenerativeModel.generate_content_async for LLM call. History length: {len(llm_messages)}")
-                api_response = await self.llm_model_instance.generate_content_async(
-                    contents=llm_messages,
-                    generation_config=gen_config_payload,
-                    safety_settings=safety_settings_payload
-                )
-                response_text = api_response.text # Simplification, assumes single part text response
-            elif client: # User's original client.models.generate_content structure
-                logging.debug(f"Using client.models.generate_content for LLM call. History length: {len(llm_messages)}")
-                # This call needs to be async or wrapped, asyncio.to_thread is used as in original
-                model_path = f"models/{self.llm_model_name}" if not self.llm_model_name.startswith("models/") else self.llm_model_name
-                api_response = await asyncio.to_thread(
-                    client.models.generate_content,
-                    model=model_path,
-                    contents=llm_messages,
-                    generation_config=gen_config_payload, # Ensure this is the correct type for client.models
-                    safety_settings=safety_settings_payload # Ensure this is the correct type
+        async def _generate_response(self, current_user_query: str) -> str:
+            """
+            Generates a response from the LLM based on the current query, system prompts,
+            data summaries, RAG context, and the agent's chat history.
+            Assumes self.chat_history is already populated by app.py and includes the current_user_query as the last entry.
+            """
+            if not self.is_ready:
+                return "Agent is not ready. Please initialize."
+            if not client and not self.llm_model_instance:
+                return "Error: AI service is not available. Check API configuration."
+        
+            try:
+                system_prompt_text = self._build_system_prompt()
+                data_summary_text = self._get_dataframes_summary()
+                rag_context_text = await self.rag_system.retrieve_relevant_info(current_user_query, top_k=2, min_similarity=0.25)
+        
+                # Construct the messages for the LLM API call
+                llm_messages = []
+        
+                # 1. System-level instructions and context (as a first "user" turn)
+                initial_context_prompt = (
+                    f"{system_prompt_text}\n\n"
+                    f"## Available Data Overview:\n{data_summary_text}\n\n"
+                    f"## Relevant Background Information (if any):\n{rag_context_text if rag_context_text else 'No specific background information retrieved for this query.'}\n\n"
+                    f"Given this context, please respond to the user queries that follow in the chat history."
                 )
-                # Parse response from client.models structure
-                if api_response.candidates and api_response.candidates[0].content and api_response.candidates[0].content.parts:
-                    response_text_parts = [part.text for part in api_response.candidates[0].content.parts if hasattr(part, 'text')]
-                    response_text = "".join(response_text_parts).strip()
-                else: # Handle blocked or empty responses from client.models
-                    if hasattr(api_response, 'prompt_feedback') and api_response.prompt_feedback and api_response.prompt_feedback.block_reason:
-                        logging.warning(f"Prompt blocked by client.models: {api_response.prompt_feedback.block_reason}")
-                        return f"I'm sorry, your request was blocked. Reason: {api_response.prompt_feedback.block_reason_message or api_response.prompt_feedback.block_reason}"
-                    if api_response.candidates and api_response.candidates[0].finish_reason != types.Candidate.FinishReason.STOP: # Assuming types.Candidate.FinishReason.STOP is valid
-                        logging.warning(f"Content generation stopped by client.models due to: {api_response.candidates[0].finish_reason}. Safety: {api_response.candidates[0].safety_ratings if hasattr(api_response.candidates[0], 'safety_ratings') else 'N/A'}")
-                        return f"I couldn't complete the response. Reason: {api_response.candidates[0].finish_reason}. Please try rephrasing."
-                    return "I apologize, but I couldn't generate a response from client.models."
-
-            else:
-                raise ConnectionError("No valid LLM client or model instance available.")
-
-            return response_text.strip()
-
-        except types.BlockedPromptException as bpe: # <--- CORRECTED LINE
-            logging.error(f"BlockedPromptException from LLM: {bpe}", exc_info=True) # Keep exc_info=True for full traceback in logs
-            return f"I'm sorry, your request was blocked by the safety filter. Please rephrase your query. Details: {bpe}"
-        except Exception as e:
-            logging.error(f"Error in _generate_response: {e}", exc_info=True) # Keep exc_info=True
-            return f"I encountered an error while processing your request: {type(e).__name__} - {str(e)}"
+                llm_messages.append({"role": "user", "parts": [{"text": initial_context_prompt}]})
+                # 2. Priming assistant message
+                llm_messages.append({"role": "model", "parts": [{"text": "Understood. I have reviewed the context and data overview. I am ready to assist with your Employer Branding analysis based on our conversation."}]})
+                
+                # 3. Append the actual conversation history (already includes the current user query)
+                for entry in self.chat_history:
+                    llm_messages.append({"role": entry["role"], "parts": [{"text": entry["content"]}]})
+        
+                # --- Make the API call ---
+                response_text = ""
+                if self.llm_model_instance:  # Standard google-generativeai usage
+                    logging.debug(f"Using google-generativeai.GenerativeModel.generate_content_async for LLM call. History length: {len(llm_messages)}")
+                    
+                    # Prepare generation config and safety settings for google-generativeai
+                    gen_config_payload = self.generation_config_dict
+                    safety_settings_payload = self.safety_settings_list
+                    
+                    if GENAI_AVAILABLE and hasattr(types, 'GenerationConfig') and not isinstance(self.generation_config_dict, types.GenerationConfig):
+                        try:
+                            gen_config_payload = types.GenerationConfig(**self.generation_config_dict)
+                        except Exception as e:
+                            logging.warning(f"Could not convert gen_config_dict to types.GenerationConfig: {e}")
+                    
+                    api_response = await self.llm_model_instance.generate_content_async(
+                        contents=llm_messages,
+                        generation_config=gen_config_payload,
+                        safety_settings=safety_settings_payload
+                    )
+                    response_text = api_response.text
+                    
+                elif client:  # google.genai client usage
+                    logging.debug(f"Using client.models.generate_content for LLM call. History length: {len(llm_messages)}")
+                    
+                    # Convert messages to the format expected by google.genai client
+                    # The client expects a simpler contents format
+                    contents = []
+                    for msg in llm_messages:
+                        if msg["role"] == "user":
+                            contents.append(msg["parts"][0]["text"])
+                        elif msg["role"] == "model":
+                            # For model responses, we might need to handle differently
+                            # but for now, let's include them as context
+                            contents.append(f"Assistant: {msg['parts'][0]['text']}")
+                    
+                    # Create the config object with both generation config and safety settings
+                    config_dict = {}
+                    
+                    # Add generation config parameters
+                    if self.generation_config_dict:
+                        for key, value in self.generation_config_dict.items():
+                            config_dict[key] = value
+                    
+                    # Add safety settings
+                    if self.safety_settings_list:
+                        # Convert safety settings to the correct format if needed
+                        safety_settings = []
+                        for ss in self.safety_settings_list:
+                            if isinstance(ss, dict):
+                                # Convert dict to types.SafetySetting
+                                safety_settings.append(types.SafetySetting(
+                                    category=ss.get('category'),
+                                    threshold=ss.get('threshold')
+                                ))
+                            else:
+                                safety_settings.append(ss)
+                        config_dict['safety_settings'] = safety_settings
+                    
+                    # Create the config object
+                    config = types.GenerateContentConfig(**config_dict)
+                    
+                    model_path = f"models/{self.llm_model_name}" if not self.llm_model_name.startswith("models/") else self.llm_model_name
+                    
+                    api_response = await asyncio.to_thread(
+                        client.models.generate_content,
+                        model=model_path,
+                        contents=contents,  # Simplified contents format
+                        config=config       # Using config parameter instead of separate generation_config and safety_settings
+                    )
+                    
+                    # Parse response from client.models structure
+                    if api_response.candidates and api_response.candidates[0].content and api_response.candidates[0].content.parts:
+                        response_text_parts = [part.text for part in api_response.candidates[0].content.parts if hasattr(part, 'text')]
+                        response_text = "".join(response_text_parts).strip()
+                    else:
+                        # Handle blocked or empty responses
+                        if hasattr(api_response, 'prompt_feedback') and api_response.prompt_feedback and api_response.prompt_feedback.block_reason:
+                            logging.warning(f"Prompt blocked by client.models: {api_response.prompt_feedback.block_reason}")
+                            return f"I'm sorry, your request was blocked. Reason: {api_response.prompt_feedback.block_reason_message or api_response.prompt_feedback.block_reason}"
+                        if api_response.candidates and hasattr(api_response.candidates[0], 'finish_reason'):
+                            finish_reason = api_response.candidates[0].finish_reason
+                            if hasattr(types.Candidate, 'FinishReason') and finish_reason != types.Candidate.FinishReason.STOP:
+                                logging.warning(f"Content generation stopped by client.models due to: {finish_reason}. Safety: {getattr(api_response.candidates[0], 'safety_ratings', 'N/A')}")
+                                return f"I couldn't complete the response. Reason: {finish_reason}. Please try rephrasing."
+                        return "I apologize, but I couldn't generate a response from client.models."
+                else:
+                    raise ConnectionError("No valid LLM client or model instance available.")
+        
+                return response_text.strip()
+        
+            except Exception as e:
+                error_message = str(e).lower()
+                
+                # Check if it's a blocked prompt error by examining the error message
+                if any(keyword in error_message for keyword in ['blocked', 'safety', 'filter', 'prohibited']):
+                    logging.error(f"Blocked prompt from LLM: {e}", exc_info=True)
+                    return f"I'm sorry, your request was blocked by the safety filter. Please rephrase your query. Details: {e}"
+                else:
+                    logging.error(f"Error in _generate_response: {e}", exc_info=True)
+                    return f"I encountered an error while processing your request: {type(e).__name__} - {str(e)}"
 
 
     def _validate_query(self, query: str) -> bool: