GuglielmoTor commited on
Commit
25e22b8
·
verified ·
1 Parent(s): e7eee4a

Create zzz_eb_agent_w_rag.py

Browse files
Files changed (1) hide show
  1. zzz_eb_agent_w_rag.py +1120 -0
zzz_eb_agent_w_rag.py ADDED
@@ -0,0 +1,1120 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pandas as pd
2
+ import json
3
+ import os
4
+ import asyncio
5
+ import logging
6
+ import numpy as np
7
+ import textwrap
8
+ from datetime import datetime
9
+ from typing import Dict, List, Optional, Union, Any, Tuple
10
+ import traceback
11
+ import pandasai as pai
12
+ from pandasai_litellm import LiteLLM
13
+
14
+ # Add this early, before matplotlib.pyplot is imported directly or by pandasai
15
+ import matplotlib
16
+ matplotlib.use('Agg') # Use a non-interactive backend for Matplotlib
17
+ import matplotlib.pyplot as plt
18
+
19
+ # Configure logging
20
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(message)s')
21
+
22
+ try:
23
+ from google import genai
24
+ from google.genai import types
25
+ from google.genai import errors
26
+ GENAI_AVAILABLE = True
27
+ logging.info("Google Generative AI library imported successfully.")
28
+ except ImportError:
29
+ logging.warning("Google Generative AI library not found. Please install it: pip install google-generativeai")
30
+ GENAI_AVAILABLE = False
31
+
32
+ # Dummy classes for graceful degradation
33
+ class genai:
34
+ Client = None
35
+
36
+ class types:
37
+ EmbedContentConfig = None
38
+ GenerationConfig = None
39
+ SafetySetting = None
40
+ Candidate = type('Candidate', (), {'FinishReason': type('FinishReason', (), {'STOP': 'STOP'})})
41
+
42
+ class HarmCategory:
43
+ HARM_CATEGORY_UNSPECIFIED = "HARM_CATEGORY_UNSPECIFIED"
44
+ HARM_CATEGORY_HATE_SPEECH = "HARM_CATEGORY_HATE_SPEECH"
45
+ HARM_CATEGORY_HARASSMENT = "HARM_CATEGORY_HARASSMENT"
46
+ HARM_CATEGORY_SEXUALLY_EXPLICIT = "HARM_CATEGORY_SEXUALLY_EXPLICIT"
47
+ HARM_CATEGORY_DANGEROUS_CONTENT = "HARM_CATEGORY_DANGEROUS_CONTENT"
48
+
49
+ class HarmBlockThreshold:
50
+ BLOCK_NONE = "BLOCK_NONE"
51
+ BLOCK_LOW_AND_ABOVE = "BLOCK_LOW_AND_ABOVE"
52
+ BLOCK_MEDIUM_AND_ABOVE = "BLOCK_MEDIUM_AND_ABOVE"
53
+ BLOCK_ONLY_HIGH = "BLOCK_ONLY_HIGH"
54
+
55
+ # --- Custom Exceptions ---
56
+ class ValidationError(Exception):
57
+ """Custom validation error for agent inputs"""
58
+ pass
59
+
60
+ class RateLimitError(Exception):
61
+ """Placeholder for rate limit errors."""
62
+ pass
63
+
64
+ class AgentNotReadyError(Exception):
65
+ """Agent is not properly initialized"""
66
+ pass
67
+
68
+ # --- Configuration Constants ---
69
+ GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', "")
70
+ LLM_MODEL_NAME = "gemini-2.5-flash-preview-05-20"
71
+ GEMINI_EMBEDDING_MODEL_NAME = "gemini-embedding-exp-03-07"
72
+
73
+ GENERATION_CONFIG_PARAMS = {
74
+ "temperature": 0.7,
75
+ "top_p": 0.95,
76
+ "top_k": 40,
77
+ "max_output_tokens": 8192,
78
+ "candidate_count": 1,
79
+ }
80
+
81
+ DEFAULT_SAFETY_SETTINGS = []
82
+
83
+ # Default RAG documents
84
+ DEFAULT_RAG_DOCUMENTS = pd.DataFrame({
85
+ 'text': [
86
+ "Employer branding focuses on how an organization is perceived as an employer by potential and current employees.",
87
+ "Key metrics for employer branding include employee engagement, candidate quality, and retention rates.",
88
+ "LinkedIn is a crucial platform for showcasing company culture and attracting talent.",
89
+ "Analyzing follower demographics and post engagement helps refine employer branding strategies.",
90
+ "Content strategy should align with company values to attract the right talent.",
91
+ "Employee advocacy programs can significantly boost employer brand reach and authenticity."
92
+ ]
93
+ })
94
+
95
+ # --- Client Initialization ---
96
+ client = None
97
+ if GEMINI_API_KEY and GENAI_AVAILABLE:
98
+ try:
99
+ client = genai.Client(api_key=GEMINI_API_KEY)
100
+ logging.info("Google GenAI client initialized successfully.")
101
+ except Exception as e:
102
+ logging.error(f"Failed to initialize Google GenAI client: {e}")
103
+ client = None
104
+ else:
105
+ if not GEMINI_API_KEY:
106
+ logging.warning("GEMINI_API_KEY environment variable not set.")
107
+ if not GENAI_AVAILABLE:
108
+ logging.warning("Google GenAI library not available.")
109
+
110
+
111
+ # --- Utility function to get DataFrame schema representation ---
112
+ def get_df_schema_representation(df: pd.DataFrame, df_name: str) -> str:
113
+ """Generates a string representation of a DataFrame's schema and a small sample."""
114
+ if not isinstance(df, pd.DataFrame):
115
+ return f"Item '{df_name}' is not a DataFrame.\n"
116
+ if df.empty:
117
+ return f"DataFrame '{df_name}': Empty\n"
118
+
119
+ # Define system columns to exclude from schema representation
120
+ system_columns = ['Created Date', 'Modified Date', '_id']
121
+
122
+ # Filter out system columns for schema representation
123
+ filtered_columns = [col for col in df.columns if col not in system_columns]
124
+
125
+ schema_parts = [f"DataFrame '{df_name}':"]
126
+ schema_parts.append(f" Shape: {df.shape}")
127
+ schema_parts.append(" Columns:")
128
+
129
+ # Show only filtered columns in schema
130
+ for col in filtered_columns:
131
+ col_type = str(df[col].dtype)
132
+ null_count = df[col].isnull().sum()
133
+ unique_count = df[col].nunique()
134
+ schema_parts.append(f" - {col} (Type: {col_type}, Nulls: {null_count}/{len(df)}, Uniques: {unique_count})")
135
+
136
+ # Add note if system columns were excluded
137
+ excluded_columns = [col for col in df.columns if col in system_columns]
138
+ if excluded_columns:
139
+ schema_parts.append(f" Note: System columns excluded from display: {', '.join(excluded_columns)}")
140
+
141
+ if not df.empty and filtered_columns:
142
+ schema_parts.append(" Sample Data (first 2 rows):")
143
+ try:
144
+ # Create sample with only filtered columns
145
+ sample_df = df[filtered_columns].head(2)
146
+ sample_df_str = sample_df.to_string(index=True, max_colwidth=50)
147
+ indented_sample_df = "\n".join([" " + line for line in sample_df_str.split('\n')])
148
+ schema_parts.append(indented_sample_df)
149
+ except Exception as e:
150
+ schema_parts.append(f" Could not generate sample data: {e}")
151
+ elif not df.empty and not filtered_columns:
152
+ schema_parts.append(" Sample Data: Only system columns present, no business data to display")
153
+
154
+ return "\n".join(schema_parts) + "\n"
155
+
156
+ def get_all_schemas_representation(dataframes: Dict[str, pd.DataFrame]) -> str:
157
+ """Generates a string representation of all DataFrame schemas."""
158
+ if not dataframes:
159
+ return "No DataFrames available to the agent."
160
+
161
+ full_representation = ["=== Available DataFrame Schemas for Analysis ==="]
162
+ for name, df_instance in dataframes.items():
163
+ full_representation.append(get_df_schema_representation(df_instance, name))
164
+ return "\n".join(full_representation)
165
+
166
+ class AdvancedRAGSystem:
167
+ def __init__(self, documents_df: pd.DataFrame, embedding_model_name: str):
168
+ self.documents_df = documents_df.copy() if not documents_df.empty else DEFAULT_RAG_DOCUMENTS.copy()
169
+ # Ensure 'text' column exists
170
+ if 'text' not in self.documents_df.columns and not self.documents_df.empty:
171
+ logging.warning("'text' column not found in RAG documents. RAG might not work.")
172
+ self.documents_df['text'] = ""
173
+
174
+ self.embedding_model_name = embedding_model_name
175
+ self.embeddings: Optional[np.ndarray] = None
176
+ self.is_initialized = False
177
+ logging.info(f"AdvancedRAGSystem initialized with {len(self.documents_df)} documents. Model: {self.embedding_model_name}")
178
+
179
+ def _embed_single_document_sync(self, text: str) -> Optional[np.ndarray]:
180
+ if not client:
181
+ raise ConnectionError("GenAI client not initialized for RAG embedding.")
182
+ if not text or not isinstance(text, str):
183
+ logging.warning("Cannot embed empty or non-string text for RAG.")
184
+ return None
185
+
186
+ try:
187
+ embed_config_payload = None
188
+ if GENAI_AVAILABLE and hasattr(types, 'EmbedContentConfig'):
189
+ embed_config_payload = types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT")
190
+
191
+ response = client.models.embed_content(
192
+ model=f"models/{self.embedding_model_name}" if not self.embedding_model_name.startswith("models/") else self.embedding_model_name,
193
+ contents=text,
194
+ config=embed_config_payload
195
+ )
196
+
197
+ # Fix: Handle ContentEmbedding objects properly
198
+ if hasattr(response, 'embeddings') and isinstance(response.embeddings, list) and len(response.embeddings) > 0:
199
+ embedding_obj = response.embeddings[0]
200
+
201
+ # Extract values from ContentEmbedding object
202
+ if hasattr(embedding_obj, 'values'):
203
+ embedding_values = embedding_obj.values
204
+ elif hasattr(embedding_obj, 'embedding'):
205
+ embedding_values = embedding_obj.embedding
206
+ elif isinstance(embedding_obj, (list, tuple)):
207
+ embedding_values = embedding_obj
208
+ else:
209
+ # Try to convert to list/array if it's a different object type
210
+ try:
211
+ embedding_values = list(embedding_obj)
212
+ except:
213
+ logging.error(f"Cannot extract embedding values from object type: {type(embedding_obj)}")
214
+ return None
215
+
216
+ return np.array(embedding_values, dtype=np.float32)
217
+ else:
218
+ logging.error(f"Unexpected response structure")
219
+ return None
220
+
221
+ except Exception as e:
222
+ logging.error(f"Error in _embed_single_document_sync for text '{text[:50]}...': {e}", exc_info=True)
223
+ raise
224
+
225
+ async def initialize_embeddings(self):
226
+ if self.documents_df.empty or 'text' not in self.documents_df.columns:
227
+ logging.warning("RAG documents DataFrame is empty or lacks 'text' column. Skipping embedding.")
228
+ self.embeddings = np.array([])
229
+ self.is_initialized = True
230
+ return
231
+
232
+ if not client and not (GENAI_AVAILABLE and os.getenv('GEMINI_API_KEY')):
233
+ logging.error("GenAI client not available for RAG embedding initialization.")
234
+ self.embeddings = np.array([])
235
+ return
236
+
237
+ logging.info(f"Starting RAG document embedding for {len(self.documents_df)} documents...")
238
+ embedded_docs_list = []
239
+
240
+ for index, row in self.documents_df.iterrows():
241
+ text_to_embed = row.get('text', '')
242
+ if not text_to_embed or not isinstance(text_to_embed, str):
243
+ logging.warning(f"Skipping RAG document at index {index} due to invalid/empty text.")
244
+ continue
245
+
246
+ try:
247
+ embedding_array = await asyncio.to_thread(self._embed_single_document_sync, text_to_embed)
248
+ if embedding_array is not None and embedding_array.size > 0:
249
+ embedded_docs_list.append(embedding_array)
250
+ else:
251
+ logging.warning(f"Empty or failed embedding for RAG document at index {index}.")
252
+ except Exception as e:
253
+ logging.error(f"Error embedding RAG document at index {index}: {e}")
254
+ continue
255
+
256
+ if not embedded_docs_list:
257
+ self.embeddings = np.array([])
258
+ logging.warning("No RAG documents were successfully embedded.")
259
+ else:
260
+ try:
261
+ first_shape = embedded_docs_list[0].shape
262
+ if not all(emb.shape == first_shape for emb in embedded_docs_list):
263
+ logging.error("Inconsistent embedding shapes found. Cannot stack for RAG.")
264
+ self.embeddings = np.array([])
265
+ return
266
+
267
+ self.embeddings = np.vstack(embedded_docs_list)
268
+ logging.info(f"Successfully embedded {len(embedded_docs_list)} RAG documents. Embeddings shape: {self.embeddings.shape}")
269
+ except ValueError as ve:
270
+ logging.error(f"Error stacking embeddings: {ve}")
271
+ self.embeddings = np.array([])
272
+
273
+ self.is_initialized = True
274
+
275
+ def _calculate_cosine_similarity(self, embeddings_matrix: np.ndarray, query_vector: np.ndarray) -> np.ndarray:
276
+ # Ensure inputs are numpy arrays with proper dtype
277
+ embeddings_matrix = np.asarray(embeddings_matrix, dtype=np.float32)
278
+ query_vector = np.asarray(query_vector, dtype=np.float32)
279
+
280
+ if embeddings_matrix.ndim == 1:
281
+ embeddings_matrix = embeddings_matrix.reshape(1, -1)
282
+ if query_vector.ndim == 1:
283
+ query_vector = query_vector.reshape(1, -1)
284
+
285
+ if embeddings_matrix.size == 0 or query_vector.size == 0:
286
+ return np.array([])
287
+
288
+ norm_matrix = np.linalg.norm(embeddings_matrix, axis=1, keepdims=True)
289
+ normalized_embeddings_matrix = np.divide(embeddings_matrix, norm_matrix + 1e-8, where=norm_matrix!=0)
290
+
291
+ norm_query = np.linalg.norm(query_vector, axis=1, keepdims=True)
292
+ normalized_query_vector = np.divide(query_vector, norm_query + 1e-8, where=norm_query!=0)
293
+
294
+ return np.dot(normalized_embeddings_matrix, normalized_query_vector.T).flatten()
295
+
296
+ async def retrieve_relevant_info(self, query: str, top_k: int = 3, min_similarity: float = 0.3) -> str:
297
+ if not self.is_initialized:
298
+ logging.debug("RAG system not initialized. Cannot retrieve info.")
299
+ return ""
300
+ if self.embeddings is None or self.embeddings.size == 0:
301
+ logging.debug("RAG embeddings not available. Cannot retrieve info.")
302
+ return ""
303
+ if not query or not isinstance(query, str):
304
+ logging.debug("Empty or invalid query for RAG retrieval.")
305
+ return ""
306
+
307
+ if not client and not (GENAI_AVAILABLE and os.getenv('GEMINI_API_KEY')):
308
+ logging.error("GenAI client not available for RAG query embedding.")
309
+ return ""
310
+
311
+ try:
312
+ query_vector = await asyncio.to_thread(self._embed_single_document_sync, query)
313
+ if query_vector is None or query_vector.size == 0:
314
+ logging.warning("Query vector embedding failed or is empty for RAG.")
315
+ return ""
316
+
317
+ similarity_scores = self._calculate_cosine_similarity(self.embeddings, query_vector)
318
+ if similarity_scores.size == 0:
319
+ return ""
320
+
321
+ relevant_indices = np.where(similarity_scores >= min_similarity)[0]
322
+ if len(relevant_indices) == 0:
323
+ logging.debug(f"No RAG documents met minimum similarity threshold of {min_similarity} for query: '{query[:50]}...'")
324
+ return ""
325
+
326
+ relevant_scores = similarity_scores[relevant_indices]
327
+ sorted_relevant_indices_of_original = relevant_indices[np.argsort(relevant_scores)[::-1]]
328
+
329
+ top_indices = sorted_relevant_indices_of_original[:top_k]
330
+
331
+ context_parts = []
332
+ if 'text' in self.documents_df.columns:
333
+ for i in top_indices:
334
+ if 0 <= i < len(self.documents_df):
335
+ context_parts.append(self.documents_df.iloc[i]['text'])
336
+
337
+ context = "\n\n---\n\n".join(context_parts)
338
+ logging.debug(f"Retrieved RAG context with {len(context_parts)} documents for query: '{query[:50]}...'")
339
+ return context
340
+
341
+ except Exception as e:
342
+ logging.error(f"Error during RAG retrieval for query '{query[:50]}...': {e}", exc_info=True)
343
+ return ""
344
+
345
+ class EmployerBrandingAgent:
346
+ def __init__(self,
347
+ all_dataframes: Optional[Dict[str, pd.DataFrame]] = None,
348
+ rag_documents_df: Optional[pd.DataFrame] = None,
349
+ llm_model_name: str = LLM_MODEL_NAME,
350
+ embedding_model_name: str = GEMINI_EMBEDDING_MODEL_NAME,
351
+ generation_config_dict: Optional[Dict] = None,
352
+ safety_settings_list: Optional[List] = None):
353
+
354
+ self.all_dataframes = {k: v.copy() for k, v in (all_dataframes or {}).items()}
355
+
356
+ _rag_docs_df = rag_documents_df if rag_documents_df is not None else DEFAULT_RAG_DOCUMENTS.copy()
357
+ self.rag_system = AdvancedRAGSystem(_rag_docs_df, embedding_model_name)
358
+
359
+ self.llm_model_name = llm_model_name
360
+ self.generation_config_dict = generation_config_dict or GENERATION_CONFIG_PARAMS
361
+ self.safety_settings_list = safety_settings_list or DEFAULT_SAFETY_SETTINGS
362
+
363
+ self.chat_history: List[Dict[str, str]] = []
364
+ self.is_ready = False
365
+
366
+ # Create charts directory
367
+ self.charts_dir = "./charts"
368
+ os.makedirs(self.charts_dir, exist_ok=True)
369
+
370
+ # Initialize PandasAI Agent
371
+ self.pandas_agent = None
372
+ self._initialize_pandas_agent()
373
+
374
+ logging.info(f"EnhancedEmployerBrandingAgent initialized. LLM: {self.llm_model_name}. RAG docs: {len(self.rag_system.documents_df)}. DataFrames: {list(self.all_dataframes.keys())}")
375
+
376
+ def _initialize_pandas_agent(self):
377
+ """Initialize PandasAI with enhanced configuration for chart generation"""
378
+ if not self.all_dataframes or not GEMINI_API_KEY:
379
+ logging.warning("Cannot initialize PandasAI agent: missing dataframes or API key")
380
+ return
381
+
382
+ self._preprocess_dataframes_for_pandas_ai()
383
+
384
+ try:
385
+ # Configure LiteLLM with Gemini
386
+ llm = LiteLLM(
387
+ model="gemini/gemini-2.5-flash-preview-05-20",
388
+ api_key=GEMINI_API_KEY
389
+ )
390
+
391
+ # Enhanced PandasAI configuration for better chart generation
392
+ pai.config.set({
393
+ "llm": llm,
394
+ "temperature": 0.3, # Lower temperature for more consistent results
395
+ "verbose": True,
396
+ "enable_cache": False, # Disable cache to avoid stale results
397
+ "save_charts": True,
398
+ "save_charts_path": "./charts",
399
+ "open_charts": False,
400
+ "custom_whitelisted_dependencies": [
401
+ "matplotlib", "seaborn", "plotly", "pandas", "numpy"
402
+ ],
403
+ "max_retries": 3, # Add retry logic
404
+ "use_error_correction_framework": True # Enable error correction
405
+ })
406
+
407
+ # Store dataframes for chat queries
408
+ self.pandas_dfs = {}
409
+ for name, df in self.all_dataframes.items():
410
+ # Skip empty dataframes
411
+ if df.empty:
412
+ continue
413
+
414
+ df_description = self._generate_dataframe_description(name, df)
415
+ pandas_df = pai.DataFrame(df, description=df_description)
416
+ self.pandas_dfs[name] = pandas_df
417
+
418
+ self.pandas_agent = True
419
+ logging.info(f"PandasAI initialized successfully with {len(self.pandas_dfs)} DataFrames")
420
+
421
+ except Exception as e:
422
+ logging.error(f"Failed to initialize PandasAI: {e}", exc_info=True)
423
+ self.pandas_agent = None
424
+ self.pandas_dfs = {}
425
+
426
+ def _generate_dataframe_description(self, name: str, df: pd.DataFrame) -> str:
427
+ """Enhanced dataframe description for better PandasAI understanding"""
428
+ description_parts = [f"This is the '{name}' dataset containing {len(df)} records."]
429
+
430
+ # Add column descriptions based on common patterns
431
+ column_descriptions = []
432
+ for col in df.columns:
433
+ col_lower = col.lower()
434
+ if 'date' in col_lower:
435
+ column_descriptions.append(f"'{col}' contains date/time information")
436
+ elif 'count' in col_lower or 'number' in col_lower:
437
+ column_descriptions.append(f"'{col}' contains numerical count data")
438
+ elif 'rate' in col_lower or 'percentage' in col_lower:
439
+ column_descriptions.append(f"'{col}' contains rate/percentage metrics")
440
+ elif 'follower' in col_lower:
441
+ column_descriptions.append(f"'{col}' contains LinkedIn follower data")
442
+ elif 'engagement' in col_lower:
443
+ column_descriptions.append(f"'{col}' contains engagement metrics")
444
+ elif 'post' in col_lower:
445
+ column_descriptions.append(f"'{col}' contains post-related information")
446
+
447
+ if column_descriptions:
448
+ description_parts.append("Key columns: " + "; ".join(column_descriptions))
449
+
450
+ # Enhanced context for specific datasets
451
+ if name.lower() in ['follower_stats', 'followers']:
452
+ description_parts.append("""
453
+ This data tracks LinkedIn company page follower growth and demographics.
454
+ For monthly growth analysis, use records where follower_count_type='follower_gains_monthly'.
455
+ The 'extracted_date' column contains properly formatted dates for time series analysis.
456
+ Use 'year_month' or 'month_name' columns for better date display in charts.
457
+ For cumulative analysis, use records where follower_count_type='follower_count_cumulative'.
458
+ """)
459
+ elif name.lower().endswith('_monthly_gains'):
460
+ description_parts.append("""
461
+ This is a filtered dataset containing only monthly follower gains data.
462
+ All records have valid dates and are sorted chronologically.
463
+ Use this for creating time series charts of monthly growth patterns.
464
+ """)
465
+ elif name.lower() in ['posts', 'post_stats']:
466
+ description_parts.append("This data contains LinkedIn post performance metrics for employer branding content analysis.")
467
+ elif name.lower() in ['mentions', 'brand_mentions']:
468
+ description_parts.append("This data tracks brand mentions and sentiment for employer branding reputation analysis.")
469
+
470
+ return " ".join(description_parts)
471
+
472
+ async def initialize(self) -> bool:
473
+ """Initializes asynchronous components of the agent"""
474
+ try:
475
+ if not client: # Fix: Remove reference to llm_model_instance
476
+ logging.error("Cannot initialize agent: GenAI client not available/configured.")
477
+ return False
478
+
479
+ await self.rag_system.initialize_embeddings()
480
+
481
+ # Verify PandasAI agent is ready
482
+ pandas_ready = self.pandas_agent is not None
483
+ if not pandas_ready:
484
+ logging.warning("PandasAI agent not initialized, attempting re-initialization")
485
+ self._initialize_pandas_agent()
486
+ pandas_ready = self.pandas_agent is not None
487
+
488
+ self.is_ready = self.rag_system.is_initialized and pandas_ready
489
+ logging.info(f"EnhancedEmployerBrandingAgent.initialize completed. RAG: {self.rag_system.is_initialized}, PandasAI: {pandas_ready}, Agent ready: {self.is_ready}")
490
+ return self.is_ready
491
+
492
+ except Exception as e:
493
+ logging.error(f"Error during EnhancedEmployerBrandingAgent.initialize: {e}", exc_info=True)
494
+ self.is_ready = False
495
+ return False
496
+
497
+
498
+ def _get_dataframes_summary(self) -> str:
499
+ return get_all_schemas_representation(self.all_dataframes)
500
+
501
+ def _preprocess_dataframes_for_pandas_ai(self):
502
+ """Enhanced preprocessing to handle date casting issues and ensure chart generation"""
503
+ if not self.all_dataframes:
504
+ return
505
+
506
+ dataframes_to_add = {} # To store newly created dataframes
507
+
508
+ # Iterate over a copy of the items to avoid runtime errors if modifying the dict
509
+ for name, df_original in list(self.all_dataframes.items()):
510
+ df_copy = df_original.copy() # Work on a copy for this iteration step
511
+
512
+ if name.lower() in ['follower_stats', 'followers']:
513
+ # Handle category_name column that contains dates for follower_gains_monthly
514
+ if 'category_name' in df_copy.columns and 'follower_count_type' in df_copy.columns:
515
+ def extract_date_from_category(row):
516
+ if row.get('follower_count_type') == 'follower_gains_monthly':
517
+ category_name = str(row.get('category_name', ''))
518
+ import re
519
+ date_pattern = r'^\d{4}-\d{2}-\d{2}$'
520
+ if re.match(date_pattern, category_name):
521
+ return category_name
522
+ return None
523
+
524
+ df_copy['extracted_date'] = df_copy.apply(extract_date_from_category, axis=1)
525
+ df_copy['extracted_date'] = pd.to_datetime(df_copy['extracted_date'], errors='coerce')
526
+
527
+ monthly_mask = df_copy['follower_count_type'] == 'follower_gains_monthly'
528
+ # Ensure extracted_date is not NaT before strftime
529
+ valid_dates_mask = monthly_mask & df_copy['extracted_date'].notna()
530
+
531
+ df_copy.loc[valid_dates_mask, 'date_for_analysis'] = df_copy.loc[valid_dates_mask, 'extracted_date']
532
+ df_copy.loc[valid_dates_mask, 'year_month'] = df_copy.loc[valid_dates_mask, 'extracted_date'].dt.strftime('%Y-%m')
533
+ df_copy.loc[valid_dates_mask, 'month_name'] = df_copy.loc[valid_dates_mask, 'extracted_date'].dt.strftime('%B %Y')
534
+
535
+ if 'follower_count' in df_copy.columns:
536
+ df_copy['follower_count'] = pd.to_numeric(df_copy['follower_count'], errors='coerce')
537
+ # df_copy['follower_count'] = df_copy['follower_count'].fillna(0) # Moved to general fillna
538
+
539
+ # Create separate monthly gains dataframe for easier analysis
540
+ if 'follower_count_type' in df_copy.columns and 'extracted_date' in df_copy.columns:
541
+ monthly_gains_df = df_copy[df_copy['follower_count_type'] == 'follower_gains_monthly'].copy()
542
+ if not monthly_gains_df.empty:
543
+ monthly_gains_df = monthly_gains_df.dropna(subset=['extracted_date'])
544
+ if not monthly_gains_df.empty: # Check again after dropna
545
+ monthly_gains_df = monthly_gains_df.sort_values('extracted_date')
546
+ # Store in the temporary dictionary
547
+ dataframes_to_add[f'{name}_monthly_gains'] = monthly_gains_df
548
+ logging.info(f"Created '{name}_monthly_gains' with {len(monthly_gains_df)} records.")
549
+
550
+ # Update the main dataframe in self.all_dataframes with these specific changes
551
+ self.all_dataframes[name] = df_copy.copy() # Save the processed df_copy
552
+ logging.info(f"Preprocessed '{name}' dataframe for date handling.")
553
+
554
+ # General preprocessing for the current dataframe (df_copy or df_original if not 'follower_stats')
555
+ # Fetch the potentially modified df_copy if it was processed above, otherwise use original df for this iteration
556
+ current_df_to_process = self.all_dataframes[name].copy()
557
+
558
+ # Convert object columns that look numeric
559
+ for col in current_df_to_process.columns:
560
+ if current_df_to_process[col].dtype == 'object':
561
+ try:
562
+ # Attempt conversion if a good portion of non-null values match numeric pattern
563
+ if current_df_to_process[col].str.match(r'^-?\d+\.?\d*$').sum() > len(current_df_to_process[col].dropna()) * 0.5:
564
+ current_df_to_process[col] = pd.to_numeric(current_df_to_process[col], errors='coerce')
565
+ logging.info(f"Converted column '{col}' in '{name}' to numeric.")
566
+ except AttributeError: # Handles cases where .str accessor fails (e.g. column has mixed types like numbers and lists)
567
+ logging.debug(f"Could not apply .str accessor to column '{col}' in '{name}'. Skipping numeric conversion for it.")
568
+
569
+
570
+ numeric_columns = current_df_to_process.select_dtypes(include=[np.number]).columns
571
+ current_df_to_process[numeric_columns] = current_df_to_process[numeric_columns].fillna(0)
572
+
573
+ text_columns = current_df_to_process.select_dtypes(include=['object']).columns
574
+ current_df_to_process[text_columns] = current_df_to_process[text_columns].fillna('')
575
+
576
+ # Update self.all_dataframes with the fully processed version for this key
577
+ self.all_dataframes[name] = current_df_to_process
578
+
579
+ # After the loop, add all newly created dataframes
580
+ if dataframes_to_add:
581
+ self.all_dataframes.update(dataframes_to_add)
582
+ logging.info(f"Added new derived dataframes: {list(dataframes_to_add.keys())}")
583
+
584
+
585
+ def _build_system_prompt(self) -> str:
586
+ """Enhanced system prompt that works with PandasAI integration"""
587
+ return textwrap.dedent("""
588
+ You are a friendly and insightful Employer Branding Analyst AI, working as a dedicated partner for HR professionals to make LinkedIn data analysis accessible, actionable, and easy to understand.
589
+
590
+ ## Your Enhanced Capabilities:
591
+ You now have advanced data analysis capabilities through PandasAI integration, allowing you to:
592
+ - Directly query and analyze DataFrames with natural language
593
+ - Generate charts and visualizations automatically (ALWAYS create charts when data visualization would be helpful)
594
+ - Perform complex statistical analysis on LinkedIn employer branding data
595
+ - Handle multi-DataFrame queries and joins seamlessly
596
+
597
+ ## Core Responsibilities:
598
+ 1. **Intelligent Data Analysis**: Use your PandasAI integration to answer data questions directly and accurately
599
+ 2. **Business Context Translation**: Convert technical analysis results into HR-friendly insights
600
+ 3. **Actionable Recommendations**: Provide specific, implementable strategies based on data findings
601
+ 4. **Educational Guidance**: Help users understand both the data insights and the LinkedIn analytics concepts
602
+
603
+ ## CRITICAL COMMUNICATION RULES:
604
+ - **NEVER show code, technical commands, or programming syntax**
605
+ - **NEVER mention dataset names, column names, or technical data structure details**
606
+ - **NEVER reference DataFrames, schemas, or database terminology**
607
+ - **Always speak in business terms**: refer to "your LinkedIn data", "follower metrics", "engagement data", etc.
608
+ - **Focus on insights, not methods**: explain what the data shows, not how it was analyzed
609
+
610
+ ## Communication Style:
611
+ - **Natural and Conversational**: Maintain a warm, supportive tone as a helpful colleague
612
+ - **HR-Focused Language**: Avoid technical jargon; explain analytics terms in business context
613
+ - **Context-Rich Responses**: Always explain what metrics mean for employer branding strategy
614
+ - **Structured Insights**: Use clear formatting with headers, bullets, and logical flow
615
+
616
+ ## Data Analysis Approach:
617
+ When users ask data questions, you will:
618
+ 1. **Leverage PandasAI**: Use your integrated data analysis capabilities to query the data directly
619
+ 2. **Interpret Results**: Translate technical findings into business insights
620
+ 3. **Add Context**: Combine data results with your RAG knowledge base for comprehensive answers
621
+ 4. **Provide Recommendations**: Suggest specific actions based on the analysis
622
+
623
+ ## Response Structure:
624
+ 1. **Executive Summary**: Key findings in business terms
625
+ 2. **Data Insights**: What the analysis reveals (charts/visualizations when helpful)
626
+ 3. **Business Impact**: What this means for employer branding strategy
627
+ 4. **Recommendations**: Specific, prioritized action items
628
+ 5. **Next Steps**: Follow-up suggestions or questions
629
+
630
+ ## Key Behaviors:
631
+ - **Data-Driven**: Always ground insights in actual data analysis when possible
632
+ - **Visual When Helpful**: Suggest or create charts that make data more understandable
633
+ - **Proactive**: Identify related insights the user might find valuable
634
+ - **Honest About Limitations**: Clearly state when data doesn't support certain analyses
635
+
636
+ ## Example Language Patterns:
637
+ - Instead of "DataFrame shows" → "Your LinkedIn data reveals"
638
+ - Instead of "follower_count column" → "follower growth metrics"
639
+ - Instead of "engagement_rate variable" → "post engagement performance"
640
+ - Instead of "dataset analysis" → "performance review"
641
+
642
+ Your goal remains to be a trusted partner, but now with powerful data analysis capabilities that enable deeper, more accurate insights for data-driven employer branding decisions.
643
+ """).strip()
644
+
645
+ def _classify_query_type(self, query: str) -> str:
646
+ """Classify whether query needs data analysis, general advice, or both"""
647
+ data_keywords = [
648
+ 'show', 'analyze', 'chart', 'graph', 'data', 'numbers', 'count', 'total',
649
+ 'average', 'trend', 'compare', 'statistics', 'performance', 'metrics',
650
+ 'followers', 'engagement', 'posts', 'growth', 'rate', 'percentage'
651
+ ]
652
+
653
+ advice_keywords = [
654
+ 'recommend', 'suggest', 'advice', 'strategy', 'improve', 'optimize',
655
+ 'best practice', 'should', 'how to', 'what to do', 'tips'
656
+ ]
657
+
658
+ query_lower = query.lower()
659
+ has_data_request = any(keyword in query_lower for keyword in data_keywords)
660
+ has_advice_request = any(keyword in query_lower for keyword in advice_keywords)
661
+
662
+ if has_data_request and has_advice_request:
663
+ return "hybrid"
664
+ elif has_data_request:
665
+ return "data"
666
+ elif has_advice_request:
667
+ return "advice"
668
+ else:
669
+ return "general"
670
+
671
+ # Replace the _generate_pandas_response method and everything after it with this properly indented code:
672
+
673
+ async def _generate_pandas_response(self, query: str) -> tuple[str, bool]:
674
+ """Generate response using PandasAI with enhanced error handling"""
675
+ if not self.pandas_agent or not hasattr(self, 'pandas_dfs'):
676
+ return "Data analysis not available - PandasAI not initialized.", False
677
+
678
+ try:
679
+ logging.info(f"Processing data query with PandasAI: {query[:100]}...")
680
+
681
+ # Clear any existing matplotlib figures
682
+ import matplotlib.pyplot as plt
683
+ plt.clf()
684
+ plt.close('all')
685
+
686
+ # Enhanced query processing based on content
687
+ processed_query = query
688
+
689
+ # Add helpful context for common chart requests
690
+ if any(word in query.lower() for word in ['chart', 'graph', 'plot', 'visualize']):
691
+ if 'monthly' in query.lower() and 'follower' in query.lower():
692
+ processed_query += """.
693
+ Use the monthly gains data (follower_count_type='follower_gains_monthly')
694
+ and use the extracted_date or month_name column for the x-axis.
695
+ Make sure to filter out any null dates and sort by date.
696
+ Create a clear line chart showing the trend over time."""
697
+ elif 'cumulative' in query.lower() and 'follower' in query.lower():
698
+ processed_query += """.
699
+ Use the cumulative data (follower_count_type='follower_count_cumulative')
700
+ and create a chart showing the total follower growth over time."""
701
+
702
+ # Execute the query
703
+ pandas_response = None
704
+ if len(self.pandas_dfs) == 1:
705
+ df = list(self.pandas_dfs.values())[0]
706
+ logging.info(f"Using single DataFrame for query with shape: {df.df.shape}")
707
+ pandas_response = df.chat(processed_query)
708
+ else:
709
+ dfs = list(self.pandas_dfs.values())
710
+ pandas_response = pai.chat(processed_query, *dfs)
711
+
712
+ # Enhanced response processing with better type handling
713
+ response_text = ""
714
+ chart_path = None
715
+
716
+ # Handle different response types from PandasAI
717
+ try:
718
+ # Case 1: Direct string response (file path)
719
+ if isinstance(pandas_response, str):
720
+ if pandas_response.endswith(('.png', '.jpg', '.jpeg', '.svg')):
721
+ chart_path = pandas_response
722
+ response_text = "Analysis completed with visualization"
723
+ else:
724
+ response_text = pandas_response
725
+
726
+ # Case 2: Chart object response
727
+ elif hasattr(pandas_response, 'value') and hasattr(pandas_response, '_get_image'):
728
+ # Handle PandasAI Chart response object
729
+ try:
730
+ # Try to get the chart path without calling show() which causes the error
731
+ if hasattr(pandas_response, 'value'):
732
+ if isinstance(pandas_response.value, str) and pandas_response.value.endswith(('.png', '.jpg', '.jpeg', '.svg')):
733
+ chart_path = pandas_response.value
734
+ response_text = "Analysis completed with visualization"
735
+ elif isinstance(pandas_response.value, dict):
736
+ # Handle dict response from Chart object
737
+ if 'path' in pandas_response.value:
738
+ chart_path = pandas_response.value['path']
739
+ response_text = "Analysis completed with visualization"
740
+ else:
741
+ response_text = "Chart generated but path not accessible"
742
+ except Exception as chart_error:
743
+ logging.warning(f"Error handling chart response: {chart_error}")
744
+ response_text = "Chart generated but encountered display issue"
745
+
746
+ # Case 3: Response with plot_path attribute
747
+ elif hasattr(pandas_response, 'plot_path') and pandas_response.plot_path:
748
+ chart_path = pandas_response.plot_path
749
+ response_text = getattr(pandas_response, 'text', "Analysis completed with visualization")
750
+
751
+ # Case 4: Other response types
752
+ else:
753
+ if pandas_response is not None:
754
+ response_text = str(pandas_response).strip()
755
+
756
+ except Exception as response_error:
757
+ logging.warning(f"Error processing PandasAI response: {response_error}")
758
+ response_text = "Analysis completed but encountered response processing issue"
759
+
760
+ # Fallback: Check charts directory for new files if no chart path found
761
+ if not chart_path and os.path.exists(self.charts_dir):
762
+ chart_files = []
763
+ for f in os.listdir(self.charts_dir):
764
+ if f.endswith(('.png', '.jpg', '.jpeg', '.svg')):
765
+ full_path = os.path.join(self.charts_dir, f)
766
+ chart_files.append((full_path, os.path.getmtime(full_path)))
767
+
768
+ if chart_files:
769
+ # Sort by modification time (newest first)
770
+ chart_files.sort(key=lambda x: x[1], reverse=True)
771
+ latest_chart_path, latest_time = chart_files[0]
772
+
773
+ # Check if created in last 60 seconds
774
+ import time
775
+ if time.time() - latest_time < 60:
776
+ chart_path = latest_chart_path
777
+ logging.info(f"Found recent chart: {chart_path}")
778
+
779
+ # Format final response
780
+ if not response_text:
781
+ response_text = "Analysis completed"
782
+
783
+ chart_info = ""
784
+ if chart_path and os.path.exists(chart_path):
785
+ chart_info = f"\n\n📊 **Chart Generated**: {os.path.basename(chart_path)}\nChart saved at: {chart_path}"
786
+ logging.info(f"Chart successfully generated: {chart_path}")
787
+
788
+ final_response = response_text + chart_info
789
+ return final_response, True
790
+
791
+ except Exception as e:
792
+ logging.error(f"Error in PandasAI processing: {e}", exc_info=True)
793
+
794
+ # Enhanced error handling
795
+ error_str = str(e).lower()
796
+ if "matplotlib" in error_str and "none" in error_str:
797
+ return "I encountered a data visualization error. This might be due to missing or null values in your data. Please try asking for the raw data first, or specify which specific columns you'd like to analyze.", False
798
+ elif "strftime" in error_str:
799
+ return "I encountered a date formatting issue. Please try asking for the data without specific date formatting, or ask me to show the raw data structure first.", False
800
+ elif "ambiguous" in error_str:
801
+ return "I encountered an ambiguous data type issue. Please try being more specific about which data you'd like to analyze (e.g., 'show monthly follower gains' vs 'show cumulative followers').", False
802
+ elif "startswith" in error_str or "dict" in error_str:
803
+ return "I encountered a response formatting issue. The analysis may have completed but I couldn't process the result properly. Please try rephrasing your query.", False
804
+ else:
805
+ return f"Error processing data query: {str(e)}", False
806
+
807
+ async def _generate_enhanced_response(self, query: str, pandas_result: str = "", query_type: str = "general") -> str:
808
+ """Generate enhanced response combining PandasAI results with RAG context"""
809
+ if not self.is_ready:
810
+ return "Agent is not ready. Please initialize."
811
+ if not client:
812
+ return "Error: AI service is not available. Check API configuration."
813
+
814
+ try:
815
+ system_prompt = self._build_system_prompt()
816
+ data_summary = self._get_dataframes_summary()
817
+ rag_context = await self.rag_system.retrieve_relevant_info(query, top_k=2, min_similarity=0.25)
818
+
819
+ # Build enhanced prompt based on query type and available results
820
+ if query_type == "data" and pandas_result:
821
+ enhanced_prompt = f"""
822
+ {system_prompt}
823
+
824
+ ## Data Analysis Context:
825
+ {data_summary}
826
+
827
+ ## PandasAI Analysis Result:
828
+ {pandas_result}
829
+
830
+ ## Additional Knowledge Context:
831
+ {rag_context if rag_context else 'No additional context retrieved.'}
832
+
833
+ ## User Query:
834
+ {query}
835
+
836
+ Please interpret the data analysis result above and provide business insights in a friendly, HR-focused manner.
837
+ Explain what the findings mean for employer branding strategy and provide actionable recommendations.
838
+ """
839
+ else:
840
+ enhanced_prompt = f"""
841
+ {system_prompt}
842
+
843
+ ## Available Data Context:
844
+ {data_summary}
845
+
846
+ ## Knowledge Base Context:
847
+ {rag_context if rag_context else 'No specific background information retrieved.'}
848
+
849
+ ## User Query:
850
+ {query}
851
+
852
+ Please provide helpful insights and recommendations for this employer branding query.
853
+ """
854
+
855
+ # Fix: Use only genai.Client approach - remove all google-generativeai code
856
+ logging.debug(f"Using genai.Client for enhanced response generation")
857
+
858
+ # Prepare config
859
+ config_dict = dict(self.generation_config_dict) if self.generation_config_dict else {}
860
+
861
+ if self.safety_settings_list:
862
+ safety_settings = []
863
+ for ss in self.safety_settings_list:
864
+ if isinstance(ss, dict):
865
+ if GENAI_AVAILABLE and hasattr(types, 'SafetySetting'):
866
+ safety_settings.append(types.SafetySetting(
867
+ category=ss.get('category'),
868
+ threshold=ss.get('threshold')
869
+ ))
870
+ else:
871
+ safety_settings.append(ss)
872
+ else:
873
+ safety_settings.append(ss)
874
+ config_dict['safety_settings'] = safety_settings
875
+
876
+ if GENAI_AVAILABLE and hasattr(types, 'GenerateContentConfig'):
877
+ config = types.GenerateContentConfig(**config_dict)
878
+ else:
879
+ config = config_dict
880
+
881
+ model_path = f"models/{self.llm_model_name}" if not self.llm_model_name.startswith("models/") else self.llm_model_name
882
+
883
+ # Fix: Use synchronous call wrapped in asyncio.to_thread
884
+ api_response = await asyncio.to_thread(
885
+ client.models.generate_content,
886
+ model=model_path,
887
+ contents=enhanced_prompt, # Fix: Pass single prompt string instead of complex message format
888
+ config=config
889
+ )
890
+
891
+ # Fix: Updated response parsing
892
+ if hasattr(api_response, 'candidates') and api_response.candidates:
893
+ candidate = api_response.candidates[0]
894
+ if hasattr(candidate, 'content') and candidate.content:
895
+ if hasattr(candidate.content, 'parts') and candidate.content.parts:
896
+ response_text_parts = [part.text for part in candidate.content.parts if hasattr(part, 'text')]
897
+ response_text = "".join(response_text_parts).strip()
898
+ else:
899
+ response_text = str(candidate.content).strip()
900
+ else:
901
+ response_text = ""
902
+ else:
903
+ response_text = ""
904
+
905
+ if not response_text:
906
+ # Handle blocked or empty responses
907
+ if hasattr(api_response, 'prompt_feedback') and api_response.prompt_feedback:
908
+ if hasattr(api_response.prompt_feedback, 'block_reason') and api_response.prompt_feedback.block_reason:
909
+ logging.warning(f"Prompt blocked: {api_response.prompt_feedback.block_reason}")
910
+ return f"I'm sorry, your request was blocked. Please try rephrasing your query."
911
+ return "I couldn't generate a response. Please try rephrasing your query."
912
+
913
+ return response_text
914
+
915
+ except Exception as e:
916
+ error_message = str(e).lower()
917
+
918
+ if any(keyword in error_message for keyword in ['blocked', 'safety', 'filter', 'prohibited']):
919
+ logging.error(f"Blocked prompt: {e}")
920
+ return "I'm sorry, your request was blocked by the safety filter. Please rephrase your query."
921
+ else:
922
+ logging.error(f"Error in _generate_enhanced_response: {e}", exc_info=True)
923
+ return f"I encountered an error while processing your request: {str(e)}"
924
+
925
+ def _validate_query(self, query: str) -> bool:
926
+ """Validate user query input"""
927
+ if not query or not isinstance(query, str) or len(query.strip()) < 3:
928
+ logging.warning(f"Invalid query: too short or not a string. Query: '{query}'")
929
+ return False
930
+ if len(query) > 3000:
931
+ logging.warning(f"Invalid query: too long. Length: {len(query)}")
932
+ return False
933
+ return True
934
+
935
+ async def process_query(self, user_query: str) -> Dict[str, Optional[str]]:
936
+ """
937
+ Main method to process user queries.
938
+ Returns a dictionary: {"text": llm_response_string, "image_path": path_to_chart_or_none}
939
+ """
940
+ if not self._validate_query(user_query):
941
+ return {"text": "Please provide a valid query (3 to 3000 characters).", "image_path": None}
942
+
943
+ if not self.is_ready:
944
+ logging.warning("process_query called but agent is not ready. Attempting re-initialization.")
945
+ init_success = await self.initialize()
946
+ if not init_success:
947
+ return {"text": "The agent is not properly initialized and could not be started. Please check configuration and logs.", "image_path": None}
948
+
949
+ try:
950
+ query_type = self._classify_query_type(user_query)
951
+ logging.info(f"Query classified as: {query_type}")
952
+
953
+ pandas_text_output: Optional[str] = None
954
+ pandas_chart_path: Optional[str] = None
955
+ pandas_success = False # Flag to track if PandasAI ran successfully
956
+
957
+ # For data-related queries, try PandasAI first
958
+ if query_type in ["data", "hybrid"] and self.pandas_agent:
959
+ logging.info("Attempting PandasAI analysis...")
960
+ pandas_text_output, pandas_success = await self._generate_pandas_response(user_query)
961
+
962
+ if pandas_success:
963
+ logging.info(f"PandasAI analysis successful. Text: '{str(pandas_text_output)[:100]}...'")
964
+ # Check for chart generation in response
965
+ if "Chart Generated" in pandas_text_output:
966
+ # Extract chart path from response if present
967
+ lines = pandas_text_output.split('\n')
968
+ for line in lines:
969
+ if "Chart saved at:" in line:
970
+ pandas_chart_path = line.split("Chart saved at: ")[1].strip()
971
+ break
972
+ else:
973
+ # pandas_text_output might contain the error message from PandasAI
974
+ logging.warning(f"PandasAI analysis failed or returned no specific result. Message from PandasAI: {pandas_text_output}")
975
+
976
+ # Prepare the context from PandasAI for the LLM
977
+ llm_context_from_pandas = ""
978
+ if pandas_text_output: # This could be a success message or an error message from PandasAI
979
+ llm_context_from_pandas += f"Data Analysis Tool Output: {pandas_text_output}\n"
980
+ if pandas_chart_path and pandas_success: # Only mention chart path if PandasAI was successful
981
+ llm_context_from_pandas += f"[A chart has been generated by the data tool and saved at '{pandas_chart_path}'. You should refer to this chart in your explanation if it's relevant to the user's query.]\n"
982
+ elif query_type in ["data", "hybrid"] and not self.pandas_agent:
983
+ llm_context_from_pandas += "Note: The data analysis tool is currently unavailable.\n"
984
+
985
+ # Always call the LLM to formulate the final response
986
+ final_llm_response = await self._generate_enhanced_response(
987
+ query=user_query,
988
+ pandas_result=llm_context_from_pandas, # Pass the textual summary from PandasAI
989
+ query_type=query_type
990
+ )
991
+
992
+ # Return the LLM's response and the chart path if PandasAI was successful and generated one.
993
+ # If PandasAI failed, pandas_chart_path would be None.
994
+ # The final_llm_response should ideally explain any failures if pandas_text_output contained an error.
995
+ return {"text": final_llm_response, "image_path": pandas_chart_path if pandas_success else None}
996
+
997
+ except Exception as e:
998
+ logging.error(f"Critical error in process_query: {e}", exc_info=True)
999
+ return {"text": f"I encountered a critical error while processing your request: {type(e).__name__}. Please check the logs.", "image_path": None}
1000
+
1001
+ def update_dataframes(self, new_dataframes: Dict[str, pd.DataFrame]):
1002
+ """Updates the agent's DataFrames and reinitializes PandasAI agent"""
1003
+ self.all_dataframes = {k: v.copy() for k, v in new_dataframes.items()}
1004
+ logging.info(f"Agent DataFrames updated. Keys: {list(self.all_dataframes.keys())}")
1005
+
1006
+ # Reinitialize PandasAI agent with new data
1007
+ self._initialize_pandas_agent()
1008
+
1009
+ # Note: RAG system uses static documents and doesn't need reinitialization
1010
+
1011
+ def update_rag_documents(self, new_rag_df: pd.DataFrame):
1012
+ """Updates RAG documents and reinitializes embeddings"""
1013
+ self.rag_system.documents_df = new_rag_df.copy()
1014
+ logging.info(f"RAG documents updated. Count: {len(new_rag_df)}")
1015
+ # Note: Embeddings will need to be reinitialized - call initialize() after this
1016
+
1017
+ def clear_chat_history(self):
1018
+ """Clears the agent's internal chat history"""
1019
+ self.chat_history = []
1020
+ logging.info("EmployerBrandingAgent internal chat history cleared.")
1021
+
1022
+ def get_status(self) -> Dict[str, Any]:
1023
+ """Returns comprehensive status information about the agent"""
1024
+ return {
1025
+ "is_ready": self.is_ready,
1026
+ "has_api_key": bool(GEMINI_API_KEY),
1027
+ "genai_available": GENAI_AVAILABLE,
1028
+ "client_type": "genai.Client" if client else "None", # Fix: Remove reference to llm_model_instance
1029
+ "rag_initialized": self.rag_system.is_initialized,
1030
+ "pandas_agent_ready": self.pandas_agent is not None,
1031
+ "num_dataframes": len(self.all_dataframes),
1032
+ "dataframe_keys": list(self.all_dataframes.keys()),
1033
+ "num_rag_documents": len(self.rag_system.documents_df) if self.rag_system.documents_df is not None else 0,
1034
+ "llm_model_name": self.llm_model_name,
1035
+ "embedding_model_name": self.rag_system.embedding_model_name,
1036
+ "chat_history_length": len(self.chat_history),
1037
+ "charts_save_path_pandasai": pai.config.save_charts_path if pai.config.llm else "PandasAI not configured"
1038
+ }
1039
+
1040
+ def get_available_analyses(self) -> List[str]:
1041
+ """Returns list of suggested analyses based on available data"""
1042
+ if not self.all_dataframes:
1043
+ return ["No data available for analysis"]
1044
+
1045
+ suggestions = []
1046
+ for df_name, df in self.all_dataframes.items():
1047
+ if 'follower' in df_name.lower():
1048
+ suggestions.extend([
1049
+ f"Show follower growth trends from {df_name}",
1050
+ f"Analyze follower demographics in {df_name}",
1051
+ f"Compare follower engagement rates"
1052
+ ])
1053
+ elif 'post' in df_name.lower():
1054
+ suggestions.extend([
1055
+ f"Analyze post performance metrics from {df_name}",
1056
+ f"Show best performing content types",
1057
+ f"Compare engagement across post categories"
1058
+ ])
1059
+ elif 'mention' in df_name.lower():
1060
+ suggestions.extend([
1061
+ f"Analyze brand mention sentiment from {df_name}",
1062
+ f"Show mention volume trends",
1063
+ f"Identify top mention sources"
1064
+ ])
1065
+
1066
+ # Add general suggestions
1067
+ suggestions.extend([
1068
+ "What are the key employer branding trends?",
1069
+ "How can I improve our LinkedIn presence?",
1070
+ "What content strategy should we adopt?",
1071
+ "How do we measure employer branding success?"
1072
+ ])
1073
+
1074
+ return suggestions[:10] # Limit to top 10 suggestions
1075
+
1076
+ # --- Helper Functions for External Integration ---
1077
+ def create_agent_instance(dataframes: Optional[Dict[str, pd.DataFrame]] = None,
1078
+ rag_docs: Optional[pd.DataFrame] = None) -> EmployerBrandingAgent:
1079
+ """Factory function to create a new agent instance"""
1080
+ logging.info("Creating new EnhancedEmployerBrandingAgent instance via helper function.")
1081
+ return EmployerBrandingAgent(all_dataframes=dataframes, rag_documents_df=rag_docs)
1082
+
1083
+ async def initialize_agent_async(agent: EmployerBrandingAgent) -> bool:
1084
+ """Async helper to initialize an agent instance"""
1085
+ logging.info("Initializing agent via async helper function.")
1086
+ return await agent.initialize()
1087
+
1088
+ def validate_dataframes(dataframes: Dict[str, pd.DataFrame]) -> Dict[str, List[str]]:
1089
+ """Validate dataframes for common issues and return validation report"""
1090
+ validation_report = {}
1091
+
1092
+ for name, df in dataframes.items():
1093
+ issues = []
1094
+
1095
+ if df.empty:
1096
+ issues.append("DataFrame is empty")
1097
+
1098
+ # Check for required columns based on data type
1099
+ if 'follower' in name.lower():
1100
+ required_cols = ['date', 'follower_count']
1101
+ missing_cols = [col for col in required_cols if col not in df.columns]
1102
+ if missing_cols:
1103
+ issues.append(f"Missing expected columns for follower data: {missing_cols}")
1104
+
1105
+ elif 'post' in name.lower():
1106
+ required_cols = ['date', 'engagement']
1107
+ missing_cols = [col for col in required_cols if col not in df.columns]
1108
+ if missing_cols:
1109
+ issues.append(f"Missing expected columns for post data: {missing_cols}")
1110
+
1111
+ # Check for data quality issues
1112
+ if not df.empty:
1113
+ null_percentages = (df.isnull().sum() / len(df) * 100).round(2)
1114
+ high_null_cols = null_percentages[null_percentages > 50].index.tolist()
1115
+ if high_null_cols:
1116
+ issues.append(f"Columns with >50% null values: {high_null_cols}")
1117
+
1118
+ validation_report[name] = issues
1119
+
1120
+ return validation_report