Update eb_agent_module.py

eb_agent_module.py

@@ -9,395 +9,693 @@ 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
+
 except ImportError:
     print("Google Generative AI library not found. Please install it: pip install google-generativeai")
     # Define dummy classes/functions if the import fails, to allow the rest of the script to be parsed
     class genai: # type: ignore
         @staticmethod
-        def configure(api_key): pass
-        
-        @staticmethod
-        def Client(api_key=None): 
-            class DummyModels:
+        def configure(api_key):
+            print(f"Dummy genai.configure called with API key: {'SET' if api_key else 'NOT SET'}")
+
+        # Dummy Client and related structures
+        class Client:
+            def __init__(self, api_key=None):
+                self.api_key = api_key
+                self.models = self._Models()
+                print(f"Dummy genai.Client initialized {'with' if api_key else 'without'} API key.")
+
+            class _Models:
                 @staticmethod
-                def generate_content(model=None, contents=None, config=None, safety_settings=None): # Added config, kept safety_settings for older dummy
-                    print(f"Dummy genai.Client.models.generate_content called for model: {model} with config: {config}, safety_settings: {safety_settings}")
+                async def generate_content_async(model=None, contents=None, generation_config=None, safety_settings=None, stream=False): # 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
                     class DummyContent:
-                        def __init__(self): self.parts = [DummyPart("# Dummy response from dummy client")]
+                        def __init__(self): self.parts = [DummyPart("# Dummy response from dummy client's async generate_content")]
                     class DummyCandidate:
                         def __init__(self):
                             self.content = DummyContent()
-                            self.finish_reason = "DUMMY"
-                            self.safety_ratings = [] # Ensure this attribute exists
+                            self.finish_reason = genai_types.FinishReason.STOP # Use dummy FinishReason
+                            self.safety_ratings = []
+                            self.token_count = 0 # Added
+                            self.index = 0 # Added
                     class DummyResponse:
                         def __init__(self):
                             self.candidates = [DummyCandidate()]
-                            self.prompt_feedback = None # Ensure this attribute exists
-                        @property
-                        def text(self): 
-                            if self.candidates and self.candidates[0].content and self.candidates[0].content.parts:
-                                return "".join(p.text for p in self.candidates[0].content.parts)
-                            return ""
+                            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
+                            def __init__(self):
+                                self.block_reason = None
+                                self.safety_ratings = []
                     return DummyResponse()
 
-            class DummyClient:
-                def __init__(self): self.models = DummyModels()
-            
-            if api_key: return DummyClient()
-            return None
+                def generate_content(self, model=None, contents=None, generation_config=None, safety_settings=None, stream=False): # 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:
+                        def __init__(self, text): self.text = text
+                    class DummyContent:
+                        def __init__(self): self.parts = [DummyPart("# Dummy response from dummy client's generate_content")]
+                    class DummyCandidate:
+                        def __init__(self):
+                            self.content = DummyContent()
+                            self.finish_reason = genai_types.FinishReason.STOP # Use dummy FinishReason
+                            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.text = "# Dummy response text from dummy client's generate_content"
+                        class _PromptFeedback:
+                            def __init__(self):
+                                self.block_reason = None
+                                self.safety_ratings = []
+                    return DummyResponse()
 
         @staticmethod
-        def GenerativeModel(model_name): 
-             print(f"Dummy genai.GenerativeModel called for model: {model_name}")
-             return None
-        
+        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}")
+            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
+                    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()
+
+            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}")
-            return {"embedding": [0.1] * 768} 
+        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 GenerateContentConfig(**kwargs): # The dummy now just returns the kwargs
-            print(f"Dummy genai_types.GenerateContentConfig called with: {kwargs}")
-            return kwargs 
-        
-        # Dummy SafetySetting to allow instantiation if real genai_types is missing
+        def GenerationConfig(**kwargs): # The dummy now just returns the kwargs as a dict
+            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
 
-        class BlockReason: 
-            SAFETY = "SAFETY"
-        class HarmCategory: 
+        # 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"
+
         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"
+            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', "") 
+GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', "")
+# Recommended: Use a standard, publicly available model name.
+# 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 (without safety settings here)
+# Base generation configuration for the LLM
 GENERATION_CONFIG_PARAMS = {
-    "temperature": 0.2,
+    "temperature": 0.3, # Slightly increased for more varied insights, adjust as needed
     "top_p": 1.0,
     "top_k": 32,
-    "max_output_tokens": 4096, 
+    "max_output_tokens": 8192, # Increased for potentially longer code with comments and insights
+    # "candidate_count": 1, # Default is 1, explicitly setting it
 }
 
 # Default safety settings list for Gemini
-# This is now a list of SafetySetting objects (or dicts if using dummy)
 try:
-    DEFAULT_SAFETY_SETTINGS = [ 
+    DEFAULT_SAFETY_SETTINGS = [
         genai_types.SafetySetting(
             category=genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
-            threshold=genai_types.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE, 
+            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_NONE, 
+            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_NONE,
+            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_NONE,
+            threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE, # Adjusted slightly
         ),
     ]
-except AttributeError as e: 
+except AttributeError as e:
     logging.warning(f"Could not define DEFAULT_SAFETY_SETTINGS using real genai_types: {e}. Using placeholder list of dicts.")
     DEFAULT_SAFETY_SETTINGS = [
-        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_LOW_AND_ABOVE"},
-        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
-        {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
-        {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
+        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+        {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+        {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
     ]
 
 
 # Logging setup
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(message)s')
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(filename)s:%(lineno)d - %(message)s')
 
 if GEMINI_API_KEY:
     try:
         genai.configure(api_key=GEMINI_API_KEY)
-        logging.info(f"Gemini API key configured globally...")
+        logging.info(f"Gemini API key configured globally.")
     except Exception as e:
         logging.error(f"Failed to configure Gemini API globally: {e}", exc_info=True)
 else:
-    logging.warning("GEMINI_API_KEY environment variable not set.")
+    logging.warning("GEMINI_API_KEY environment variable not set. Agent will use dummy responses if real genai library is not fully mocked or if API calls fail.")
 
 
-# --- RAG Documents Definition ---
+# --- RAG Documents Definition (Example) ---
 rag_documents_data = {
-    'Title': ["Employer Branding Best Practices 2024", "Attracting Tech Talent"],
-    'Text': ["Focus on authentic employee stories...", "Tech candidates value challenging projects..."]
-} 
+    '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."
+    ]
+}
 df_rag_documents = pd.DataFrame(rag_documents_data)
 
 # --- Schema Representation ---
 def get_schema_representation(df_name: str, df: pd.DataFrame) -> str:
-    if df.empty: return f"Schema for DataFrame '{df_name}': Empty.\n"
-    # Truncated for brevity in example, keep your full version
-    return f"Schema for DataFrame 'df_{df_name}': {df.columns.tolist()[:5]}...\nSample:\n{df.head(1).to_string()}\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 (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"
+    return schema_str
+
 def get_all_schemas_representation(dataframes_dict: dict) -> str:
-    # Truncated for brevity in example, keep your full version
-    return "".join(get_schema_representation(name, df) for name, df in dataframes_dict.items() if isinstance(df, pd.DataFrame))
+    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: # Truncated for brevity, assume correct from previous versions
+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 
-        if GEMINI_API_KEY and hasattr(genai, 'embed_content') and not (hasattr(genai.embed_content, '__func__') and genai.embed_content.__func__.__qualname__.startswith('genai.embed_content')):
+        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
+
+        if GEMINI_API_KEY and self.client_available:
             try:
-                self._precompute_embeddings() 
+                self._precompute_embeddings()
                 self.embeddings_generated = True
-                logging.info("RAG embeddings precomputed.")
-            except Exception as e: logging.error(f"RAG precomputation error: {e}")
+                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)
         else:
-            logging.warning("RAG embeddings not precomputed (API key or genai.embed_content issue).")
+            logging.warning(f"RAG embeddings not precomputed. GEMINI_API_KEY set: {bool(GEMINI_API_KEY)}, genai.embed_content available: {self.client_available}.")
 
-    def _embed_fn(self, title: str, text: str) -> list[float]: 
-        if not self.embeddings_generated: return [0.0] * 768
+    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
         try:
-            return genai.embed_content(model=self.embedding_model_name, content=text, task_type="retrieval_document", title=title)["embedding"]
+            # 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"]
         except Exception as e:
-            logging.error(f"Error in _embed_fn for '{title}': {e}")
+            logging.error(f"Error in _embed_fn for '{title}': {e}", exc_info=True)
             return [0.0] * 768
-    def _precompute_embeddings(self): 
+
+    def _precompute_embeddings(self):
         if 'Embeddings' not in self.documents_df.columns:
             self.documents_df['Embeddings'] = pd.Series(dtype='object')
-        self.documents_df['Embeddings'] = self.documents_df.apply(lambda row: self._embed_fn(row['Title'], row['Text']), axis=1)
-    def retrieve_relevant_info(self, query_text: str, top_k: int = 1) -> str: 
-        if not self.embeddings_generated or self.documents_df['Embeddings'].isnull().all():
-            return "\n[RAG Context]\nEmbeddings not generated or all are null.\n"
-        # Simplified retrieval logic for brevity
+        
+        # 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 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"
+
         try:
-            query_embedding = np.array(genai.embed_content(model=self.embedding_model_name, content=query_text, task_type="retrieval_query")["embedding"])
-            # Filter out rows with invalid embeddings before stacking
+            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) and len(x) > 0)]
+            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"
 
-            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]: return "\n[RAG Context]\nEmbedding dimension mismatch.\n"
+
+            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)
-            idx = np.argsort(dot_products)[-min(top_k, len(valid_embeddings_df)):][::-1]
-            
-            relevant_passages = "".join([f"\n[RAG Context from: '{valid_embeddings_df.iloc[i]['Title']}']\n{valid_embeddings_df.iloc[i]['Text']}\n" for i in idx])
-            return relevant_passages if relevant_passages else "\n[RAG Context]\nNo relevant passages found.\n"
+            # 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"
         except Exception as e:
-            logging.error(f"Error in RAG retrieve_relevant_info: {e}")
-            return f"\n[RAG Context]\nError during RAG retrieval: {e}\n"
+            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) ---
 class PandasLLM:
-    def __init__(self, llm_model_name: str, 
-                 generation_config_dict: dict, 
-                 safety_settings_list: list,   
+    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
+        self.safety_settings_list = safety_settings_list
         self.data_privacy = data_privacy
         self.force_sandbox = force_sandbox
-        self.client = None
-        self.generative_model_service = None
+        self.generative_model = None # Will hold the GenerativeModel instance
 
         if not GEMINI_API_KEY:
-            logging.warning("PandasLLM: GEMINI_API_KEY not set.")
-        else:
+            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:
-                self.client = genai.Client(api_key=GEMINI_API_KEY) 
-                if self.client and hasattr(self.client, 'models') and hasattr(self.client.models, 'generate_content'):
-                    self.generative_model_service = self.client.models
-                    logging.info(f"PandasLLM: Using client.models for '{self.llm_model_name}'.")
-                elif self.client and hasattr(self.client, 'generate_content'):
-                    self.generative_model_service = self.client
-                    logging.info(f"PandasLLM: Using client.generate_content for '{self.llm_model_name}'.")
-                else:
-                    logging.warning(f"PandasLLM: genai.Client suitable 'generate_content' not found.")
+                # 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.Client: {e}", exc_info=True)
+                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.")
 
-    async def _call_gemini_api_async(self, prompt_text: str, history: list = None) -> str:
-        if not self.generative_model_service:
-            return "# Error: Gemini client/service not available."
 
-        contents_for_api = []
+    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.
+        
+        # Convert simple history to Gemini's expected format
+        gemini_history = []
         if history:
             for entry in history:
                 role = "model" if entry.get("role") == "assistant" else entry.get("role", "user")
-                contents_for_api.append({"role": role, "parts": [{"text": entry.get("content", "")}]})
-        contents_for_api.append({"role": "user", "parts": [{"text": prompt_text}]})
+                gemini_history.append({"role": role, "parts": [{"text": entry.get("content", "")}]})
         
-        api_config_object = None
-        try:
-            api_config_object = genai_types.GenerateContentConfig(
-                **self.generation_config_dict, 
-                safety_settings=self.safety_settings_list 
-            )
-        except Exception as e_cfg:
-            logging.error(f"Error creating GenerateContentConfig object: {e_cfg}.")
-            api_config_object = {**self.generation_config_dict, "safety_settings": self.safety_settings_list}
-
-        logging.info(f"\n--- Calling Gemini API via Client (model: {self.llm_model_name}) with config: {api_config_object} ---\n")
+        # Add current prompt as the last user message
+        current_content = [{"role": "user", "parts": [{"text": prompt_text}]}]
+        
+        # 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")
         
         try:
-            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}"
-            
-            response = await asyncio.to_thread(
-                self.generative_model_service.generate_content,
-                model=model_id_for_api,
+            # The GenerativeModel instance already has config and safety settings.
+            # We just pass the 'contents'.
+            response = await self.generative_model.generate_content_async(
                 contents=contents_for_api,
-                config=api_config_object 
+                # generation_config, safety_settings are already part of self.generative_model
             )
 
-            if hasattr(response, 'prompt_feedback') and response.prompt_feedback and response.prompt_feedback.block_reason:
-                return f"# Error: Prompt blocked by API: {response.prompt_feedback.block_reason}."
+            if hasattr(response, 'prompt_feedback') and response.prompt_feedback and \
+               hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason:
+                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}")
+                return f"# Error: Prompt blocked by API. Reason: {block_reason_str}."
 
             llm_output = ""
-            if hasattr(response, 'text') and response.text:
+            # Standard way to get text from Gemini response
+            if hasattr(response, 'text') and isinstance(response.text, str):
                 llm_output = response.text
-            elif hasattr(response, 'candidates') and response.candidates: 
+            elif response.candidates:
                 candidate = response.candidates[0]
-                if hasattr(candidate, 'content') and candidate.content and hasattr(candidate.content, 'parts') and candidate.content.parts:
+                if candidate.content and candidate.content.parts:
                     llm_output = "".join(part.text for part in candidate.content.parts if hasattr(part, 'text'))
-                if not llm_output and hasattr(candidate, 'finish_reason'):
-                     return f"# Error: Empty response. Finish reason: {candidate.finish_reason}."
+                
+                if not llm_output and candidate.finish_reason:
+                    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
+                    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)}"
+                    
+                    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}."
             else:
+                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 Exception as e:
-            logging.error(f"Error calling Gemini API via Client: {e}", exc_info=True)
+            logging.error(f"Error calling Gemini API: {e}", exc_info=True)
             return f"# Error during API call: {type(e).__name__} - {str(e)[:100]}."
 
+
     async def query(self, prompt_with_query_and_context: str, dataframes_dict: dict, history: list = None) -> str:
         llm_response_text = await self._call_gemini_api_async(prompt_with_query_and_context, history)
+
         if self.force_sandbox:
             code_to_execute = ""
+            # Robust code extraction
             if "```python" in llm_response_text:
                 try:
-                    code_to_execute = llm_response_text.split("```python\n", 1)[1].split("\n```", 1)[0]
-                except IndexError: 
-                    try:
-                        code_to_execute = llm_response_text.split("```python", 1)[1].split("```", 1)[0]
-                        if code_to_execute.startswith("\n"): code_to_execute = code_to_execute[1:]
-                        if code_to_execute.endswith("\n"): code_to_execute = code_to_execute[:-1]
-                    except IndexError: code_to_execute = ""
-            
-            if llm_response_text.startswith("# Error:") or not code_to_execute:
+                    # 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
+                        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 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.
-                # The user will see the LLM's error message or its non-code response.
-                logging.warning(f"LLM response is an error or not code: {llm_response_text}")
-                return llm_response_text 
-
-            logging.info(f"\n--- Code to Execute: ---\n{code_to_execute}\n----------------------\n")
+                # Or if it's a polite non-code refusal (e.g. "# Hello there! ...")
+                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")
+            
             from io import StringIO
             import sys
-            old_stdout = sys.stdout; sys.stdout = captured_output = StringIO()
-            # Ensure dataframes_dict is correctly populated for exec_globals
-            exec_globals = {'pd': pd, 'np': np} 
-            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.")
-
+            old_stdout = sys.stdout
+            sys.stdout = captured_output = StringIO()
+            
+            # Prepare globals for exec. Prefix DataFrames with 'df_' as per prompt.
+            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.")
+            
             try:
-                exec(code_to_execute, exec_globals, {})
+                exec(code_to_execute, exec_globals, {}) # Using empty dict for locals
                 final_output_str = captured_output.getvalue()
-                # Check if the output is just whitespace or truly empty
-                if not final_output_str.strip(): # If only whitespace or empty
-                    # This is where the "no print output" message originates.
-                    # We can now add a more informative message if the code itself ran without error.
+                
+                if not final_output_str.strip():
                     logging.info("Code executed successfully, but no explicit print() output was generated by the LLM's code.")
-                    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."
+                    # 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."
                 return final_output_str
             except Exception as e:
-                logging.error(f"Sandbox Execution Error: {e}\nCode was:\n{code_to_execute}", exc_info=False)
-                return f"# Sandbox Execution Error: {type(e).__name__}: {e}\n# --- Code that caused error: ---\n{textwrap.indent(code_to_execute, '# ')}"
-            finally: 
+                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: 
+        else: # Not force_sandbox
             return llm_response_text
 
 # --- Employer Branding Agent ---
 class EmployerBrandingAgent:
-    def __init__(self, llm_model_name: str, 
-                 generation_config_dict: dict, 
-                 safety_settings_list: list,   
-                 all_dataframes: dict, 
-                 rag_documents_df: pd.DataFrame, 
+    def __init__(self, llm_model_name: str,
+                 generation_config_dict: dict,
+                 safety_settings_list: list,
+                 all_dataframes: dict,
+                 rag_documents_df: pd.DataFrame,
                  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, 
+            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 
+        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 with updated safety settings handling.")
+        self.chat_history = []
+        logging.info("EmployerBrandingAgent Initialized.")
 
-    def _build_prompt(self, user_query: str, role="Employer Branding Analyst", task_decomposition_hint=None, cot_hint=True) -> str:
-        prompt = f"You are a helpful and expert '{role}'. Your primary goal is to assist with analyzing LinkedIn-related data using Pandas DataFrames.\n"
+    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: Be mindful of data privacy. Do not output raw Personally Identifiable Information (PII) like names or specific user details unless explicitly asked and absolutely necessary for the query. Summarize or aggregate data where possible.\n"
+            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"
 
         if self.pandas_llm.force_sandbox:
-            prompt += "Your main task is to GENERATE PYTHON CODE using the Pandas library to answer the user query based on the provided DataFrames. Output ONLY the Python code block.\n"
+            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 += "CRITICAL INSTRUCTION: Your Python code MUST include `print()` statements for ANY results, DataFrames, or values that should be displayed as the answer to the user's query. The output of these `print()` statements will be the final answer shown to the user.\n"
-            prompt += "If you define a function to perform the analysis, you MUST call this function with the appropriate DataFrame(s) and `print()` its returned value. Do not just define functions without executing them and printing their results.\n"
-            prompt += "If the query is simple and the result is a single value or a small piece of information, compute it and `print()` it directly.\n"
-            prompt += "For example, if asked for 'total followers', your code should end with something like `print(total_followers)` or `print(df_result.to_string())`.\n"
-
-            prompt += "If a column contains lists (e.g., 'skills' in a hypothetical 'df_employees'), you might need to use methods like `.explode()` or `.apply(pd.Series)` or `.apply(lambda x: ...)` for analysis.\n"
-            prompt += "If the query is ambiguous or requires clarification, ask for it instead of making assumptions. If the query cannot be answered with the given data, state that clearly in a comment within the code block (e.g. `# Cannot answer: data not available`).\n"
-            prompt += "If the query is not about data analysis or code generation (e.g. 'hello', 'how are you?'), respond politely and briefly in a comment, do not attempt to generate code (e.g. `# Hello there! How can I help you with data analysis today?`).\n"
-            prompt += "Structure your code clearly. Add comments (#) to explain each step of your logic.\n"
-        else: 
-            prompt += "Your task is to analyze the data and provide a comprehensive textual answer to the user query. You can explain your reasoning step-by-step.\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 += "\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 += "\n--- AVAILABLE DATA AND SCHEMAS ---\n"
-        prompt += self.schemas_representation
+        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
         
         rag_context = self.rag_system.retrieve_relevant_info(user_query)
-        if rag_context and "[RAG Context]" in rag_context and "No specific pre-defined context found" not in rag_context and "No highly relevant passages found" not in rag_context and "Embeddings not generated" not in rag_context:
-            prompt += f"\n--- ADDITIONAL CONTEXT (from internal knowledge base, consider this information) ---\n{rag_context}\n"
+        # 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"
+
 
         prompt += f"\n--- USER QUERY ---\n{user_query}\n"
 
@@ -406,40 +704,163 @@ class EmployerBrandingAgent:
         
         if cot_hint:
             if self.pandas_llm.force_sandbox:
-                prompt += "\n--- INSTRUCTIONS FOR PYTHON CODE GENERATION (Chain of Thought & Output) ---\n"
-                prompt += "1. Understand the query: What specific information is requested?\n"
-                prompt += "2. Identify relevant DataFrame(s) and column(s) from the schemas provided.\n"
-                prompt += "3. Plan the steps: Outline the Pandas operations needed (filtering, grouping, aggregation, merging, etc.) as comments in your code.\n"
-                prompt += "4. Write the code: Implement the steps using Pandas. Remember to use `df_name_of_dataframe` (e.g. `df_follower_stats`).\n"
-                prompt += "5. CRITICAL - Ensure output: Call any functions you define and use `print()` for ALL results that should be displayed. For DataFrames, you can print the DataFrame directly (e.g., `print(my_result_df)`), or `print(df.to_string())` if it might be large. For single values, `print(my_value)`.\n"
-                prompt += "6. Review: Check for correctness, efficiency, and adherence to the prompt (especially the CRITICAL `print()` requirement for the final answer).\n"
-                prompt += "7. Generate ONLY the Python code block starting with ```python and ending with ```. No explanations outside the code block's comments.\n"
-            else: 
-                prompt += "\n--- INSTRUCTIONS FOR RESPONSE (Chain of Thought) ---\n"
-                prompt += "Please provide a step-by-step explanation of your analysis before giving the final answer.\n"
+                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"
+            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"
         
         return prompt
 
-    async def process_query(self, user_query: str, role="Employer Branding Analyst", task_decomposition_hint=None, cot_hint=True) -> str:
-        self.chat_history.append({"role": "user", "content": user_query})
+    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
+        
         full_prompt = self._build_prompt(user_query, role, task_decomposition_hint, cot_hint)
         
-        logging.info(f"Full prompt to LLM (last 300 chars of user query part for brevity in log): ... {full_prompt[-500:]}") # Log end of prompt
+        # 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]}...")
 
-        response_text = await self.pandas_llm.query(full_prompt, self.all_dataframes, history=self.chat_history[:-1]) 
-        self.chat_history.append({"role": "assistant", "content": response_text})
+
+        # Pass the history *before* the current user query to the LLM
+        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 
+        MAX_HISTORY_TURNS = 5 # Each turn has a user and a model message
         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.")
             
         return response_text
 
-    def update_dataframes(self, new_dataframes: dict): 
-        self.all_dataframes = new_dataframes
+    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("EmployerBrandingAgent DataFrames updated.")
-    def clear_chat_history(self): 
+        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.
+
+    def clear_chat_history(self):
         self.chat_history = []
         logging.info("EmployerBrandingAgent chat history cleared.")
 
+# --- Example Usage (Conceptual - for testing the module structure) ---
+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."
+    ]
+
+    for query in queries:
+        logging.info(f"\n\n--- Processing 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")
+
+
+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())
+        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.")