Update eb_agent_module.py

eb_agent_module.py

@@ -425,141 +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)
-        
-                # 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."
+    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."
+            )
+            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
                 )
-                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."}]})
+                response_text = api_response.text
                 
-                # 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()
+            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)
                 
-                # 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}"
+                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:
-                    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)}"
+                    # 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: