Update eb_agent_module.py

eb_agent_module.py

@@ -9,10 +9,8 @@ import textwrap
 
 # Attempt to import Google Generative AI and related types
 try:
-
-    from google import genai
-
-    from google.genai import types as genai_types
+    from google import generativeai as genai # Renamed for clarity to avoid conflict
+    from google.generativeai import types as genai_types
     # from google.generativeai import GenerationConfig # For direct use if needed
     # from google.generativeai.types import HarmCategory, HarmBlockThreshold, SafetySetting # For direct use
 
@@ -26,14 +24,13 @@ except ImportError:
 
         # Dummy Client and related structures
         class Client:
-            def __init__(self, api_key=None):
+            def __init__(self, api_key=None): # api_key is optional for Client constructor
                 self.api_key = api_key
-                self.models = self._Models()
-                print(f"Dummy genai.Client initialized {'with' if api_key else 'without'} API key.")
+                self.models = self._Models() # This is the service client for models
+                print(f"Dummy genai.Client initialized {'with api_key' if api_key else '(global API key expected)'}.")
 
-            class _Models:
-                @staticmethod
-                async def generate_content_async(model=None, contents=None, generation_config=None, safety_settings=None, stream=False): # Matched real signature better
+            class _Models: # Represents the model service client
+                async def generate_content_async(self, model=None, contents=None, generation_config=None, safety_settings=None, stream=False, tools=None, tool_config=None): # Matched real signature better
                     print(f"Dummy genai.Client.models.generate_content_async called for model: {model} with config: {generation_config}, safety_settings: {safety_settings}, stream: {stream}")
                     class DummyPart:
                         def __init__(self, text): self.text = text
@@ -44,20 +41,20 @@ except ImportError:
                             self.content = DummyContent()
                             self.finish_reason = genai_types.FinishReason.STOP # Use dummy FinishReason
                             self.safety_ratings = []
-                            self.token_count = 0 # Added
-                            self.index = 0 # Added
+                            self.token_count = 0
+                            self.index = 0
                     class DummyResponse:
                         def __init__(self):
                             self.candidates = [DummyCandidate()]
-                            self.prompt_feedback = self._PromptFeedback() # Use dummy PromptFeedback
-                            self.text = "# Dummy response text from dummy client's async generate_content" # for easier access
-                        class _PromptFeedback: # Nested dummy class
+                            self.prompt_feedback = self._PromptFeedback()
+                            self.text = "# Dummy response text from dummy client's async generate_content"
+                        class _PromptFeedback:
                             def __init__(self):
                                 self.block_reason = None
                                 self.safety_ratings = []
                     return DummyResponse()
 
-                def generate_content(self, model=None, contents=None, generation_config=None, safety_settings=None, stream=False): # Matched real signature better
+                def generate_content(self, model=None, contents=None, generation_config=None, safety_settings=None, stream=False, tools=None, tool_config=None): # Matched real signature better
                     print(f"Dummy genai.Client.models.generate_content called for model: {model} with config: {generation_config}, safety_settings: {safety_settings}, stream: {stream}")
                     # Re-using the async dummy structure for simplicity
                     class DummyPart:
@@ -67,14 +64,14 @@ except ImportError:
                     class DummyCandidate:
                         def __init__(self):
                             self.content = DummyContent()
-                            self.finish_reason = genai_types.FinishReason.STOP # Use dummy FinishReason
+                            self.finish_reason = genai_types.FinishReason.STOP
                             self.safety_ratings = []
                             self.token_count = 0
                             self.index = 0
                     class DummyResponse:
                         def __init__(self):
                             self.candidates = [DummyCandidate()]
-                            self.prompt_feedback = self._PromptFeedback() # Use dummy PromptFeedback
+                            self.prompt_feedback = self._PromptFeedback()
                             self.text = "# Dummy response text from dummy client's generate_content"
                         class _PromptFeedback:
                             def __init__(self):
@@ -83,134 +80,78 @@ except ImportError:
                     return DummyResponse()
 
         @staticmethod
-        def GenerativeModel(model_name, generation_config=None, safety_settings=None, system_instruction=None): # Matched real signature
-            print(f"Dummy genai.GenerativeModel called for model: {model_name} with config: {generation_config}, safety: {safety_settings}, system_instruction: {system_instruction}")
+        def GenerativeModel(model_name, generation_config=None, safety_settings=None, system_instruction=None): # Kept for AdvancedRAGSystem if it uses it, or if user switches back
+            print(f"Dummy genai.GenerativeModel called for model: {model_name} (This might be unused if Client approach is preferred)")
+            # ... (rest of DummyGenerativeModel as before, for completeness) ...
             class DummyGenerativeModel:
                 def __init__(self, model_name_in, generation_config_in, safety_settings_in, system_instruction_in):
                     self.model_name = model_name_in
-                    self.generation_config = generation_config_in
-                    self.safety_settings = safety_settings_in
-                    self.system_instruction = system_instruction_in
-                async def generate_content_async(self, contents, stream=False): # Matched real signature
-                    print(f"Dummy GenerativeModel.generate_content_async called for {self.model_name}")
-                    # Simplified response, similar to Client's dummy
-                    class DummyPart:
-                        def __init__(self, text): self.text = text
-                    class DummyContent:
-                        def __init__(self): self.parts = [DummyPart(f"# Dummy response from dummy GenerativeModel ({self.model_name})")]
-                    class DummyCandidate:
-                        def __init__(self):
-                            self.content = DummyContent()
-                            self.finish_reason = genai_types.FinishReason.STOP
-                            self.safety_ratings = []
-                    class DummyResponse:
-                        def __init__(self):
-                            self.candidates = [DummyCandidate()]
-                            self.prompt_feedback = None
-                            self.text = f"# Dummy response text from dummy GenerativeModel ({self.model_name})"
-                    return DummyResponse()
-
-                def generate_content(self, contents, stream=False): # Matched real signature
-                    print(f"Dummy GenerativeModel.generate_content called for {self.model_name}")
-                    # Simplified response, similar to Client's dummy
+                async def generate_content_async(self, contents, stream=False):
                     class DummyPart:
                         def __init__(self, text): self.text = text
                     class DummyContent:
                         def __init__(self): self.parts = [DummyPart(f"# Dummy response from dummy GenerativeModel ({self.model_name})")]
                     class DummyCandidate:
                         def __init__(self):
-                            self.content = DummyContent()
-                            self.finish_reason = genai_types.FinishReason.STOP
-                            self.safety_ratings = []
+                            self.content = DummyContent(); self.finish_reason = genai_types.FinishReason.STOP; self.safety_ratings = []
                     class DummyResponse:
                         def __init__(self):
-                            self.candidates = [DummyCandidate()]
-                            self.prompt_feedback = None
-                            self.text = f"# Dummy response text from dummy GenerativeModel ({self.model_name})"
+                            self.candidates = [DummyCandidate()]; self.prompt_feedback = None; self.text = f"# Dummy GM response"
                     return DummyResponse()
-
             return DummyGenerativeModel(model_name, generation_config, safety_settings, system_instruction)
 
+
         @staticmethod
         def embed_content(model, content, task_type, title=None):
             print(f"Dummy genai.embed_content called for model: {model}, task_type: {task_type}, title: {title}")
-            # Ensure the dummy embedding matches typical dimensions (e.g., 768 for many models)
             return {"embedding": [0.1] * 768}
 
     class genai_types: # type: ignore
-        # Using dicts for dummy GenerationConfig and SafetySetting for simplicity
         @staticmethod
-        def GenerationConfig(**kwargs): # The dummy now just returns the kwargs as a dict
+        def GenerationConfig(**kwargs):
             print(f"Dummy genai_types.GenerationConfig created with: {kwargs}")
             return dict(kwargs)
 
         @staticmethod
         def SafetySetting(category, threshold):
             print(f"Dummy SafetySetting created: category={category}, threshold={threshold}")
-            return {"category": category, "threshold": threshold} # Return a dict for dummy
+            return {"category": category, "threshold": threshold}
 
-        # Dummy Enums (can be simple string attributes)
         class HarmCategory:
-            HARM_CATEGORY_UNSPECIFIED = "HARM_CATEGORY_UNSPECIFIED"
-            HARM_CATEGORY_HARASSMENT = "HARM_CATEGORY_HARASSMENT"
-            HARM_CATEGORY_HATE_SPEECH = "HARM_CATEGORY_HATE_SPEECH"
-            HARM_CATEGORY_SEXUALLY_EXPLICIT = "HARM_CATEGORY_SEXUALLY_EXPLICIT"
-            HARM_CATEGORY_DANGEROUS_CONTENT = "HARM_CATEGORY_DANGEROUS_CONTENT"
-
+            HARM_CATEGORY_UNSPECIFIED = "HARM_CATEGORY_UNSPECIFIED"; HARM_CATEGORY_HARASSMENT = "HARM_CATEGORY_HARASSMENT"; HARM_CATEGORY_HATE_SPEECH = "HARM_CATEGORY_HATE_SPEECH"; HARM_CATEGORY_SEXUALLY_EXPLICIT = "HARM_CATEGORY_SEXUALLY_EXPLICIT"; HARM_CATEGORY_DANGEROUS_CONTENT = "HARM_CATEGORY_DANGEROUS_CONTENT"
         class HarmBlockThreshold:
-            BLOCK_NONE = "BLOCK_NONE"
-            BLOCK_LOW_AND_ABOVE = "BLOCK_LOW_AND_ABOVE"
-            BLOCK_MEDIUM_AND_ABOVE = "BLOCK_MEDIUM_AND_ABOVE"
-            BLOCK_ONLY_HIGH = "BLOCK_ONLY_HIGH"
-
-        class FinishReason: # Added dummy FinishReason
-            FINISH_REASON_UNSPECIFIED = "FINISH_REASON_UNSPECIFIED"
-            STOP = "STOP"
-            MAX_TOKENS = "MAX_TOKENS"
+            BLOCK_NONE = "BLOCK_NONE"; BLOCK_LOW_AND_ABOVE = "BLOCK_LOW_AND_ABOVE"; BLOCK_MEDIUM_AND_ABOVE = "BLOCK_MEDIUM_AND_ABOVE"; BLOCK_ONLY_HIGH = "BLOCK_ONLY_HIGH"
+        class FinishReason:
+            FINISH_REASON_UNSPECIFIED = "UNSPECIFIED"; STOP = "STOP"; MAX_TOKENS = "MAX_TOKENS"; SAFETY = "SAFETY"; RECITATION = "RECITATION"; OTHER = "OTHER"
+        
+        # Dummy for BlockedReason if needed by response parsing
+        class BlockedReason:
+            BLOCKED_REASON_UNSPECIFIED = "BLOCKED_REASON_UNSPECIFIED"
             SAFETY = "SAFETY"
-            RECITATION = "RECITATION"
             OTHER = "OTHER"
 
-        # Placeholder for other types if needed by the script
-        # class BlockReason:
-        #     SAFETY = "SAFETY"
-
-
 # --- Configuration ---
 GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', "")
-# Recommended: Use a standard, publicly available model name.
+# User-specified model names:
 # LLM_MODEL_NAME = "gemini-2.0-flash" # Original
 LLM_MODEL_NAME = "gemini-2.0-flash" 
 GEMINI_EMBEDDING_MODEL_NAME = "gemini-embedding-exp-03-07"
 
 # Base generation configuration for the LLM
 GENERATION_CONFIG_PARAMS = {
-    "temperature": 0.3, # Slightly increased for more varied insights, adjust as needed
+    "temperature": 0.3,
     "top_p": 1.0,
     "top_k": 32,
-    "max_output_tokens": 8192, # Increased for potentially longer code with comments and insights
-    # "candidate_count": 1, # Default is 1, explicitly setting it
+    "max_output_tokens": 8192,
 }
 
 # Default safety settings list for Gemini
 try:
     DEFAULT_SAFETY_SETTINGS = [
-        genai_types.SafetySetting(
-            category=genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
-            threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE, # Adjusted slightly
-        ),
-        genai_types.SafetySetting(
-            category=genai_types.HarmCategory.HARM_CATEGORY_HARASSMENT,
-            threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE, # Adjusted slightly
-        ),
-        genai_types.SafetySetting(
-            category=genai_types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
-            threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE, # Adjusted slightly
-        ),
-        genai_types.SafetySetting(
-            category=genai_types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
-            threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE, # Adjusted slightly
-        ),
+        genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
+        genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
+        genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
+        genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
     ]
 except AttributeError as e:
     logging.warning(f"Could not define DEFAULT_SAFETY_SETTINGS using real genai_types: {e}. Using placeholder list of dicts.")
@@ -221,7 +162,6 @@ except AttributeError as e:
         {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
     ]
 
-
 # Logging setup
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(filename)s:%(lineno)d - %(message)s')
 
@@ -237,268 +177,166 @@ else:
 
 # --- RAG Documents Definition (Example) ---
 rag_documents_data = {
-    'Title': [
-        "Employer Branding Best Practices 2024",
-        "Attracting Tech Talent in Competitive Markets",
-        "The Power of Employee Advocacy",
-        "Understanding Gen Z Workforce Expectations"
-    ],
-    'Text': [
-        "Focus on authentic employee stories and showcase company culture. Highlight diversity and inclusion initiatives. Use video content for higher engagement. Clearly articulate your Employee Value Proposition (EVP).",
-        "Tech candidates value challenging projects, continuous learning opportunities, and a flexible work environment. Competitive compensation and modern tech stacks are crucial. Highlight your company's innovation and impact.",
-        "Encourage employees to share their positive experiences on social media. Provide them with shareable content and guidelines. Employee-generated content is often perceived as more trustworthy than corporate messaging.",
-        "Gen Z values purpose-driven work, transparency, mental health support, and opportunities for growth. They are digital natives and expect seamless online application processes. They also care deeply about social responsibility."
-    ]
+    'Title': ["Employer Branding Best Practices 2024", "Attracting Tech Talent", "Employee Advocacy", "Gen Z Expectations"],
+    'Text': ["Focus on authentic employee stories...", "Tech candidates value challenging projects...", "Encourage employees to share experiences...", "Gen Z values purpose-driven work..."]
 }
 df_rag_documents = pd.DataFrame(rag_documents_data)
 
 # --- Schema Representation ---
 def get_schema_representation(df_name: str, df: pd.DataFrame) -> str:
-    if not isinstance(df, pd.DataFrame):
-        return f"Schema for item '{df_name}': Not a DataFrame.\n"
-    if df.empty:
-        return f"Schema for DataFrame 'df_{df_name}': Empty (no columns or rows).\n"
-    
-    schema_str = f"DataFrame 'df_{df_name}':\n"
-    schema_str += f"  Columns: {df.columns.tolist()}\n"
-    schema_str += f"  Shape: {df.shape}\n"
-    # Add dtypes for more clarity
-    # schema_str += "  Data Types:\n"
-    # for col in df.columns:
-    #     schema_str += f"    {col}: {df[col].dtype}\n"
-    
-    # Sample data (first 2 rows)
-    if not df.empty:
-        sample_str = df.head(2).to_string()
-        # Indent sample string for better readability in the prompt
-        indented_sample = "\n".join(["    " + line for line in sample_str.splitlines()])
-        schema_str += f"  Sample Data (first 2 rows):\n{indented_sample}\n"
-    else:
-        schema_str += "  Sample Data: DataFrame is empty.\n"
+    if not isinstance(df, pd.DataFrame): return f"Schema for item '{df_name}': Not a DataFrame.\n"
+    if df.empty: return f"Schema for DataFrame 'df_{df_name}': Empty.\n"
+    schema_str = f"DataFrame 'df_{df_name}':\n  Columns: {df.columns.tolist()}\n  Shape: {df.shape}\n"
+    if not df.empty: schema_str += f"  Sample Data (first 2 rows):\n{textwrap.indent(df.head(2).to_string(), '    ')}\n"
+    else: schema_str += "  Sample Data: DataFrame is empty.\n"
     return schema_str
 
 def get_all_schemas_representation(dataframes_dict: dict) -> str:
-    if not dataframes_dict:
-        return "No DataFrames provided.\n"
+    if not dataframes_dict: return "No DataFrames provided.\n"
     return "".join(get_schema_representation(name, df) for name, df in dataframes_dict.items())
 
-
 # --- Advanced RAG System ---
 class AdvancedRAGSystem:
     def __init__(self, documents_df: pd.DataFrame, embedding_model_name: str):
         self.embedding_model_name = embedding_model_name
         self.documents_df = documents_df.copy()
         self.embeddings_generated = False
-        self.client_available = hasattr(genai, 'embed_content') and not (hasattr(genai.embed_content, '__func__') and genai.embed_content.__func__.__qualname__.startswith('genai.embed_content')) # Check if it's not the dummy
+        # Check if genai.embed_content is the real one or our dummy
+        self.client_available = hasattr(genai, 'embed_content') and not (hasattr(genai.embed_content, '__func__') and genai.embed_content.__func__.__qualname__.startswith('genai.embed_content'))
 
         if GEMINI_API_KEY and self.client_available:
             try:
                 self._precompute_embeddings()
                 self.embeddings_generated = True
                 logging.info(f"RAG embeddings precomputed using '{self.embedding_model_name}'.")
-            except Exception as e:
-                logging.error(f"RAG precomputation error: {e}", exc_info=True)
+            except Exception as e: logging.error(f"RAG precomputation error: {e}", exc_info=True)
         else:
-            logging.warning(f"RAG embeddings not precomputed. GEMINI_API_KEY set: {bool(GEMINI_API_KEY)}, genai.embed_content available: {self.client_available}.")
+            logging.warning(f"RAG embeddings not precomputed. Key: {bool(GEMINI_API_KEY)}, embed_content_ok: {self.client_available}.")
 
     def _embed_fn(self, title: str, text: str) -> list[float]:
-        if not self.embeddings_generated: # Should rely on self.client_available too
-            # logging.debug(f"Skipping embedding for '{title}' as embeddings are not active.")
-            return [0.0] * 768 # Default dimension, adjust if your model differs
+        if not self.client_available: return [0.0] * 768 
         try:
-            # logging.debug(f"Embedding '{title}' with model '{self.embedding_model_name}'")
-            # Ensure content is not empty
             content_to_embed = text if text else title
-            if not content_to_embed:
-                logging.warning(f"Cannot embed '{title}' because both title and text are empty.")
-                return [0.0] * 768
-
-            embedding_result = genai.embed_content(
-                model=self.embedding_model_name,
-                content=content_to_embed,
-                task_type="retrieval_document",
-                title=title if title else None # Pass title only if it exists
-            )
-            return embedding_result["embedding"]
+            if not content_to_embed: return [0.0] * 768
+            return genai.embed_content(model=self.embedding_model_name, content=content_to_embed, task_type="retrieval_document", title=title if title else None)["embedding"]
         except Exception as e:
             logging.error(f"Error in _embed_fn for '{title}': {e}", exc_info=True)
             return [0.0] * 768
 
     def _precompute_embeddings(self):
-        if 'Embeddings' not in self.documents_df.columns:
-            self.documents_df['Embeddings'] = pd.Series(dtype='object')
-        
-        # Ensure there's text to embed
-        mask = (self.documents_df['Text'].notna() & (self.documents_df['Text'] != '')) | \
-               (self.documents_df['Title'].notna() & (self.documents_df['Title'] != ''))
-        
-        if not mask.any():
-            logging.warning("No content found in 'Text' or 'Title' columns to generate embeddings.")
-            return
-
-        self.documents_df.loc[mask, 'Embeddings'] = self.documents_df[mask].apply(
-            lambda row: self._embed_fn(row.get('Title', ''), row.get('Text', '')), axis=1
-        )
-        logging.info(f"Applied embedding function to {mask.sum()} rows.")
-
-
-    def retrieve_relevant_info(self, query_text: str, top_k: int = 2) -> str: # Increased top_k for more context
-        if not self.client_available:
-             return "\n[RAG Context]\nEmbedding client not available. Cannot retrieve RAG context.\n"
+        if 'Embeddings' not in self.documents_df.columns: self.documents_df['Embeddings'] = pd.Series(dtype='object')
+        mask = (self.documents_df['Text'].notna() & (self.documents_df['Text'] != '')) | (self.documents_df['Title'].notna() & (self.documents_df['Title'] != ''))
+        if not mask.any(): logging.warning("No content for RAG embeddings."); return
+        self.documents_df.loc[mask, 'Embeddings'] = self.documents_df[mask].apply(lambda row: self._embed_fn(row.get('Title', ''), row.get('Text', '')), axis=1)
+        logging.info(f"Applied RAG embedding function to {mask.sum()} rows.")
+
+    def retrieve_relevant_info(self, query_text: str, top_k: int = 2) -> str:
+        if not self.client_available: return "\n[RAG Context]\nEmbedding client not available.\n"
         if not self.embeddings_generated or 'Embeddings' not in self.documents_df.columns or self.documents_df['Embeddings'].isnull().all():
-            return "\n[RAG Context]\nEmbeddings not generated or all are null. No RAG context available.\n"
-
+            return "\n[RAG Context]\nEmbeddings not ready for RAG.\n"
         try:
-            query_embedding_response = genai.embed_content(
-                model=self.embedding_model_name,
-                content=query_text,
-                task_type="retrieval_query"
-            )
-            query_embedding = np.array(query_embedding_response["embedding"])
-
-            valid_embeddings_df = self.documents_df.dropna(subset=['Embeddings'])
-            valid_embeddings_df = valid_embeddings_df[valid_embeddings_df['Embeddings'].apply(lambda x: isinstance(x, (list, np.ndarray)) and len(x) > 0)]
-
-            if valid_embeddings_df.empty:
-                return "\n[RAG Context]\nNo valid document embeddings for RAG.\n"
-
-            document_embeddings = np.stack(valid_embeddings_df['Embeddings'].apply(np.array).values)
-
-            if query_embedding.shape[0] != document_embeddings.shape[1]:
-                logging.error(f"Embedding dimension mismatch. Query: {query_embedding.shape[0]}, Docs: {document_embeddings.shape[1]}")
-                return "\n[RAG Context]\nEmbedding dimension mismatch. Cannot calculate similarity.\n"
-
-            dot_products = np.dot(document_embeddings, query_embedding)
-            # Get indices of top_k largest dot products
-            # If fewer valid documents than top_k, take all of them
-            num_to_retrieve = min(top_k, len(valid_embeddings_df))
-            if num_to_retrieve == 0: # Should be caught by valid_embeddings_df.empty earlier
-                 return "\n[RAG Context]\nNo relevant passages found (num_to_retrieve is 0).\n"
-
-            # Ensure indices are within bounds
-            idx = np.argsort(dot_products)[-num_to_retrieve:][::-1] # Top N, descending order
-
-            relevant_passages = ""
-            for i in idx:
-                if i < len(valid_embeddings_df): # Defensive check
-                    doc = valid_embeddings_df.iloc[i]
-                    relevant_passages += f"\n[RAG Context from: '{doc['Title']}']\n{doc['Text']}\n"
-                else:
-                    logging.warning(f"Index {i} out of bounds for valid_embeddings_df (len {len(valid_embeddings_df)})")
-
-
-            return relevant_passages if relevant_passages else "\n[RAG Context]\nNo relevant passages found after similarity search.\n"
+            query_embedding = np.array(genai.embed_content(model=self.embedding_model_name, content=query_text, task_type="retrieval_query")["embedding"])
+            valid_df = self.documents_df.dropna(subset=['Embeddings'])
+            valid_df = valid_df[valid_df['Embeddings'].apply(lambda x: isinstance(x, (list, np.ndarray)) and len(x) > 0)]
+            if valid_df.empty: return "\n[RAG Context]\nNo valid document embeddings.\n"
+            
+            doc_embeddings = np.stack(valid_df['Embeddings'].apply(np.array).values)
+            if query_embedding.shape[0] != doc_embeddings.shape[1]: return "\n[RAG Context]\nEmbedding dimension mismatch.\n"
+            
+            dot_products = np.dot(doc_embeddings, query_embedding)
+            num_to_retrieve = min(top_k, len(valid_df))
+            if num_to_retrieve == 0: return "\n[RAG Context]\nNo relevant passages found (num_to_retrieve is 0).\n"
+            
+            idx = np.argsort(dot_products)[-num_to_retrieve:][::-1]
+            passages = "".join([f"\n[RAG Context from: '{valid_df.iloc[i]['Title']}']\n{valid_df.iloc[i]['Text']}\n" for i in idx if i < len(valid_df)])
+            return passages if passages else "\n[RAG Context]\nNo relevant passages found after search.\n"
         except Exception as e:
             logging.error(f"Error in RAG retrieve_relevant_info: {e}", exc_info=True)
             return f"\n[RAG Context]\nError during RAG retrieval: {type(e).__name__} - {e}\n"
 
-
-# --- PandasLLM Class (Gemini-Powered) ---
+# --- PandasLLM Class (Gemini-Powered using genai.Client) ---
 class PandasLLM:
     def __init__(self, llm_model_name: str,
                  generation_config_dict: dict,
                  safety_settings_list: list,
                  data_privacy=True, force_sandbox=True):
         self.llm_model_name = llm_model_name
-        self.generation_config_dict = generation_config_dict
-        self.safety_settings_list = safety_settings_list
+        self.generation_config_dict = generation_config_dict # Will be passed to API call
+        self.safety_settings_list = safety_settings_list     # Will be passed to API call
         self.data_privacy = data_privacy
         self.force_sandbox = force_sandbox
-        self.generative_model = None # Will hold the GenerativeModel instance
-
-        if not GEMINI_API_KEY:
-            logging.warning(f"PandasLLM: GEMINI_API_KEY not set. Using dummy model if real 'genai' is not fully mocked.")
-            # Even if API key is not set, we might be using a dummy genai
-            # So, initialize the dummy model if genai.GenerativeModel is the dummy one
-            if hasattr(genai, 'GenerativeModel') and hasattr(genai.GenerativeModel, '__func__') and genai.GenerativeModel.__func__.__qualname__.startswith('genai.GenerativeModel'): # Heuristic for dummy
-                 self.generative_model = genai.GenerativeModel(
-                    model_name=self.llm_model_name,
-                    generation_config=genai_types.GenerationConfig(**self.generation_config_dict) if self.generation_config_dict else None,
-                    safety_settings=self.safety_settings_list
-                )
-                 logging.info(f"PandasLLM: Initialized with DUMMY genai.GenerativeModel for '{self.llm_model_name}'.")
-
-        else: # GEMINI_API_KEY is set
-            try:
-                # Use genai_types.GenerationConfig for real API
-                config_for_model = genai_types.GenerationConfig(**self.generation_config_dict) if self.generation_config_dict else None
-                
-                self.generative_model = genai.GenerativeModel(
-                    model_name=self.llm_model_name, # The SDK handles the "models/" prefix
-                    generation_config=config_for_model,
-                    safety_settings=self.safety_settings_list
-                    # system_instruction can be added here if needed globally for this model
-                )
-                logging.info(f"PandasLLM: Initialized with REAL genai.GenerativeModel for '{self.llm_model_name}'.")
-            except Exception as e:
-                logging.error(f"Failed to initialize PandasLLM with genai.GenerativeModel: {e}", exc_info=True)
-                # Fallback to dummy if real initialization fails, to prevent crashes
-                if hasattr(genai, 'GenerativeModel') and hasattr(genai.GenerativeModel, '__func__') and genai.GenerativeModel.__func__.__qualname__.startswith('genai.GenerativeModel'):
-                    self.generative_model = genai.GenerativeModel(model_name=self.llm_model_name) # Basic dummy
-                    logging.warning("PandasLLM: Falling back to DUMMY genai.GenerativeModel due to real initialization error.")
+        self.client = None
+        self.model_service = None # This will be client.models
 
+        # Check if genai.Client is the real one or our dummy
+        is_real_genai_client = hasattr(genai, 'Client') and not (hasattr(genai.Client, '__func__') and genai.Client.__func__.__qualname__.startswith('genai.Client'))
 
-    async def _call_gemini_api_async(self, prompt_text: str, history: list = None) -> str:
-        if not self.generative_model:
-            logging.error("PandasLLM: GenerativeModel not available (or not initialized). Cannot call API.")
-            return "# Error: Gemini model not available for API call."
-
-        # Gemini API expects chat history in a specific format
-        # The 'contents' parameter should be a list of Content objects (dicts)
-        # For chat, this list often alternates between 'user' and 'model' roles.
-        # The final part of 'contents' should be the current user prompt.
+        if not GEMINI_API_KEY and is_real_genai_client: # Real client but no API key
+            logging.warning(f"PandasLLM: GEMINI_API_KEY not set, but real 'genai.Client' seems available. API calls may fail if global config is not sufficient.")
+            # Proceed to initialize client; it might work if genai.configure() was successful without explicit key here
+            # or if the environment provides credentials in another way.
         
-        # Convert simple history to Gemini's expected format
+        try:
+            self.client = genai.Client() # API key is usually set via genai.configure or environment
+            self.model_service = self.client.models
+            logging.info(f"PandasLLM: Initialized with genai.Client().models for '{self.llm_model_name}'.")
+        except Exception as e:
+            logging.error(f"Failed to initialize PandasLLM with genai.Client: {e}", exc_info=True)
+            # Fallback to dummy if real initialization fails, to prevent crashes
+            if not is_real_genai_client: # If it was already the dummy, re-initialize dummy
+                 self.client = genai.Client() 
+                 self.model_service = self.client.models
+                 logging.warning("PandasLLM: Falling back to DUMMY genai.Client due to real initialization error or it was already dummy.")
+
+
+    async def _call_gemini_api_async(self, prompt_text: str, history: list = None) -> str:
+        if not self.model_service:
+            logging.error("PandasLLM: Model service (client.models) not available. Cannot call API.")
+            return "# Error: Gemini model service not available for API call."
+
         gemini_history = []
         if history:
             for entry in history:
                 role = "model" if entry.get("role") == "assistant" else entry.get("role", "user")
                 gemini_history.append({"role": role, "parts": [{"text": entry.get("content", "")}]})
         
-        # Add current prompt as the last user message
         current_content = [{"role": "user", "parts": [{"text": prompt_text}]}]
+        contents_for_api = gemini_history + current_content
         
-        # The 'contents' for generate_content should be the full conversation history + current prompt
-        # If there's history, it's usually passed to start_chat, but for one-off generate_content,
-        # we might need to construct it carefully.
-        # For non-chat models or simpler generate_content, just the prompt might be enough.
-        # The GenerativeModel().generate_content typically takes 'contents' which can be a string,
-        # a Part, or a list of Parts for the current turn.
-        # For chat-like behavior with generate_content, the 'contents' list should represent the conversation.
-        
-        # Let's assume the prompt_text is the primary input for this turn.
-        # If the model is a chat model, the history needs to be managed carefully.
-        # For now, we'll pass the prompt_text as the main content for this turn.
-        # The `history` parameter here is for our internal tracking,
-        # the `GenerativeModel` might handle history differently (e.g. via `start_chat`).
-        # For a direct `generate_content` call, we typically provide the current turn's content.
-        # If we want to provide history, it should be part of the `contents` list.
-
-        contents_for_api = gemini_history + current_content # This forms the conversation
-        
-        logging.info(f"\n--- Calling Gemini API (model: {self.llm_model_name}) ---\nContent (last part): {contents_for_api[-1]['parts'][0]['text'][:200]}...\n")
+        # Prepare model ID (e.g., "models/gemini-2.0-flash")
+        model_id_for_api = self.llm_model_name
+        if not model_id_for_api.startswith("models/"):
+            model_id_for_api = f"models/{model_id_for_api}"
+
+        # Prepare generation config object
+        api_generation_config = None
+        if self.generation_config_dict:
+            try:
+                api_generation_config = genai_types.GenerationConfig(**self.generation_config_dict)
+            except Exception as e_cfg:
+                logging.error(f"Error creating GenerationConfig object: {e_cfg}. Using dict as fallback.")
+                api_generation_config = self.generation_config_dict # Fallback to dict
+
+        logging.info(f"\n--- Calling Gemini API via Client (model: {model_id_for_api}) ---\nConfig: {api_generation_config}\nSafety: {bool(self.safety_settings_list)}\nContent (last part text): {contents_for_api[-1]['parts'][0]['text'][:100]}...\n")
         
         try:
-            # The GenerativeModel instance already has config and safety settings.
-            # We just pass the 'contents'.
-            response = await self.generative_model.generate_content_async(
+            response = await self.model_service.generate_content_async(
+                model=model_id_for_api,
                 contents=contents_for_api,
-                # generation_config, safety_settings are already part of self.generative_model
+                generation_config=api_generation_config,
+                safety_settings=self.safety_settings_list
             )
 
+            # ... (Response parsing logic remains largely the same as before) ...
             if hasattr(response, 'prompt_feedback') and response.prompt_feedback and \
                hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason:
+                # ... block reason handling ...
                 block_reason_val = response.prompt_feedback.block_reason
-                # Try to get enum name if available
-                try:
-                    block_reason_str = genai_types.BlockedReason(block_reason_val).name
-                except:
-                    block_reason_str = str(block_reason_val)
-                logging.warning(f"Prompt blocked by API. Reason: {block_reason_str}. Ratings: {response.prompt_feedback.safety_ratings}")
+                block_reason_str = str(block_reason_val.name if hasattr(block_reason_val, 'name') else block_reason_val)
+                logging.warning(f"Prompt blocked by API. Reason: {block_reason_str}.")
                 return f"# Error: Prompt blocked by API. Reason: {block_reason_str}."
 
             llm_output = ""
-            # Standard way to get text from Gemini response
             if hasattr(response, 'text') and isinstance(response.text, str):
                 llm_output = response.text
             elif response.candidates:
@@ -507,29 +345,14 @@ class PandasLLM:
                     llm_output = "".join(part.text for part in candidate.content.parts if hasattr(part, 'text'))
                 
                 if not llm_output and candidate.finish_reason:
+                    # ... finish reason handling ...
                     finish_reason_val = candidate.finish_reason
-                    try:
-                        # finish_reason_str = genai_types.FinishReason(finish_reason_val).name # This might fail if finish_reason_val is an int
-                        finish_reason_str = str(finish_reason_val) # Safer for now
-                        # For real API, finish_reason is an enum member, so .name would work.
-                        # For dummy, it might be a string already.
-                        if hasattr(genai_types.FinishReason, '_enum_map_') and finish_reason_val in genai_types.FinishReason._enum_map_: # Check if it's a valid enum value
-                             finish_reason_str = genai_types.FinishReason(finish_reason_val).name
-
-                    except Exception as fre:
-                        logging.debug(f"Could not get FinishReason name: {fre}")
-                        finish_reason_str = str(finish_reason_val)
-
-                    # Check if blocked due to safety
+                    finish_reason_str = str(finish_reason_val.name if hasattr(finish_reason_val, 'name') else finish_reason_val)
+
                     if finish_reason_str == "SAFETY": # or candidate.finish_reason == genai_types.FinishReason.SAFETY:
-                         safety_messages = []
-                         if candidate.safety_ratings:
-                             for rating in candidate.safety_ratings:
-                                 cat_name = rating.category.name if hasattr(rating.category, 'name') else str(rating.category)
-                                 prob_name = rating.probability.name if hasattr(rating.probability, 'name') else str(rating.probability)
-                                 safety_messages.append(f"Category: {cat_name}, Probability: {prob_name}")
-                         logging.warning(f"Content generation stopped due to safety. Finish reason: {finish_reason_str}. Details: {'; '.join(safety_messages)}")
-                         return f"# Error: Content generation stopped by API due to safety. Finish Reason: {finish_reason_str}. Details: {'; '.join(safety_messages)}"
+                         # ... safety message handling ...
+                         logging.warning(f"Content generation stopped due to safety. Finish reason: {finish_reason_str}.")
+                         return f"# Error: Content generation stopped by API due to safety. Finish Reason: {finish_reason_str}."
                     
                     logging.warning(f"Empty response from LLM. Finish reason: {finish_reason_str}.")
                     return f"# Error: LLM returned an empty response. Finish reason: {finish_reason_str}."
@@ -537,17 +360,16 @@ class PandasLLM:
                 logging.error(f"Unexpected API response structure: {str(response)[:500]}")
                 return f"# Error: Unexpected API response structure: {str(response)[:200]}"
             
-            # logging.debug(f"LLM Raw Output:\n{llm_output}")
             return llm_output
 
-        except genai_types.BlockedPromptException as bpe: # Specific exception for blocked prompts
-            logging.error(f"Prompt was blocked by the API (BlockedPromptException): {bpe}", exc_info=True)
-            return f"# Error: Your prompt was blocked by the API. Please revise. Details: {bpe.prompt_feedback}"
-        except genai_types.StopCandidateException as sce: # Specific exception for candidate stopped
-            logging.error(f"Candidate generation stopped (StopCandidateException): {sce}", exc_info=True)
-            return f"# Error: Content generation was stopped. Details: {sce.candidate}"
+        except genai_types.BlockedPromptException as bpe:
+            logging.error(f"Prompt blocked (BlockedPromptException): {bpe}", exc_info=True)
+            return f"# Error: Prompt blocked. Details: {bpe}"
+        except genai_types.StopCandidateException as sce:
+            logging.error(f"Candidate stopped (StopCandidateException): {sce}", exc_info=True)
+            return f"# Error: Content generation stopped. Details: {sce}"
         except Exception as e:
-            logging.error(f"Error calling Gemini API: {e}", exc_info=True)
+            logging.error(f"Error calling Gemini API via Client: {e}", exc_info=True)
             return f"# Error during API call: {type(e).__name__} - {str(e)[:100]}."
 
 
@@ -556,75 +378,47 @@ class PandasLLM:
 
         if self.force_sandbox:
             code_to_execute = ""
-            # Robust code extraction
             if "```python" in llm_response_text:
                 try:
-                    # Standard ```python\nCODE\n```
                     code_block_match = llm_response_text.split("```python\n", 1)
-                    if len(code_block_match) > 1:
-                        code_to_execute = code_block_match[1].split("\n```", 1)[0]
-                    else: # Try without newline after ```python
+                    if len(code_block_match) > 1: code_to_execute = code_block_match[1].split("\n```", 1)[0]
+                    else:
                         code_block_match = llm_response_text.split("```python", 1)
                         if len(code_block_match) > 1:
                            code_to_execute = code_block_match[1].split("```", 1)[0]
-                           if code_to_execute.startswith("\n"): # Remove leading newline if present
-                               code_to_execute = code_to_execute[1:]
-
-                except IndexError:
-                    code_to_execute = "" # Should not happen with proper split logic
+                           if code_to_execute.startswith("\n"): code_to_execute = code_to_execute[1:]
+                except IndexError: code_to_execute = ""
 
             if llm_response_text.startswith("# Error:") or not code_to_execute.strip():
-                logging.warning(f"LLM response is an error, or no valid Python code block found. Raw LLM response: {llm_response_text}")
-                # If LLM returns an error or no code, pass that through directly.
-                # Or if it's a polite non-code refusal (e.g. "# Hello there! ...")
+                logging.warning(f"LLM error or no code: {llm_response_text[:200]}")
                 if not code_to_execute.strip() and not llm_response_text.startswith("# Error:"):
-                    # This means LLM might have responded with natural language instead of code.
-                    # If force_sandbox is true, we expect code. If it's not code, it's a deviation.
-                    # However, the prompt allows for comments like "# Hello there..."
-                    # So, if it's just comments, that's fine.
-                    # If it's substantial text without code delimiters, that's an issue for sandbox mode.
-                    if "```" not in llm_response_text and len(llm_response_text.strip()) > 0: # Heuristic for non-code text
-                        logging.info(f"LLM produced text output instead of Python code in sandbox mode. Passing through: {llm_response_text}")
-                        # This might be desired if the LLM is explaining why it can't generate code.
-                        return llm_response_text # Pass through LLM's direct response
-                return llm_response_text # Pass through LLM's error or its non-code (comment-only) response
-
-            logging.info(f"\n--- Code to Execute (extracted from LLM response): ---\n{code_to_execute}\n----------------------\n")
-            
+                    if "```" not in llm_response_text and len(llm_response_text.strip()) > 0:
+                        logging.info(f"LLM text output in sandbox mode: {llm_response_text[:200]}")
+                return llm_response_text
+
+            logging.info(f"\n--- Code to Execute: ---\n{code_to_execute}\n----------------------\n")
             from io import StringIO
             import sys
-            old_stdout = sys.stdout
-            sys.stdout = captured_output = StringIO()
-            
-            # Prepare globals for exec. Prefix DataFrames with 'df_' as per prompt.
+            old_stdout, sys.stdout = sys.stdout, StringIO()
             exec_globals = {'pd': pd, 'np': np}
             if dataframes_dict:
                 for name, df_instance in dataframes_dict.items():
-                    if isinstance(df_instance, pd.DataFrame):
-                        exec_globals[f"df_{name}"] = df_instance
-                    else:
-                        logging.warning(f"Item '{name}' in dataframes_dict is not a DataFrame. Skipping for exec_globals.")
-            
+                    if isinstance(df_instance, pd.DataFrame): exec_globals[f"df_{name}"] = df_instance
+                    else: logging.warning(f"Item '{name}' not a DataFrame.")
             try:
-                exec(code_to_execute, exec_globals, {}) # Using empty dict for locals
-                final_output_str = captured_output.getvalue()
-                
+                exec(code_to_execute, exec_globals, {})
+                final_output_str = sys.stdout.getvalue()
                 if not final_output_str.strip():
-                    logging.info("Code executed successfully, but no explicit print() output was generated by the LLM's code.")
-                    # Check if the code was just comments or an empty block
-                    if not any(line.strip() and not line.strip().startswith("#") for line in code_to_execute.splitlines()):
-                        return "# LLM generated only comments or an empty code block. No output produced."
-                    return "# Code executed successfully, but it did not produce any printed output. Please ensure the LLM's Python code includes print() statements for the desired results, insights, or answers."
+                    if not any(ln.strip() and not ln.strip().startswith("#") for ln in code_to_execute.splitlines()):
+                        return "# LLM generated only comments or empty code. No output."
+                    return "# Code executed, but no print() output. Ensure print() for results."
                 return final_output_str
             except Exception as e:
-                logging.error(f"Sandbox Execution Error: {e}\nCode was:\n{code_to_execute}", exc_info=True) # Log full traceback for sandbox error
-                # Indent the problematic code for better display in the error message
-                indented_code = textwrap.indent(code_to_execute, '# ', predicate=lambda line: True)
-                return f"# Sandbox Execution Error: {type(e).__name__}: {e}\n# --- Code that caused error: ---\n{indented_code}"
-            finally:
-                sys.stdout = old_stdout
-        else: # Not force_sandbox
-            return llm_response_text
+                logging.error(f"Sandbox Exec Error: {e}\nCode:\n{code_to_execute}", exc_info=True)
+                indented_code = textwrap.indent(code_to_execute, '# ')
+                return f"# Sandbox Exec Error: {type(e).__name__}: {e}\n# Code:\n{indented_code}"
+            finally: sys.stdout = old_stdout
+        else: return llm_response_text
 
 # --- Employer Branding Agent ---
 class EmployerBrandingAgent:
@@ -636,233 +430,96 @@ class EmployerBrandingAgent:
                  embedding_model_name: str,
                  data_privacy=True, force_sandbox=True):
         
-        self.pandas_llm = PandasLLM(
-            llm_model_name,
-            generation_config_dict,
-            safety_settings_list,
-            data_privacy,
-            force_sandbox
-        )
+        self.pandas_llm = PandasLLM(llm_model_name, generation_config_dict, safety_settings_list, data_privacy, force_sandbox)
         self.rag_system = AdvancedRAGSystem(rag_documents_df, embedding_model_name)
         self.all_dataframes = all_dataframes if all_dataframes else {}
         self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
         self.chat_history = []
-        logging.info("EmployerBrandingAgent Initialized.")
+        logging.info("EmployerBrandingAgent Initialized (using Client API approach).")
 
     def _build_prompt(self, user_query: str, role="Employer Branding Analyst & Strategist", task_decomposition_hint=None, cot_hint=True) -> str:
-        # System Instruction part of the prompt (can also be passed to GenerativeModel directly if API supports it well)
-        # This initial instruction sets the persona and overall goal.
-        prompt = f"You are a highly skilled '{role}'. Your primary goal is to provide actionable employer branding insights and strategic recommendations by analyzing provided data (Pandas DataFrames) and contextual information (RAG documents).\n"
-        prompt += "You will be provided with schemas for available Pandas DataFrames and a user query.\n"
-
-        if self.pandas_llm.data_privacy:
-            prompt += "IMPORTANT: Adhere to data privacy. Do not output raw Personally Identifiable Information (PII) like individual names or specific user contact details. Summarize, aggregate, or anonymize data in your insights.\n"
+        prompt = f"You are a highly skilled '{role}'. Your goal is to provide actionable employer branding insights by analyzing Pandas DataFrames and RAG documents.\n"
+        if self.pandas_llm.data_privacy: prompt += "IMPORTANT: Adhere to data privacy. Summarize/aggregate PII.\n"
 
         if self.pandas_llm.force_sandbox:
             prompt += "\n--- TASK: PYTHON CODE GENERATION FOR INSIGHTS ---\n"
-            prompt += "Your main task is to GENERATE PYTHON CODE. This code should use the Pandas library to analyze the provided DataFrames and incorporate insights from any RAG context. The code's `print()` statements MUST output the final textual insights, analyses, or answers to the user's query.\n"
-            prompt += "Output ONLY the Python code block, starting with ```python and ending with ```.\n"
-            prompt += "The available DataFrames are already loaded and can be accessed by their dictionary keys prefixed with 'df_' (e.g., df_follower_stats, df_posts) within the execution environment.\n"
-            prompt += "Example of accessing a DataFrame: `df_follower_stats['country']`.\n"
-            
+            prompt += "GENERATE PYTHON CODE using Pandas. The code's `print()` statements MUST output final textual insights/answers.\n"
+            prompt += "Output ONLY the Python code block (```python ... ```).\n"
+            prompt += "Access DataFrames as 'df_name' (e.g., `df_follower_stats`).\n"
             prompt += "\n--- CRITICAL INSTRUCTIONS FOR PYTHON CODE OUTPUT ---\n"
-            prompt += "1.  **Print Insights, Not Just Data:** Your Python code's `print()` statements are the agent's final response. These prints MUST articulate clear, actionable insights or answers. Do NOT just print raw DataFrames or intermediate variables unless the query *specifically* asks for a table of data (e.g., 'Show me the first 5 posts').\n"
-            prompt += "    Example of good insight print: `print(f'Key Insight: Content related to {top_theme} receives {avg_engagement_increase}% higher engagement, suggesting a focus on this area.')`\n"
-            prompt += "    Example of what to AVOID for insight queries: `print(df_analysis_result)` (unless df_analysis_result is the specific table requested).\n"
-            prompt += "2.  **Synthesize with RAG Context:** If RAG context is provided, weave takeaways from it into your printed insights. Example: `print(f'Data shows X. Combined with RAG best practice Y, we recommend Z.')`\n"
-            prompt += "3.  **Structure and Comments:** Write clean, commented Python code. Explain your logic for each step.\n"
-            prompt += "4.  **Handle Ambiguity/Errors in Code:**\n"
-            prompt += "    - If the query is ambiguous, `print()` a clarifying question as a string. Do not generate analysis code.\n"
-            prompt += "    - If the query cannot be answered with the given data/schemas, `print()` a statement explaining this. Example: `print('Insight: Cannot determine X as the required data Y is not available in the provided DataFrames.')`\n"
-            prompt += "    - For non-analytical queries (e.g., 'hello'), respond politely with a `print()` statement. Example: `print('Hello! How can I assist with your employer branding data analysis today?')`\n"
-            prompt += "5.  **Function Usage:** If you define functions, ENSURE they are called and their results (or insights derived from them) are `print()`ed.\n"
-            prompt += "6.  **DataFrame Naming:** Remember to use the `df_` prefix for DataFrame names in your code (e.g., `df_your_data`).\n"
-
-        else: # Not force_sandbox - LLM provides direct textual answer
+            prompt += "1.  **Print Insights, Not Just Data:** `print()` clear, actionable insights. NOT raw DataFrames unless specifically asked for a table.\n"
+            prompt += "    Good: `print(f'Insight: Theme {top_theme} has {engagement_increase}% higher engagement.')`\n"
+            prompt += "    Avoid: `print(df_result)` (for insight queries).\n"
+            prompt += "2.  **Synthesize with RAG:** Weave RAG takeaways into printed insights. Ex: `print(f'Data shows X. RAG says Y. Recommend Z.')`\n"
+            prompt += "3.  **Comments & Clarity:** Write clean, commented code.\n"
+            prompt += "4.  **Handle Issues in Code:** If ambiguous, `print()` a question. If data unavailable, `print()` explanation. For non-analytical queries, `print()` polite reply.\n"
+            prompt += "5.  **Function Usage:** Call functions and `print()` their (insightful) results.\n"
+        else: # Not force_sandbox
             prompt += "\n--- TASK: DIRECT TEXTUAL INSIGHT GENERATION ---\n"
-            prompt += "Your task is to analyze the data (described by schemas) and RAG context, then provide a comprehensive textual answer with actionable insights and strategic recommendations. Explain your reasoning step-by-step.\n"
+            prompt += "Analyze data and RAG, then provide a comprehensive textual answer with insights. Explain step-by-step.\n"
 
         prompt += "\n--- AVAILABLE DATA AND SCHEMAS ---\n"
-        if self.schemas_representation.strip() == "No DataFrames provided.":
-             prompt += "No specific DataFrames are currently loaded. Please rely on general knowledge and any provided RAG context for your response, or ask for data to be loaded.\n"
-        else:
-            prompt += self.schemas_representation
+        prompt += self.schemas_representation if self.schemas_representation.strip() != "No DataFrames provided." else "No DataFrames loaded.\n"
         
         rag_context = self.rag_system.retrieve_relevant_info(user_query)
-        # Check if RAG context is meaningful before appending
         meaningful_rag_keywords = ["Error", "No valid", "No relevant", "Cannot retrieve", "not available", "not generated"]
         is_meaningful_rag = bool(rag_context.strip()) and not any(keyword in rag_context for keyword in meaningful_rag_keywords)
-
-        if is_meaningful_rag:
-            prompt += f"\n--- ADDITIONAL CONTEXT (from Employer Branding Knowledge Base - consider this for your insights) ---\n{rag_context}\n"
-        else:
-            prompt += "\n--- ADDITIONAL CONTEXT (from Employer Branding Knowledge Base) ---\nNo specific pre-defined context found highly relevant to this query, or RAG system encountered an issue. Rely on general knowledge and DataFrame analysis.\n"
-
+        if is_meaningful_rag: prompt += f"\n--- RAG CONTEXT ---\n{rag_context}\n"
+        else: prompt += "\n--- RAG CONTEXT ---\nNo specific RAG context found or RAG error.\n"
 
         prompt += f"\n--- USER QUERY ---\n{user_query}\n"
-
-        if task_decomposition_hint:
-            prompt += f"\n--- GUIDANCE FOR ANALYSIS (Task Decomposition) ---\n{task_decomposition_hint}\n"
+        if task_decomposition_hint: prompt += f"\n--- GUIDANCE ---\n{task_decomposition_hint}\n"
         
         if cot_hint:
             if self.pandas_llm.force_sandbox:
-                prompt += "\n--- THOUGHT PROCESS FOR PYTHON CODE GENERATION (Follow these steps) ---\n"
-                prompt += "1.  **Understand Query & Goal:** What specific employer branding insight or answer is the user seeking?\n"
-                prompt += "2.  **Identify Data Sources:** Which DataFrame(s) and column(s) are relevant? Is there RAG context to incorporate?\n"
-                prompt += "3.  **Plan Analysis (Mental Outline / Code Comments):**\n"
-                prompt += "    a.  What calculations, aggregations, or transformations are needed?\n"
-                prompt += "    b.  How will RAG context be integrated into the final printed insight?\n"
-                prompt += "    c.  What is the exact textual insight/answer to be `print()`ed?\n"
-                prompt += "4.  **Write Python Code:** Implement the plan. Use `df_name_of_dataframe`.\n"
-                prompt += "5.  **CRITICAL - Formulate and `print()` Insights:** Construct the final textual insight(s) as strings and use `print()` statements for them. These prints are the agent's entire response. Ensure they are clear, actionable, and directly address the user's query, incorporating RAG if applicable.\n"
-                prompt += "6.  **Review Code:** Check for correctness, clarity, and adherence to ALL instructions, especially the `print()` requirements for insightful text.\n"
-                prompt += "7.  **Final Output:** Ensure ONLY the Python code block (```python...```) is generated.\n"
+                prompt += "\n--- PYTHON CODE GENERATION THOUGHT PROCESS ---\n"
+                prompt += "1. Goal? 2. Data sources (DFs, RAG)? 3. Analysis plan (comments)? 4. Write Python code. 5. CRITICAL: Formulate & `print()` textual insights. 6. Review. 7. Output ONLY ```python ... ```.\n"
             else: # Not force_sandbox
-                prompt += "\n--- THOUGHT PROCESS FOR DIRECT TEXTUAL RESPONSE (Follow these steps) ---\n"
-                prompt += "1.  **Understand Query & Goal:** What specific employer branding insight or answer is the user seeking?\n"
-                prompt += "2.  **Identify Data Sources:** Analyze the DataFrame schemas. Consider relevant RAG context.\n"
-                prompt += "3.  **Formulate Insights:** Synthesize information from data and RAG to derive key insights and recommendations.\n"
-                prompt += "4.  **Structure Response:** Provide a step-by-step explanation of your analysis, followed by the clear, actionable insights and strategic advice.\n"
-        
+                prompt += "\n--- TEXTUAL RESPONSE THOUGHT PROCESS ---\n"
+                prompt += "1. Goal? 2. Data sources? 3. Formulate insights (data + RAG). 4. Structure: explanation, then insights.\n"
         return prompt
 
     async def process_query(self, user_query: str, role="Employer Branding Analyst & Strategist", task_decomposition_hint=None, cot_hint=True) -> str:
-        # Add user query to history before building prompt, so RAG can use the latest query
-        # However, the LLM call itself should get history *excluding* the current query in its history part.
-        
-        current_turn_history_for_llm = self.chat_history[:] # History *before* this turn
-
-        self.chat_history.append({"role": "user", "parts": [{"text": user_query}]}) # Use 'parts' for Gemini
-        
+        current_turn_history_for_llm = self.chat_history[:]
+        self.chat_history.append({"role": "user", "parts": [{"text": user_query}]})
         full_prompt = self._build_prompt(user_query, role, task_decomposition_hint, cot_hint)
-        
-        # Log only a part of the prompt to avoid overly verbose logs
-        # logging.info(f"Full prompt to LLM (showing first 300 and last 300 chars for brevity):\n{full_prompt[:300]}...\n...\n{full_prompt[-300:]}")
-        logging.info(f"Built prompt for user query: {user_query[:100]}...")
-
-
-        # Pass the history *before* the current user query to the LLM
+        logging.info(f"Built prompt for query: {user_query[:100]}...")
         response_text = await self.pandas_llm.query(full_prompt, self.all_dataframes, history=current_turn_history_for_llm)
-        
-        self.chat_history.append({"role": "model", "parts": [{"text": response_text}]}) # Use 'parts' for Gemini
-        
-        MAX_HISTORY_TURNS = 5 # Each turn has a user and a model message
+        self.chat_history.append({"role": "model", "parts": [{"text": response_text}]})
+        MAX_HISTORY_TURNS = 5
         if len(self.chat_history) > MAX_HISTORY_TURNS * 2:
-            # Keep the most recent turns. The history is [user1, model1, user2, model2,...]
             self.chat_history = self.chat_history[-(MAX_HISTORY_TURNS * 2):]
-            logging.info(f"Chat history truncated to last {MAX_HISTORY_TURNS} turns.")
-            
+            logging.info(f"Chat history truncated.")
         return response_text
 
     def update_dataframes(self, new_dataframes: dict):
         self.all_dataframes = new_dataframes if new_dataframes else {}
         self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
-        logging.info(f"EmployerBrandingAgent DataFrames updated. New schemas: {self.schemas_representation[:200]}...")
-        # Potentially clear RAG embeddings if they depend on the old dataframes, or recompute.
-        # For now, RAG is independent of these dataframes.
+        logging.info(f"Agent DataFrames updated. Schemas: {self.schemas_representation[:100]}...")
 
-    def clear_chat_history(self):
-        self.chat_history = []
-        logging.info("EmployerBrandingAgent chat history cleared.")
+    def clear_chat_history(self): self.chat_history = []; logging.info("Agent chat history cleared.")
 
-# --- Example Usage (Conceptual - for testing the module structure) ---
+# --- Example Usage (Conceptual) ---
 async def main_test():
     logging.info("Starting main_test for EmployerBrandingAgent...")
-    
-    # Dummy DataFrames for testing
-    followers_data = {
-        'date': pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-01', '2023-01-03']),
-        'country': ['USA', 'USA', 'Canada', 'UK'],
-        'new_followers': [10, 12, 5, 8]
-    }
-    df_follower_stats = pd.DataFrame(followers_data)
-
-    posts_data = {
-        'post_id': [1, 2, 3, 4],
-        'post_date': pd.to_datetime(['2023-01-01', '2023-01-01', '2023-01-02', '2023-01-03']),
-        'theme': ['Culture', 'Tech', 'Culture', 'Jobs'],
-        'impressions': [1000, 1500, 1200, 2000],
-        'engagements': [50, 90, 60, 120]
-    }
-    df_posts = pd.DataFrame(posts_data)
-    df_posts['engagement_rate'] = df_posts['engagements'] / df_posts['impressions']
-
-    test_dataframes = {
-        "follower_stats": df_follower_stats,
-        "posts": df_posts,
-        "empty_df": pd.DataFrame(), # Test empty df representation
-        "non_df_item": "This is not a dataframe" # Test non-df item
-    }
-
-    # Initialize the agent
-    # Ensure GEMINI_API_KEY is set in your environment for real calls
-    if not GEMINI_API_KEY:
-        logging.warning("GEMINI_API_KEY not found in environment. Testing with dummy/mocked functionality.")
-
-    agent = EmployerBrandingAgent(
-        llm_model_name=LLM_MODEL_NAME,
-        generation_config_dict=GENERATION_CONFIG_PARAMS,
-        safety_settings_list=DEFAULT_SAFETY_SETTINGS,
-        all_dataframes=test_dataframes,
-        rag_documents_df=df_rag_documents, # Using the example RAG data
-        embedding_model_name=GEMINI_EMBEDDING_MODEL_NAME,
-        force_sandbox=True # Set to True to test code generation, False for direct LLM text
-    )
-    
-    logging.info(f"Schema representation:\n{agent.schemas_representation}")
-
-    queries = [
-        "What are the key trends in follower growth by country based on the first few days of January 2023?",
-        "Which post theme has the highest average engagement rate? Provide an insight.",
-        "Hello there!",
-        "Can you tell me the average salary for software engineers? (This should state data is not available)",
-        "Summarize the best practices for attracting tech talent and combine it with an analysis of our top performing post themes."
-    ]
+    df_follower_stats = pd.DataFrame({'date': pd.to_datetime(['2023-01-01']), 'country': ['USA'], 'new_followers': [10]})
+    df_posts = pd.DataFrame({'post_id': [1], 'theme': ['Culture'], 'engagement_rate': [0.05]})
+    test_dataframes = {"follower_stats": df_follower_stats, "posts": df_posts}
 
+    if not GEMINI_API_KEY: logging.warning("GEMINI_API_KEY not set. Testing with dummy functionality.")
+    agent = EmployerBrandingAgent(LLM_MODEL_NAME, GENERATION_CONFIG_PARAMS, DEFAULT_SAFETY_SETTINGS, test_dataframes, df_rag_documents, GEMINI_EMBEDDING_MODEL_NAME, force_sandbox=True)
+    
+    queries = ["Which post theme has the highest average engagement rate? Provide an insight.", "Hello!"]
     for query in queries:
-        logging.info(f"\n\n--- Processing Query: {query} ---")
+        logging.info(f"\n\n--- Query: {query} ---")
         response = await agent.process_query(user_query=query)
-        logging.info(f"--- Agent Response for '{query}': ---\n{response}\n---------------------------\n")
-        # Small delay if making actual API calls to avoid rate limits during testing
-        if GEMINI_API_KEY: await asyncio.sleep(1) 
-
-    # Test updating dataframes
-    new_posts_data = {
-        'post_id': [5, 6], 'post_date': pd.to_datetime(['2023-01-04', '2023-01-05']),
-        'theme': ['Innovation', 'Team'], 'impressions': [2500, 1800], 'engagements': [150, 100]
-    }
-    df_new_posts = pd.DataFrame(new_posts_data)
-    df_new_posts['engagement_rate'] = df_new_posts['engagements'] / df_new_posts['impressions']
-    
-    updated_dataframes = {
-        "follower_stats": df_follower_stats, # unchanged
-        "posts": pd.concat([df_posts, df_new_posts]), # updated
-        "company_values": pd.DataFrame({'value': ['Innovation', 'Collaboration'], 'description': ['...', '...']}) # new df
-    }
-    agent.update_dataframes(updated_dataframes)
-    logging.info(f"\n--- Processing Query after DataFrame Update ---")
-    response_after_update = await agent.process_query("What's the latest top performing post theme now?")
-    logging.info(f"--- Agent Response for 'What's the latest top performing post theme now?': ---\n{response_after_update}\n---------------------------\n")
-
+        logging.info(f"--- Response for '{query}': ---\n{response}\n---------------------------\n")
+        if GEMINI_API_KEY: await asyncio.sleep(1)
 
 if __name__ == "__main__":
-    # This allows running the test if the script is executed directly
-    # Note: For real API calls, ensure GEMINI_API_KEY is set in your environment.
-    # Example: export GEMINI_API_KEY="your_api_key_here"
-    
-    # To run the async main_test:
-    # asyncio.run(main_test())
-    # Or, if you're in a Jupyter environment that has its own loop:
-    # await main_test()
-    
-    # For simplicity in a standard Python script:
-    if GEMINI_API_KEY: # Only run full async test if API key likely present
-        try:
-            asyncio.run(main_test())
+    if GEMINI_API_KEY:
+        try: asyncio.run(main_test())
         except RuntimeError as e:
-            if " asyncio.run() cannot be called from a running event loop" in str(e):
-                print("Skipping asyncio.run(main_test()) as it seems to be in an existing event loop (e.g., Jupyter). Call 'await main_test()' instead if appropriate.")
-            else:
-                raise
-    else:
-        print("GEMINI_API_KEY not set. Skipping main_test() which might make real API calls. The module can be imported and used elsewhere.")
+            if "asyncio.run() cannot be called from a running event loop" in str(e): print("Skip asyncio.run in existing loop.")
+            else: raise
+    else: print("GEMINI_API_KEY not set. Skipping main_test().")