Update eb_agent_module.py

eb_agent_module.py

@@ -20,14 +20,16 @@ except ImportError:
         SafetySetting = None
         # Define HarmCategory and HarmBlockThreshold as inner classes or attributes for the dummy types
         class HarmCategory: # type: ignore
+            HARM_CATEGORY_UNSPECIFIED = "HARM_CATEGORY_UNSPECIFIED"
             HARM_CATEGORY_HATE_SPEECH = "HARM_CATEGORY_HATE_SPEECH"
             HARM_CATEGORY_HARASSMENT = "HARM_CATEGORY_HARASSMENT"
             HARM_CATEGORY_SEXUALLY_EXPLICIT = "HARM_CATEGORY_SEXUALLY_EXPLICIT"
             HARM_CATEGORY_DANGEROUS_CONTENT = "HARM_CATEGORY_DANGEROUS_CONTENT"
         class HarmBlockThreshold: # type: ignore
+            BLOCK_NONE = "BLOCK_NONE"
             BLOCK_LOW_AND_ABOVE = "BLOCK_LOW_AND_ABOVE"
             BLOCK_MEDIUM_AND_ABOVE = "BLOCK_MEDIUM_AND_ABOVE"
-            BLOCK_NONE = "BLOCK_NONE"
+            BLOCK_ONLY_HIGH = "BLOCK_ONLY_HIGH" # Added for completeness, adjust if needed
 
 # --- Custom Exceptions ---
 class ValidationError(Exception):
@@ -54,17 +56,27 @@ GENERATION_CONFIG_PARAMS = {
     "candidate_count": 1,
 }
 
-# Updated to use types.HarmCategory and types.HarmBlockThreshold
-DEFAULT_SAFETY_SETTINGS = [
-    {"category": types.HarmCategory.HARM_CATEGORY_HATE_SPEECH if types and hasattr(types, 'HarmCategory') else "HARM_CATEGORY_HATE_SPEECH",
-     "threshold": types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE if types and hasattr(types, 'HarmBlockThreshold') else "BLOCK_MEDIUM_AND_ABOVE"},
-    {"category": types.HarmCategory.HARM_CATEGORY_HARASSMENT if types and hasattr(types, 'HarmCategory') else "HARM_CATEGORY_HARASSMENT",
-     "threshold": types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE if types and hasattr(types, 'HarmBlockThreshold') else "BLOCK_MEDIUM_AND_ABOVE"},
-    {"category": types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT if types and hasattr(types, 'HarmCategory') else "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-     "threshold": types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE if types and hasattr(types, 'HarmBlockThreshold') else "BLOCK_MEDIUM_AND_ABOVE"},
-    {"category": types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT if types and hasattr(types, 'HarmCategory') else "HARM_CATEGORY_DANGEROUS_CONTENT",
-     "threshold": types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE if types and hasattr(types, 'HarmBlockThreshold') else "BLOCK_MEDIUM_AND_ABOVE"},
-]
+# Corrected to use direct enum members when available
+DEFAULT_SAFETY_SETTINGS = []
+if types and hasattr(types, 'HarmCategory') and hasattr(types, 'HarmBlockThreshold'):
+    DEFAULT_SAFETY_SETTINGS = [
+        {"category": types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+         "threshold": types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE},
+        {"category": types.HarmCategory.HARM_CATEGORY_HARASSMENT,
+         "threshold": types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE},
+        {"category": types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+         "threshold": types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE},
+        {"category": types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+         "threshold": types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE},
+    ]
+else: # Fallback to strings if types or enums are not properly imported
+    logging.warning("Falling back to string representations for DEFAULT_SAFETY_SETTINGS due to missing types.HarmCategory or types.HarmBlockThreshold.")
+    DEFAULT_SAFETY_SETTINGS = [
+        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+        {"category": "HARM_CATEGORY_HARASSMENT", "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"},
+    ]
 
 
 df_rag_documents = pd.DataFrame({
@@ -101,7 +113,6 @@ class AdvancedRAGSystem:
         if not text or not isinstance(text, str):
             raise ValueError("Cannot embed empty or non-string text.")
         
-        # Ensure types.EmbedContentConfig is available before using it
         embed_config = None
         if types and hasattr(types, 'EmbedContentConfig'):
             embed_config = types.EmbedContentConfig(task_type="SEMANTIC_SIMILARITY")
@@ -148,7 +159,6 @@ class AdvancedRAGSystem:
                 self.embeddings = np.array([])
 
     def _calculate_cosine_similarity(self, embeddings_matrix: np.ndarray, query_vector: np.ndarray) -> np.ndarray:
-        """Calculate normalized cosine similarity between a matrix of embeddings and a query vector."""
         query_vector = query_vector.flatten()
         norm_matrix = np.linalg.norm(embeddings_matrix, axis=1, keepdims=True)
         normalized_embeddings_matrix = embeddings_matrix / (norm_matrix + 1e-8)
@@ -179,21 +189,15 @@ class AdvancedRAGSystem:
         
         try:
             similarity_scores = self._calculate_cosine_similarity(self.embeddings, query_vector)
-            
-            if similarity_scores.size == 0:
-                return ""
-
+            if similarity_scores.size == 0: return ""
             relevant_indices_after_threshold = np.where(similarity_scores >= min_similarity)[0]
             if len(relevant_indices_after_threshold) == 0:
                 logging.debug(f"No documents met the minimum similarity threshold of {min_similarity} for query: {query[:50]}")
                 return ""
-            
             relevant_scores = similarity_scores[relevant_indices_after_threshold]
             sorted_relevant_indices_local = np.argsort(relevant_scores)[::-1]
             top_original_indices = relevant_indices_after_threshold[sorted_relevant_indices_local[:top_k]]
-
             if len(top_original_indices) == 0: return ""
-
             context_parts = [self.documents_df.iloc[i]['text'] for i in top_original_indices if 'text' in self.documents_df.columns]
             context = "\n\n---\n\n".join(context_parts)
             logging.debug(f"Retrieved RAG context for query '{str(query)[:50]}...':\n{context[:200]}...")
@@ -210,15 +214,14 @@ class EmployerBrandingAgent:
                  llm_model_name: str,
                  embedding_model_name: str,
                  generation_config_dict: dict,
-                 safety_settings_list_of_dicts: list,
+                 safety_settings_list_of_dicts: list, # This list now contains dicts with ENUM values or STRINGS
                  force_sandbox: bool = False):
         self.all_dataframes = {k: df.copy() for k, df in all_dataframes.items()}
         self.schemas_representation = self._get_enhanced_schemas_representation()
-        
         self.chat_history = []
         self.llm_model_name = llm_model_name
         self.generation_config_dict = generation_config_dict
-        self.safety_settings_list_of_dicts = safety_settings_list_of_dicts # These are dicts
+        self.safety_settings_list_of_dicts = safety_settings_list_of_dicts
         self.embedding_model_name = embedding_model_name
         self.rag_system = AdvancedRAGSystem(rag_documents_df, self.embedding_model_name)
         self.force_sandbox = force_sandbox
@@ -236,7 +239,7 @@ class EmployerBrandingAgent:
         return "N/A"
 
     def _calculate_growth_rate(self, df: pd.DataFrame) -> str:
-        logging.debug("_calculate_growth_rate is a placeholder.") # Changed to debug
+        logging.debug("_calculate_growth_rate is a placeholder.")
         return "Growth rate calculation not implemented."
     def _analyze_engagement_trends(self, df: pd.DataFrame) -> str:
         logging.debug("_analyze_engagement_trends is a placeholder.")
@@ -257,21 +260,12 @@ class EmployerBrandingAgent:
     def _calculate_key_metrics(self, df: pd.DataFrame, df_type: str) -> dict:
         metrics = {}
         if 'follower' in df_type.lower():
-            metrics.update({
-                'follower_growth_rate': self._calculate_growth_rate(df),
-                'engagement_trends': self._analyze_engagement_trends(df),
-                'demographic_distribution': self._analyze_demographics(df)
-            })
+            metrics.update({'follower_growth_rate': self._calculate_growth_rate(df), 'engagement_trends': self._analyze_engagement_trends(df), 'demographic_distribution': self._analyze_demographics(df)})
         elif 'post' in df_type.lower():
-            metrics.update({
-                'post_performance': self._analyze_post_performance(df),
-                'content_themes': self._extract_content_themes(df),
-                'optimal_posting_times': self._find_optimal_times(df)
-            })
+            metrics.update({'post_performance': self._analyze_post_performance(df), 'content_themes': self._extract_content_themes(df), 'optimal_posting_times': self._find_optimal_times(df)})
         elif 'mention' in df_type.lower():
             metrics['mention_volume_trend'] = "Mention volume trend not implemented."
             metrics['mention_sentiment_overview'] = "Mention sentiment overview not implemented."
-        
         if not metrics:
             logging.debug(f"No specific key metrics defined for df_type: {df_type}")
             return {"info": "Standard metrics applicable."}
@@ -283,7 +277,7 @@ class EmployerBrandingAgent:
                 try:
                     max_date = df[col].max()
                     if pd.notna(max_date):
-                        days_diff = (datetime.now(max_date.tzinfo) - max_date).days # tz aware
+                        days_diff = (datetime.now(max_date.tzinfo if max_date.tzinfo else None) - max_date).days
                         return f"Data up to {max_date.strftime('%Y-%m-%d')} ({days_diff} days old)"
                 except Exception: pass
         return "Freshness N/A (no clear date column)"
@@ -296,28 +290,17 @@ class EmployerBrandingAgent:
 
     def _assess_data_quality(self, df: pd.DataFrame) -> dict:
         completeness = (1 - (df.isnull().sum().sum() / (len(df) * len(df.columns)))) if len(df) > 0 and len(df.columns) > 0 else 0
-        return {
-            'completeness_score': f"{completeness:.2%}",
-            'freshness_info': self._calculate_data_freshness(df),
-            'consistency_check': self._check_data_consistency(df),
-            'accuracy_flags_summary': self._identify_accuracy_issues(df),
-            'sample_size_notes': f"{len(df)} records. {'Adequate for basic analysis.' if len(df) >= 100 else 'Limited sample size; insights may be indicative.'}"
-        }
+        return {'completeness_score': f"{completeness:.2%}", 'freshness_info': self._calculate_data_freshness(df), 'consistency_check': self._check_data_consistency(df), 'accuracy_flags_summary': self._identify_accuracy_issues(df), 'sample_size_notes': f"{len(df)} records. {'Adequate for basic analysis.' if len(df) >= 100 else 'Limited sample size; insights may be indicative.'}"}
 
     def _identify_patterns(self, df: pd.DataFrame, key: str) -> str:
         logging.debug(f"_identify_patterns for {key} is a placeholder.")
         return "Pattern identification not implemented."
 
     def _format_df_analysis(self, df_key: str, analysis: dict) -> str:
-        formatted_parts = [f"\n--- DataFrame: df_{df_key} ---"]
-        formatted_parts.append(f"  Shape: {analysis['shape']}")
-        formatted_parts.append(f"  Date Range: {analysis['date_range']}")
-        formatted_parts.append("  Key Metrics:")
-        for metric, value in analysis['key_metrics'].items():
-            formatted_parts.append(f"    - {metric.replace('_', ' ').title()}: {value}")
+        formatted_parts = [f"\n--- DataFrame: df_{df_key} ---", f"  Shape: {analysis['shape']}", f"  Date Range: {analysis['date_range']}", "  Key Metrics:"]
+        for metric, value in analysis['key_metrics'].items(): formatted_parts.append(f"    - {metric.replace('_', ' ').title()}: {value}")
         formatted_parts.append("  Data Quality Assessment:")
-        for aspect, value in analysis['data_quality'].items():
-            formatted_parts.append(f"    - {aspect.replace('_', ' ').title()}: {value}")
+        for aspect, value in analysis['data_quality'].items(): formatted_parts.append(f"    - {aspect.replace('_', ' ').title()}: {value}")
         formatted_parts.append(f"  Notable Patterns: {analysis['notable_patterns']}")
         return "\n".join(formatted_parts)
 
@@ -330,13 +313,7 @@ class EmployerBrandingAgent:
             if df.empty:
                 schema_descriptions.append(f"\n--- DataFrame: df_{key} ---\nStatus: Empty. No analysis possible.")
                 continue
-            analysis = {
-                'shape': df.shape,
-                'date_range': self._get_date_range(df),
-                'key_metrics': self._calculate_key_metrics(df, key),
-                'data_quality': self._assess_data_quality(df),
-                'notable_patterns': self._identify_patterns(df, key)
-            }
+            analysis = {'shape': df.shape, 'date_range': self._get_date_range(df), 'key_metrics': self._calculate_key_metrics(df, key), 'data_quality': self._assess_data_quality(df), 'notable_patterns': self._identify_patterns(df, key)}
             schema_descriptions.append(self._format_df_analysis(key, analysis))
         return "\n".join(schema_descriptions)
     
@@ -367,108 +344,48 @@ Relevant Benchmarks: {benchmarks_val}"""
         return enhanced_context
 
     async def _build_prompt_for_current_turn(self, raw_user_query: str) -> str:
-        prompt_parts = [
-            "You are an expert Employer Branding Analyst and a helpful AI assistant. "
-            "Your goal is to provide insightful analysis based on the provided LinkedIn data. "
-            "When asked to generate Pandas code, ensure it is correct, runnable, and clearly explained. "
-            "When providing insights, be specific and refer to the data where possible. "
-            "Use the detailed data overview and any contextual information provided."
-        ]
-        prompt_parts.append("\n\n--- DETAILED DATA OVERVIEW ---")
-        prompt_parts.append(self.schemas_representation)
-
+        prompt_parts = ["You are an expert Employer Branding Analyst...", "--- DETAILED DATA OVERVIEW ---", self.schemas_representation] # Truncated for brevity
         if self.rag_system.embeddings is not None and self.rag_system.embeddings.size > 0:
-            logging.debug(f"Retrieving base RAG context for query: {raw_user_query[:50]}...")
             base_rag_context = await self.rag_system.retrieve_relevant_info(raw_user_query)
             if base_rag_context:
-                logging.debug(f"Enhancing RAG context for query: {raw_user_query[:50]}...")
                 enhanced_rag_context = await self._enhance_rag_context(raw_user_query, base_rag_context)
-                prompt_parts.append("\n\n--- RELEVANT CONTEXTUAL INFORMATION (from documents & business knowledge) ---")
-                prompt_parts.append(enhanced_rag_context)
-            else: logging.debug("No base RAG context found.")
-        else: logging.debug("RAG system not initialized or embeddings not available, skipping RAG context retrieval.")
-
-        prompt_parts.append("\n\n--- USER REQUEST ---")
-        prompt_parts.append(f"Based on all the information above, please respond to the following user query:\n{raw_user_query}")
+                prompt_parts.extend(["--- RELEVANT CONTEXTUAL INFORMATION (from documents & business knowledge) ---", enhanced_rag_context])
+        prompt_parts.extend(["--- USER REQUEST ---", f"Based on all the information above, please respond to the following user query:\n{raw_user_query}"])
         final_prompt = "\n".join(prompt_parts)
         logging.debug(f"Built prompt for current turn (first 300 chars): {final_prompt[:300]}")
         return final_prompt
 
     async def _process_structured_query(self, prompt: str) -> dict:
-        logging.debug("_process_structured_query is a placeholder. Returning dummy structure.")
-        return {
-            "Key Findings": ["Placeholder finding 1", "Placeholder finding 2"],
-            "Performance Metrics": ["Placeholder metric performance"],
-            "Actionable Recommendations": {
-                "Immediate Actions (0-30 days)": ["Placeholder immediate action"],
-                "Short-term Strategy (1-3 months)": ["Placeholder short-term strategy"],
-                "Long-term Vision (3-12 months)": ["Placeholder long-term vision"]
-            },
-            "Risk Assessment": ["Placeholder risk"],
-            "Success Metrics to Track": ["Placeholder KPI"]
-        }
+        logging.debug("_process_structured_query is a placeholder.")
+        return {"Key Findings": ["Placeholder finding 1"], "Performance Metrics": ["Placeholder metric"], "Actionable Recommendations": {"Immediate Actions (0-30 days)": ["Placeholder action"]}, "Risk Assessment": ["Placeholder risk"], "Success Metrics to Track": ["Placeholder KPI"]}
 
     async def _generate_hr_insights(self, query: str, context: str) -> str:
-        insight_prompt = f"""
-As an expert HR analytics consultant, analyze the following LinkedIn employer branding data:
-{context}
-User Query: {query}
-Please provide insights in this structured format:
-## Key Findings
-- [3-5 bullet points of main discoveries]
-## Performance Metrics
-- Current performance vs industry benchmarks
-- Trend analysis (improving/declining/stable)
-## Actionable Recommendations
-1. **Immediate Actions (0-30 days)**
-  - [Specific, measurable actions]
-2. **Short-term Strategy (1-3 months)**
-  - [Strategic initiatives]
-3. **Long-term Vision (3-12 months)**
-  - [Comprehensive improvements]
-## Risk Assessment
-- Potential challenges or red flags
-- Mitigation strategies
-## Success Metrics to Track
-- KPIs to monitor progress
-- Reporting frequency recommendations
-"""
+        insight_prompt = f"As an expert HR analytics consultant...\n{context}\nUser Query: {query}\nPlease provide insights in this structured format:\n## Key Findings\n- ...\n..." # Truncated
         if not client: return "Error: AI client not configured for generating HR insights."
         api_call_contents = [{"role": "user", "parts": [{"text": insight_prompt}]}]
         
-        # Construct SafetySetting objects if types.SafetySetting is available
         api_safety_settings_objects = []
         if types and hasattr(types, 'SafetySetting'):
             for ss_dict in self.safety_settings_list_of_dicts:
                 try:
-                    # Use types.HarmCategory and types.HarmBlockThreshold directly
-                    category = getattr(types.HarmCategory, ss_dict['category'].split('.')[-1] if isinstance(ss_dict['category'], str) else ss_dict['category'].name, types.HarmCategory.HARM_CATEGORY_UNSPECIFIED)
-                    threshold = getattr(types.HarmBlockThreshold, ss_dict['threshold'].split('.')[-1] if isinstance(ss_dict['threshold'], str) else ss_dict['threshold'].name, types.HarmBlockThreshold.BLOCK_NONE)
-                    api_safety_settings_objects.append(types.SafetySetting(category=category, threshold=threshold))
-                except Exception as e_ss:
-                    logging.warning(f"Could not create SafetySetting object from {ss_dict}: {e_ss}. Using dict.")
-                    api_safety_settings_objects.append(ss_dict) # Fallback to dict if creation fails
-        else: # Fallback if types.SafetySetting is not available
+                    # Directly use the category and threshold from the dict,
+                    # which should be enum members if types was available at DEFAULT_SAFETY_SETTINGS definition,
+                    # or strings otherwise.
+                    api_safety_settings_objects.append(types.SafetySetting(category=ss_dict['category'], threshold=ss_dict['threshold']))
+                except Exception as e_ss: # Catch if ss_dict values are not valid for SafetySetting
+                    logging.warning(f"Could not create SafetySetting object from {ss_dict} for HR insights: {e_ss}. Using raw dict.")
+                    api_safety_settings_objects.append(ss_dict) 
+        else:
             api_safety_settings_objects = self.safety_settings_list_of_dicts
 
         api_generation_config_obj = None
         if types and hasattr(types, 'GenerateContentConfig'):
-            api_generation_config_obj = types.GenerateContentConfig(
-                **self.generation_config_dict,
-                safety_settings=api_safety_settings_objects
-            )
-        else: # Fallback if types.GenerateContentConfig is not available
-            config_dict_for_api = {**self.generation_config_dict, "safety_settings": api_safety_settings_objects}
-            api_generation_config_obj = config_dict_for_api
-
+            api_generation_config_obj = types.GenerateContentConfig(**self.generation_config_dict, safety_settings=api_safety_settings_objects)
+        else:
+            api_generation_config_obj = {**self.generation_config_dict, "safety_settings": api_safety_settings_objects}
 
         try:
-            response = await asyncio.to_thread(
-                client.models.generate_content,
-                model=self.llm_model_name,
-                contents=api_call_contents,
-                config=api_generation_config_obj
-            )
+            response = await asyncio.to_thread(client.models.generate_content, model=self.llm_model_name, contents=api_call_contents, config=api_generation_config_obj)
             if not response.candidates: return "HR insights generation failed: No response from AI."
             return response.text.strip()
         except Exception as e:
@@ -476,29 +393,23 @@ Please provide insights in this structured format:
             return f"Error generating HR insights: {str(e)}"
 
     def _validate_query(self, query: str) -> bool:
-        if not query or len(query.strip()) < 3:
-            logging.warning(f"Query too short: '{query}'")
-            return False
+        if not query or len(query.strip()) < 3: logging.warning(f"Query too short: '{query}'"); return False
         hr_keywords = ['employee', 'talent', 'hiring', 'culture', 'brand', 'engagement', 'retention', 'follower', 'post', 'mention', 'linkedin']
-        if not any(keyword in query.lower() for keyword in hr_keywords):
-            logging.warning(f"Query may not be HR/LinkedIn-relevant: {query[:50]}")
+        if not any(keyword in query.lower() for keyword in hr_keywords): logging.warning(f"Query may not be HR/LinkedIn-relevant: {query[:50]}")
         return True
 
     def _get_query_help_message(self) -> str:
-        return ("I'm here to help with Employer Branding analysis on LinkedIn data. "
-                "Please ask specific questions about your followers, posts, or mentions. "
-                "For example: 'What are the top industries of my followers?' or 'Analyze the engagement trend of my recent posts.'")
+        return "I'm here to help with Employer Branding analysis... Example: 'What are the top industries of my followers?'"
 
     async def _check_system_readiness(self) -> dict:
         logging.debug("_check_system_readiness is a placeholder.")
         if not client: return {'ready': False, 'reason': 'AI Client not initialized.'}
-        if self.rag_system.embeddings is None:
-             logging.warning("RAG embeddings not yet initialized. Proceeding, but RAG context will be unavailable.")
+        if self.rag_system.embeddings is None: logging.warning("RAG embeddings not yet initialized.")
         return {'ready': True, 'reason': 'System appears ready.'}
 
     def _get_fallback_response(self, query: str) -> str:
         logging.error(f"Executing fallback response for query: {query[:50]}")
-        return "I encountered an unexpected issue while processing your request. Please try rephrasing your query or try again later."
+        return "I encountered an unexpected issue..."
 
     async def _core_query_processing(self, raw_user_query_this_turn: str) -> str:
         augmented_current_user_prompt_text = await self._build_prompt_for_current_turn(raw_user_query_this_turn)
@@ -510,69 +421,40 @@ Please provide insights in this structured format:
         if types and hasattr(types, 'SafetySetting'):
             for ss_dict in self.safety_settings_list_of_dicts:
                 try:
-                    category_enum_val = ss_dict['category']
-                    threshold_enum_val = ss_dict['threshold']
-                    # If they are already enum members, use them directly
-                    if not isinstance(category_enum_val, str): # Assumes it's an enum member
-                        category = category_enum_val
-                    else: # If string, try to get from types.HarmCategory
-                         category = getattr(types.HarmCategory, category_enum_val.split('.')[-1], types.HarmCategory.HARM_CATEGORY_UNSPECIFIED)
-
-                    if not isinstance(threshold_enum_val, str): # Assumes it's an enum member
-                        threshold = threshold_enum_val
-                    else: # If string, try to get from types.HarmBlockThreshold
-                        threshold = getattr(types.HarmBlockThreshold, threshold_enum_val.split('.')[-1], types.HarmBlockThreshold.BLOCK_NONE)
-                    
-                    api_safety_settings_objects.append(types.SafetySetting(category=category, threshold=threshold))
+                    # Directly use category/threshold from ss_dict. They should be enums or valid strings.
+                    api_safety_settings_objects.append(types.SafetySetting(category=ss_dict['category'], threshold=ss_dict['threshold']))
                 except Exception as e_ss_core:
-                    logging.warning(f"Could not create SafetySetting object from {ss_dict} in core: {e_ss_core}. Using dict.")
-                    api_safety_settings_objects.append(ss_dict) # Fallback
-        else:
+                    logging.warning(f"Could not create SafetySetting object from {ss_dict} in core: {e_ss_core}. Using raw dict.")
+                    api_safety_settings_objects.append(ss_dict) # Fallback to passing the dict itself
+        else: # Fallback if types.SafetySetting is not available
             api_safety_settings_objects = self.safety_settings_list_of_dicts
 
-
         api_generation_config_obj = None
         if types and hasattr(types, 'GenerateContentConfig'):
-             api_generation_config_obj = types.GenerateContentConfig(
-                **self.generation_config_dict,
-                safety_settings=api_safety_settings_objects
-            )
+             api_generation_config_obj = types.GenerateContentConfig(**self.generation_config_dict, safety_settings=api_safety_settings_objects)
         else:
             logging.error("GenerateContentConfig type not available. API call might fail.")
-            config_dict_for_api = {**self.generation_config_dict, "safety_settings": api_safety_settings_objects}
-            api_generation_config_obj = config_dict_for_api
-
-        response = await asyncio.to_thread(
-            client.models.generate_content,
-            model=self.llm_model_name,
-            contents=api_call_contents,
-            config=api_generation_config_obj
-        )
+            api_generation_config_obj = {**self.generation_config_dict, "safety_settings": api_safety_settings_objects}
 
+        response = await asyncio.to_thread(client.models.generate_content, model=self.llm_model_name, contents=api_call_contents, config=api_generation_config_obj)
         if not response.candidates:
             block_reason = response.prompt_feedback.block_reason if response.prompt_feedback else "Unknown"
             block_message = response.prompt_feedback.block_reason_message if response.prompt_feedback else ""
-            logging.warning(f"AI response blocked or empty. Reason: {block_reason}, Msg: {block_message}")
-            error_message = f"The AI's response was blocked. Reason: {block_reason}."
-            if block_message: error_message += f" Details: {block_message}"
+            error_message = f"The AI's response was blocked. Reason: {block_reason}." + (f" Details: {block_message}" if block_message else "")
             return error_message
         return response.text.strip()
 
     async def _process_query_with_timeout(self, raw_user_query_this_turn: str, timeout_seconds: int = 60) -> str:
-        try:
-            return await asyncio.wait_for(self._core_query_processing(raw_user_query_this_turn), timeout=timeout_seconds)
+        try: return await asyncio.wait_for(self._core_query_processing(raw_user_query_this_turn), timeout=timeout_seconds)
         except asyncio.TimeoutError:
-            logging.error(f"Query processing timed out after {timeout_seconds} seconds for query: {raw_user_query_this_turn[:50]}")
-            return "I'm sorry, but your request took too long to process. Please try a simpler query or try again later."
+            logging.error(f"Query processing timed out for {timeout_seconds} seconds...")
+            return "I'm sorry, but your request took too long..."
 
     async def process_query(self, raw_user_query_this_turn: str) -> str:
-        if not client:
-            logging.error("GenAI client not initialized. Cannot process query.")
-            return "Error: The AI Agent is not available due to a configuration issue with the AI service."
+        if not client: return "Error: The AI Agent is not available..."
         if not self._validate_query(raw_user_query_this_turn): return self._get_query_help_message()
         readiness_check = await self._check_system_readiness()
         if not readiness_check['ready']: return f"System not ready: {readiness_check['reason']}"
-
         max_retries = 2
         for attempt in range(max_retries + 1):
             try:
@@ -581,22 +463,15 @@ Please provide insights in this structured format:
                 logging.info(f"Successfully received AI response (attempt {attempt+1}): {response_text[:100]}")
                 return response_text
             except RateLimitError as rle:
-                logging.warning(f"Rate limit encountered on attempt {attempt + 1}: {rle}. Retrying after backoff...")
-                if attempt == max_retries:
-                    logging.error("Max retries reached due to rate limiting.")
-                    return "The AI service is currently busy. Please try again in a few moments."
+                if attempt == max_retries: return "The AI service is currently busy..."
                 await asyncio.sleep(2 ** attempt)
-            except ValidationError as ve:
-                logging.warning(f"Validation error during processing: {ve}")
-                return f"Query validation failed: {str(ve)}"
+            except ValidationError as ve: return f"Query validation failed: {str(ve)}"
             except Exception as e:
-                logging.error(f"Error during GenAI call on attempt {attempt + 1}: {e}", exc_info=True)
-                if attempt == max_retries:
-                    logging.error("Max retries reached due to general errors.")
-                    return self._get_fallback_response(raw_user_query_this_turn)
+                if attempt == max_retries: return self._get_fallback_response(raw_user_query_this_turn)
         return self._get_fallback_response(raw_user_query_this_turn)
 
     def _classify_query_type(self, query: str) -> str:
+        # ... (implementation unchanged)
         query_lower = query.lower()
         if any(word in query_lower for word in ['trend', 'growth', 'change', 'time']): return 'trend_analysis'
         elif any(word in query_lower for word in ['compare', 'benchmark', 'versus']): return 'comparative_analysis'
@@ -605,11 +480,13 @@ Please provide insights in this structured format:
         elif any(word in query_lower for word in ['what is', 'explain', 'define']): return 'definition_explanation'
         else: return 'general_inquiry'
 
+
     def clear_chat_history(self):
         self.chat_history = []
         logging.info("EmployerBrandingAgent chat history cleared by request.")
 
 def get_all_schemas_representation(all_dataframes: dict) -> str:
+    # ... (implementation unchanged)
     if not all_dataframes: return "No DataFrames are currently loaded."
     schema_descriptions = ["DataFrames currently available in the application state:"]
     for key, df in all_dataframes.items():
@@ -624,7 +501,9 @@ def get_all_schemas_representation(all_dataframes: dict) -> str:
         schema_descriptions.append(schema)
     return "\n".join(schema_descriptions)
 
+
 async def test_rag_retrieval_accuracy():
+    # ... (implementation unchanged, ensure client and types are checked if used here)
     logging.info("Running RAG retrieval accuracy test...")
     test_embedding_model = GEMINI_EMBEDDING_MODEL_NAME 
     if not client: