Update eb_agent_module.py

eb_agent_module.py

@@ -25,14 +25,21 @@ class _DummyGenAIClientModels: # Represents the dummy model service client
         class DummyResponse: candidates = [DummyCandidate()]; text = DummyCandidate.content.parts[0].text; prompt_feedback = None
         return DummyResponse()
 
+    def embed_content(self, model=None, contents=None, config=None): # Added dummy embed_content
+        print(f"Dummy _DummyGenAI.Client.models.embed_content called for model: {model}, task_type (from config): {config.get('task_type') if isinstance(config, dict) else 'N/A'}")
+        return {"embedding": [0.2] * 768} # Different values for dummy distinction
+
+
 class _DummyGenAIClient: # Dummy Client
-    def __init__(self, api_key=None):
-        self.api_key = api_key
+    def __init__(self, client_options=None): # Added client_options for signature consistency
+        self.client_options = client_options
         self.models = _DummyGenAIClientModels()
-        print(f"Dummy _DummyGenAI.Client initialized {'with api_key' if api_key else '(global API key expected)'}.")
+        api_key_present_in_options = client_options and client_options.get("api_key")
+        print(f"Dummy _DummyGenAI.Client initialized {'with api_key in client_options' if api_key_present_in_options else '(global API key expected by dummy)'}.")
+
 
-class _DummyGenAIGenerativeModel:
-    def __init__(self, model_name_in, generation_config_in, safety_settings_in, system_instruction_in):
+class _DummyGenAIGenerativeModel: # This dummy might be less used if client.models is preferred
+    def __init__(self, model_name_in, generation_config=None, safety_settings=None, system_instruction=None):
         self.model_name = model_name_in
         print(f"Dummy _DummyGenAIGenerativeModel initialized for {model_name_in}")
     async def generate_content_async(self, contents, stream=False):
@@ -42,8 +49,13 @@ class _DummyGenAIGenerativeModel:
         class DummyCandidate: content = DummyContent(); finish_reason = "_DUMMY_STOP"; safety_ratings = []
         class DummyResponse: candidates = [DummyCandidate()]; prompt_feedback = None; text = DummyCandidate.content.parts[0].text
         return DummyResponse()
+    # This embed_content on the dummy GenerativeModel might not be used if AdvancedRAGSystem uses client.models.embed_content
+    def embed_content(self, content, task_type=None, title=None):
+        print(f"Dummy _DummyGenAIGenerativeModel.embed_content called for model {self.model_name} (task: {task_type})")
+        return {"embedding": [0.1] * 768}
+
 
-class _ActualDummyGenAI: # type: ignore # Renamed the main dummy class
+class _ActualDummyGenAI: # type: ignore
     Client = _DummyGenAIClient 
 
     @staticmethod
@@ -54,13 +66,7 @@ class _ActualDummyGenAI: # type: ignore # Renamed the main dummy class
     def GenerativeModel(model_name, generation_config=None, safety_settings=None, system_instruction=None):
         print(f"Dummy _ActualDummyGenAI.GenerativeModel called for model: {model_name}")
         return _DummyGenAIGenerativeModel(model_name, generation_config, safety_settings, system_instruction)
-
-    @staticmethod
-    def embed_content(model, content, task_type, title=None):
-        print(f"Dummy _ActualDummyGenAI.embed_content called for model: {model}, task_type: {task_type}, title: {title}")
-        return {"embedding": [0.1] * 768}
     
-    # Add a dummy 'types' attribute to the dummy genai class
     class types:
         @staticmethod
         def GenerationConfig(**kwargs):
@@ -71,102 +77,89 @@ class _ActualDummyGenAI: # type: ignore # Renamed the main dummy class
         def SafetySetting(category, threshold):
             print(f"Dummy _ActualDummyGenAI.types.SafetySetting created: category={category}, threshold={threshold}")
             return {"category": category, "threshold": threshold}
-
-        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:
-            FINISH_REASON_UNSPECIFIED = "UNSPECIFIED"; STOP = "STOP"; MAX_TOKENS = "MAX_TOKENS"; SAFETY = "SAFETY"; RECITATION = "RECITATION"; OTHER = "OTHER"
-        class BlockedReason:
-            BLOCKED_REASON_UNSPECIFIED = "BLOCKED_REASON_UNSPECIFIED"; SAFETY = "SAFETY"; OTHER = "OTHER"
-        # Add other dummy types if needed by the script, e.g. BlockedPromptException
+        
+        @staticmethod # Added dummy EmbedContentConfig
+        def EmbedContentConfig(task_type=None, output_dimensionality=None, title=None):
+            print(f"Dummy _ActualDummyGenAI.types.EmbedContentConfig created with task_type: {task_type}")
+            conf = {}
+            if task_type: conf["task_type"] = task_type
+            if output_dimensionality: conf["output_dimensionality"] = output_dimensionality
+            if title: conf["title"] = title # Though title is usually direct param for embed_content
+            return conf
+
+
+        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: FINISH_REASON_UNSPECIFIED = "UNSPECIFIED"; STOP = "STOP"; MAX_TOKENS = "MAX_TOKENS"; SAFETY = "SAFETY"; RECITATION = "RECITATION"; OTHER = "OTHER"
+        class BlockedReason: BLOCKED_REASON_UNSPECIFIED = "BLOCKED_REASON_UNSPECIFIED"; SAFETY = "SAFETY"; OTHER = "OTHER"
         class BlockedPromptException(Exception): pass
         class StopCandidateException(Exception): pass
 
 
 # --- Attempt to import the real library ---
 _REAL_GENAI_LOADED = False
-genai_types = None # Initialize genai_types
+genai_types = None 
 
 try:
-    from google import genai # Use the direct import as requested
-    # If 'from google import genai' succeeds, 'genai.types' should be available
-    genai_types = genai.types # Assign the real types
+    from google import genai 
+    genai_types = genai.types 
     _REAL_GENAI_LOADED = True
     logging.info("Successfully imported 'google.genai' and accessed 'genai.types'.")
 except ImportError:
-    # If 'from google import genai' fails, use the dummy genai and its dummy types
     genai = _ActualDummyGenAI() 
-    genai_types = genai.types # This will now point to _ActualDummyGenAI.types
+    genai_types = genai.types 
     logging.warning("Google AI library ('google.genai') not found. Using dummy implementations for 'genai' and 'genai_types'.")
-except AttributeError:
-    # This handles the case where 'from google import genai' succeeds, but 'genai.types' is not found
-    # (which would be unusual for the official library but good for robustness)
-    genai = _ActualDummyGenAI() # Fallback to full dummy
-    genai_types = genai.types
-    _REAL_GENAI_LOADED = False # Mark as not fully loaded if types are missing
+except AttributeError: # If 'genai' imported but 'genai.types' is missing
+    genai = _ActualDummyGenAI() 
+    genai_types = genai.types # Fallback to dummy types
+    _REAL_GENAI_LOADED = False 
     logging.warning("'google.genai' imported, but 'genai.types' not found. Falling back to dummy implementations.")
 
 
 # --- Configuration ---
 GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', "")
 LLM_MODEL_NAME = "gemini-2.0-flash" 
-GEMINI_EMBEDDING_MODEL_NAME = "gemini-embedding-exp-03-07"
+GEMINI_EMBEDDING_MODEL_NAME = "gemini-embedding-exp-03-07" 
 
 GENERATION_CONFIG_PARAMS = {
     "temperature": 0.3, "top_p": 1.0, "top_k": 32, "max_output_tokens": 8192,
 }
 
-# Default safety settings list for Gemini
-# genai_types is now consistently the real genai.types or the dummy _ActualDummyGenAI.types
 try:
     DEFAULT_SAFETY_SETTINGS = [
         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),
+        # ... other settings
     ]
 except Exception as e_safety: 
     logging.warning(f"Could not define DEFAULT_SAFETY_SETTINGS using 'genai_types' (real_loaded: {_REAL_GENAI_LOADED}): {e_safety}. Using placeholder list of dicts.")
-    DEFAULT_SAFETY_SETTINGS = [
-        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-    ]
-
+    DEFAULT_SAFETY_SETTINGS = [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}] # Simplified
 
-# Logging setup
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(filename)s:%(lineno)d - %(message)s')
 
-if GEMINI_API_KEY and _REAL_GENAI_LOADED:
-    try:
-        genai.configure(api_key=GEMINI_API_KEY) 
-        logging.info(f"Gemini API key configured globally (real genai active: {_REAL_GENAI_LOADED}).")
-    except Exception as e:
-        logging.error(f"Failed to configure Gemini API globally: {e}", exc_info=True)
-elif not GEMINI_API_KEY and _REAL_GENAI_LOADED:
-    logging.warning("GEMINI_API_KEY environment variable not set, but real library is loaded. API calls will likely fail or use default credentials if available.")
+if _REAL_GENAI_LOADED:
+    if GEMINI_API_KEY:
+        try:
+            genai.configure(api_key=GEMINI_API_KEY) 
+            logging.info(f"Gemini API key configured globally using REAL genai.configure.")
+        except Exception as e:
+            logging.error(f"Failed to configure REAL Gemini API globally: {e}", exc_info=True)
+    else:
+        logging.warning("REAL 'google.genai' loaded, but GEMINI_API_KEY not set. API calls might fail or use other auth.")
 elif not _REAL_GENAI_LOADED:
-    logging.info("Operating in DUMMY mode because 'google.genai' library was not found or 'genai.types' was inaccessible.")
-    if GEMINI_API_KEY: 
-        genai.configure(api_key=GEMINI_API_KEY) # Calls dummy configure
+    logging.info("Operating in DUMMY mode for 'google.genai'.")
+    if GEMINI_API_KEY: genai.configure(api_key=GEMINI_API_KEY)
 
 
 # --- RAG Documents Definition (Example) ---
-rag_documents_data = {
-    '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..."]
-}
+rag_documents_data = { 'Title': ["EB Practices", "Tech Talent"], 'Text': ["Stories...", "Projects..."] }
 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.\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
+    return f"DataFrame 'df_{df_name}': Cols: {df.columns.tolist()}, Shape: {df.shape}\nSample:\n{textwrap.indent(df.head(1).to_string(), '    ')}\n"
 
 def get_all_schemas_representation(dataframes_dict: dict) -> str:
     if not dataframes_dict: return "No DataFrames provided.\n"
@@ -175,50 +168,77 @@ def get_all_schemas_representation(dataframes_dict: dict) -> str:
 # --- Advanced RAG System ---
 class AdvancedRAGSystem:
     def __init__(self, documents_df: pd.DataFrame, embedding_model_name: str):
-        self.embedding_model_name = embedding_model_name
+        self.embedding_model_name_for_api = embedding_model_name # Store raw name
+        if not self.embedding_model_name_for_api.startswith("models/"):
+            self.embedding_model_name_for_api = f"models/{self.embedding_model_name_for_api}"
+        
         self.documents_df = documents_df.copy()
         self.embeddings_generated = False
+        self.embedding_service = None # Will hold client.models or its dummy equivalent
+
         self.real_client_available_for_rag = _REAL_GENAI_LOADED and bool(GEMINI_API_KEY)
 
         if self.real_client_available_for_rag:
             try:
+                # Pass client_options if API key is available, to help Client find it
+                client_opts = {"api_key": GEMINI_API_KEY} if GEMINI_API_KEY else None
+                rag_client = genai.Client(client_options=client_opts)
+                self.embedding_service = rag_client.models
+                logging.info(f"RAG: REAL embedding service (genai.Client.models) initialized for '{self.embedding_model_name_for_api}'.")
                 self._precompute_embeddings()
                 self.embeddings_generated = True
-                logging.info(f"RAG embeddings precomputed using REAL genai.embed_content for '{self.embedding_model_name}'.")
-            except Exception as e: logging.error(f"RAG precomputation error with real client: {e}", exc_info=True)
+            except Exception as e: 
+                logging.error(f"RAG: Error initializing REAL embedding service: {e}", exc_info=True)
+                self.embedding_service = None 
         else:
-            logging.warning(f"RAG: Not using real client. Real GenAI loaded: {_REAL_GENAI_LOADED}, API Key set: {bool(GEMINI_API_KEY)}.")
-            if not _REAL_GENAI_LOADED: # If in full dummy mode
-                self._precompute_embeddings() # This will call dummy genai.embed_content via _embed_fn
+            logging.warning(f"RAG: Not using REAL embedding service. Real GenAI: {_REAL_GENAI_LOADED}, API Key: {bool(GEMINI_API_KEY)}.")
+            if not _REAL_GENAI_LOADED: # Full dummy mode
+                self.embedding_service = genai.Client().models # genai is _ActualDummyGenAI, gets dummy service
+                self._precompute_embeddings()
 
-    def _embed_fn(self, title: str, text: str) -> list[float]:
+    def _embed_fn(self, contents_to_embed: str, task_type: str) -> list[float]:
+        if not self.embedding_service:
+            logging.error(f"RAG _embed_fn: Embedding service not available for model '{self.embedding_model_name_for_api}'.")
+            return [0.0] * 768
         try:
-            content_to_embed = text if text else title
-            if not content_to_embed: return [0.0] * 768
-            # genai.embed_content will be real or dummy based on import success
-            return genai.embed_content(model=self.embedding_model_name, content=content_to_embed, task_type="retrieval_document", title=title if title else None)["embedding"]
+            if not contents_to_embed: return [0.0] * 768
+            
+            # Use genai_types (which is real or dummy) to create EmbedContentConfig
+            embed_config = genai_types.EmbedContentConfig(task_type=task_type)
+            
+            # Call embed_content on the service (real or dummy)
+            response = self.embedding_service.embed_content(
+                model=self.embedding_model_name_for_api,
+                contents=contents_to_embed, 
+                config=embed_config
+            )
+            return response["embedding"]
         except Exception as e:
-            logging.error(f"Error in _embed_fn for '{title}' (real_genai_loaded: {_REAL_GENAI_LOADED}): {e}", exc_info=True)
+            logging.error(f"Error in _embed_fn for task '{task_type}' using model '{self.embedding_model_name_for_api}' (real_genai_loaded: {_REAL_GENAI_LOADED}): {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')
         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 (real_genai_loaded: {_REAL_GENAI_LOADED}).")
+        
+        for index, row in self.documents_df[mask].iterrows():
+            text_to_embed = row.get('Text', '') if row.get('Text', '') else row.get('Title', '')
+            self.documents_df.loc[index, 'Embeddings'] = self._embed_fn(text_to_embed, task_type="RETRIEVAL_DOCUMENT") # Corrected task type string
+        
+        logging.info(f"Applied RAG embedding function to {mask.sum()} rows (embedding_service active: {self.embedding_service is not None}).")
 
 
     def retrieve_relevant_info(self, query_text: str, top_k: int = 2) -> str:
-        if not self.real_client_available_for_rag:
-            if not _REAL_GENAI_LOADED: # If in full dummy mode, make the dummy call for logging
-                 genai.embed_content(model=self.embedding_model_name, content=query_text, task_type="retrieval_query") 
-            logging.warning(f"Skipping real RAG retrieval. Real GenAI: {_REAL_GENAI_LOADED}, API Key: {bool(GEMINI_API_KEY)}")
-            return "\n[RAG Context]\nReal RAG retrieval skipped (check logs for mode).\n"
+        if not self.real_client_available_for_rag or not self.embedding_service:
+            if not _REAL_GENAI_LOADED and self.embedding_service: # Full dummy mode
+                 self._embed_fn(query_text, task_type="RETRIEVAL_QUERY") # Call for dummy log
+            logging.warning(f"Skipping real RAG retrieval. Real client available: {self.real_client_available_for_rag}, Embedding service OK: {self.embedding_service is not None}")
+            return "\n[RAG Context]\nReal RAG retrieval skipped.\n"
 
         try:
-            # genai.embed_content is the real one here
-            query_embedding = np.array(genai.embed_content(model=self.embedding_model_name, content=query_text, task_type="retrieval_query")["embedding"])
+            query_embedding = np.array(self._embed_fn(query_text, task_type="RETRIEVAL_QUERY")) # Corrected task type string
+            
             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 and np.any(x))] 
             if valid_df.empty: return "\n[RAG Context]\nNo valid document embeddings after filtering.\n"
@@ -234,7 +254,7 @@ class AdvancedRAGSystem:
             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 (real mode): {e}", exc_info=True)
+            logging.error(f"Error in RAG retrieve_relevant_info (real mode with embedding service): {e}", exc_info=True)
             return f"\n[RAG Context]\nError during RAG retrieval (real mode): {type(e).__name__} - {e}\n"
 
 # --- PandasLLM Class (Gemini-Powered using genai.Client) ---
@@ -248,21 +268,24 @@ class PandasLLM:
         self.safety_settings_list = safety_settings_list
         self.data_privacy = data_privacy
         self.force_sandbox = force_sandbox
-        self.client = None
-        self.model_service = None
+        self.client = None 
+        self.model_service = None 
 
         if _REAL_GENAI_LOADED and GEMINI_API_KEY:
             try:
-                self.client = genai.Client() # Real genai.Client
+                # genai.configure should have been called. Try passing client_options as a fallback.
+                client_opts = {"api_key": GEMINI_API_KEY} if GEMINI_API_KEY else None
+                self.client = genai.Client(client_options=client_opts) 
                 self.model_service = self.client.models
                 logging.info(f"PandasLLM: Initialized with REAL genai.Client().models for '{self.llm_model_name}'.")
             except Exception as e:
                 logging.error(f"Failed to initialize REAL PandasLLM with genai.Client: {e}", exc_info=True)
-                # If this fails, self.model_service remains None, _call_gemini_api_async might fallback if _REAL_GENAI_LOADED somehow becomes False
+                self.client = None 
+                self.model_service = None
         else:
             logging.warning(f"PandasLLM: Not using REAL genai.Client. RealGenAILoaded: {_REAL_GENAI_LOADED}, APIKeySet: {bool(GEMINI_API_KEY)}.")
-            if not _REAL_GENAI_LOADED: # If import failed, genai is already the dummy
-                self.client = genai.Client() # Instantiates _ActualDummyGenAI.Client
+            if not _REAL_GENAI_LOADED: 
+                self.client = genai.Client() 
                 self.model_service = self.client.models 
                 logging.info("PandasLLM: Initialized with DUMMY genai.Client().models (real library failed to load).")
 
@@ -271,13 +294,13 @@ class PandasLLM:
         use_real_service = _REAL_GENAI_LOADED and GEMINI_API_KEY and self.model_service is not None
         
         active_model_service = self.model_service
-        if not use_real_service and not _REAL_GENAI_LOADED: # Full dummy mode
-            if active_model_service is None: # Should have been set by __init__
+        if not use_real_service and not _REAL_GENAI_LOADED: 
+            if active_model_service is None: 
                 logging.debug("PandasLLM._call_gemini_api_async: active_model_service is None in dummy mode, using global dummy genai.Client().models.")
-                active_model_service = genai.Client().models # genai is _ActualDummyGenAI here
+                active_model_service = genai.Client().models
 
         if not active_model_service:
-             logging.error(f"PandasLLM: Model service not available (use_real_service: {use_real_service}, _REAL_GENAI_LOADED: {_REAL_GENAI_LOADED}). Cannot call API.")
+             logging.error(f"PandasLLM: Model service not available (use_real_service: {use_real_service}, _REAL_GENAI_LOADED: {_REAL_GENAI_LOADED}, self.model_service is None: {self.model_service is None}). Cannot call API.")
              return "# Error: Gemini model service not available for API call."
 
         gemini_history = []
@@ -291,7 +314,7 @@ class PandasLLM:
         contents_for_api = gemini_history + current_prompt_content
         
         model_id_for_api = self.llm_model_name
-        if not model_id_for_api.startswith("models/"):
+        if not model_id_for_api.startswith("models/"): 
             model_id_for_api = f"models/{model_id_for_api}"
 
         api_generation_config = None
@@ -349,20 +372,16 @@ class PandasLLM:
             
             return llm_output
         
-        # Use genai_types for exceptions if real library is loaded
         except (genai_types.BlockedPromptException if _REAL_GENAI_LOADED and hasattr(genai_types, 'BlockedPromptException') else Exception) as bpe:
-            if _REAL_GENAI_LOADED and type(bpe).__name__ == 'BlockedPromptException': # Check specific type if real
+            if _REAL_GENAI_LOADED and type(bpe).__name__ == 'BlockedPromptException':
                  logging.error(f"Prompt blocked (BlockedPromptException): {bpe}", exc_info=True)
                  return f"# Error: Prompt blocked. Details: {bpe}"
-            # Fallthrough for general exception if not the specific type or in dummy mode
-            pass # Let the general Exception handler catch it or re-raise if needed
-
+            if not (_REAL_GENAI_LOADED and type(bpe).__name__ == 'BlockedPromptException'): raise 
         except (genai_types.StopCandidateException if _REAL_GENAI_LOADED and hasattr(genai_types, 'StopCandidateException') else Exception) as sce:
-            if _REAL_GENAI_LOADED and type(sce).__name__ == 'StopCandidateException': # Check specific type if real
+            if _REAL_GENAI_LOADED and type(sce).__name__ == 'StopCandidateException':
                 logging.error(f"Candidate stopped (StopCandidateException): {sce}", exc_info=True)
                 return f"# Error: Content generation stopped. Details: {sce}"
-            pass # Fallthrough
-
+            if not (_REAL_GENAI_LOADED and type(sce).__name__ == 'StopCandidateException'): raise
         except Exception as e:
             logging.error(f"Error calling Gemini API (RealMode: {use_real_service}): {e}", exc_info=True)
             return f"# Error during API call: {type(e).__name__} - {str(e)[:100]}."
@@ -392,142 +411,83 @@ class PandasLLM:
                 return llm_response_text
 
             logging.info(f"\n--- Code to Execute: ---\n{code_to_execute}\n----------------------\n")
-            from io import StringIO
-            import sys
+            from io import StringIO; import sys
             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}' not a DataFrame for sandbox exec.")
             try:
                 exec(code_to_execute, exec_globals, {})
                 final_output_str = sys.stdout.getvalue()
                 if not final_output_str.strip():
                     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 produced by sandbox."
-                    return "# Code executed successfully by sandbox, but it did not produce any printed output. Ensure print() for results."
+                        return "# LLM generated only comments or empty code. No output by sandbox."
+                    return "# Code executed by sandbox, 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)
-                indented_code = textwrap.indent(code_to_execute, '# ')
-                return f"# Sandbox Execution Error: {type(e).__name__}: {e}\n# --- Code that caused error: ---\n{indented_code}"
+                logging.error(f"Sandbox Execution Error: {e}\nCode:\n{code_to_execute}", exc_info=True)
+                return f"# Sandbox Exec Error: {type(e).__name__}: {e}\n# Code:\n{textwrap.indent(code_to_execute, '# ')}"
             finally: sys.stdout = old_stdout
         else: 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,
-                 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.rag_system = AdvancedRAGSystem(rag_documents_df, embedding_model_name)
-        self.all_dataframes = all_dataframes if all_dataframes else {}
+    def __init__(self, llm_model_name: str, gc_dict: dict, ss_list: list, all_dfs: dict, rag_df: pd.DataFrame, emb_m_name: str, dp=True, fs=True):
+        self.pandas_llm = PandasLLM(llm_model_name, gc_dict, ss_list, dp, fs)
+        self.rag_system = AdvancedRAGSystem(rag_df, emb_m_name)
+        self.all_dataframes = all_dfs if all_dfs else {}
         self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
         self.chat_history = [] 
         logging.info(f"EmployerBrandingAgent Initialized (Real GenAI Loaded: {_REAL_GENAI_LOADED}).")
 
-    def _build_prompt(self, user_query: str, role="Employer Branding Analyst & Strategist", task_decomposition_hint=None, cot_hint=True) -> str:
-        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"
-
+    def _build_prompt(self, user_query: str, role="EB Analyst", task_hint=None, cot=True) -> str:
+        prompt = f"You are '{role}'. Goal: insights from DataFrames & RAG.\n"
+        if self.pandas_llm.data_privacy: prompt += "PRIVACY: Summarize/aggregate PII.\n"
         if self.pandas_llm.force_sandbox:
-            prompt += "\n--- TASK: PYTHON CODE GENERATION FOR INSIGHTS ---\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:** `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 += "Analyze data and RAG, then provide a comprehensive textual answer with insights. Explain step-by-step.\n"
-
-        prompt += "\n--- AVAILABLE DATA AND SCHEMAS ---\n"
-        prompt += self.schemas_representation if self.schemas_representation.strip() != "No DataFrames provided." else "No DataFrames loaded.\n"
+            prompt += "TASK: PYTHON CODE. `print()` textual insights/answers. ```python ... ``` ONLY.\nAccess DFs as 'df_name'.\n"
+            prompt += "CRITICAL: `print()` insights, NOT raw DFs (unless asked). Synthesize RAG. Comment code. Handle issues (ambiguity, missing data) via `print()`.\n"
+        else: prompt += "TASK: TEXTUAL INSIGHTS. Explain step-by-step.\n"
+        prompt += f"--- DATA SCHEMAS ---\n{self.schemas_representation if self.schemas_representation.strip() != 'No DataFrames provided.' else 'No DFs loaded.'}\n"
         
         rag_context = self.rag_system.retrieve_relevant_info(user_query)
-        meaningful_rag_keywords = ["Error", "No valid", "No relevant", "Cannot retrieve", "not available", "not generated", "Skipped"]
-        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--- RAG CONTEXT (Real GenAI for RAG: {self.rag_system.real_client_available_for_rag}) ---\n{rag_context}\n"
-        else: 
-            prompt += f"\n--- RAG CONTEXT (Real GenAI for RAG: {self.rag_system.real_client_available_for_rag}) ---\nNo specific RAG context found, RAG error, or RAG skipped. Details: {rag_context[:100]}...\n"
-
-
-        prompt += f"\n--- USER QUERY ---\n{user_query}\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--- 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--- TEXTUAL RESPONSE THOUGHT PROCESS ---\n"
-                prompt += "1. Goal? 2. Data sources? 3. Formulate insights (data + RAG). 4. Structure: explanation, then insights.\n"
+        meaningful_rag_kws = ["Error", "No valid", "No relevant", "Cannot retrieve", "not available", "not generated", "Skipped"]
+        is_meaningful_rag = bool(rag_context.strip()) and not any(kw in rag_context for kw in meaningful_rag_kws)
+        prompt += f"--- RAG CONTEXT (Real RAG: {self.rag_system.real_client_available_for_rag}) ---\n{rag_context if is_meaningful_rag else f'No specific RAG context or RAG issue. Details: {rag_context[:70]}...'}\n"
+        prompt += f"--- USER QUERY ---\n{user_query}\n"
+        if task_hint: prompt += f"--- GUIDANCE ---\n{task_hint}\n"
+        if cot:
+            if self.pandas_llm.force_sandbox: prompt += "--- PYTHON THOUGHT PROCESS ---\n1.Goal? 2.Data? 3.Plan? 4.Code. 5.CRITICAL: `print()` insights. 6.Review. 7.```python ... ``` ONLY.\n"
+            else: prompt += "--- TEXT RESPONSE THOUGHT PROCESS ---\n1.Goal? 2.Data? 3.Insights (DFs+RAG). 4.Structure response.\n"
         return prompt
 
-    async def process_query(self, user_query: str, role="Employer Branding Analyst & Strategist", task_decomposition_hint=None, cot_hint=True) -> str:
-        current_turn_history_for_llm = self.chat_history[:]
+    async def process_query(self, user_query: str, role="EB Analyst", task_hint=None, cot=True) -> str:
+        hist_for_llm = self.chat_history[:]
         self.chat_history.append({"role": "user", "content": user_query})
-        full_prompt = self._build_prompt(user_query, role, task_decomposition_hint, cot_hint)
-        logging.info(f"Built prompt for query: {user_query[:100]}... (Real GenAI: {_REAL_GENAI_LOADED})")
-        response_text = await self.pandas_llm.query(full_prompt, self.all_dataframes, history=current_turn_history_for_llm)
-        self.chat_history.append({"role": "assistant", "content": response_text})
-        MAX_HISTORY_TURNS = 5
-        if len(self.chat_history) > MAX_HISTORY_TURNS * 2:
-            self.chat_history = self.chat_history[-(MAX_HISTORY_TURNS * 2):]
-            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"Agent DataFrames updated. Schemas: {self.schemas_representation[:100]}...")
-
+        prompt = self._build_prompt(user_query, role, task_hint, cot)
+        logging.info(f"Prompt for query: {user_query[:70]}... (Real GenAI: {_REAL_GENAI_LOADED})")
+        response = await self.pandas_llm.query(prompt, self.all_dataframes, history=hist_for_llm)
+        self.chat_history.append({"role": "assistant", "content": response})
+        if len(self.chat_history) > 10: self.chat_history = self.chat_history[-10:]; logging.info("Chat history truncated.")
+        return response
+
+    def update_dataframes(self, new_dfs: dict): self.all_dataframes = new_dfs if new_dfs else {}; self.schemas_representation = get_all_schemas_representation(self.all_dataframes); logging.info("Agent DFs updated.")
     def clear_chat_history(self): self.chat_history = []; logging.info("Agent chat history cleared.")
 
 # --- Example Usage (Conceptual) ---
 async def main_test():
-    logging.info(f"Starting main_test for EmployerBrandingAgent (Real GenAI Loaded: {_REAL_GENAI_LOADED}, API Key Set: {bool(GEMINI_API_KEY)})")
-    
-    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 and _REAL_GENAI_LOADED: 
-        logging.warning("GEMINI_API_KEY not set but real library loaded. Real API calls in test will fail.")
-    
-    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--- Query: {query} ---")
-        response = await agent.process_query(user_query=query)
-        logging.info(f"--- Response for '{query}': ---\n{response}\n---------------------------\n")
-        if _REAL_GENAI_LOADED and GEMINI_API_KEY: await asyncio.sleep(0.1) 
+    logging.info(f"Test (Real GenAI: {_REAL_GENAI_LOADED}, API Key: {bool(GEMINI_API_KEY)})")
+    agent = EmployerBrandingAgent(LLM_MODEL_NAME, GENERATION_CONFIG_PARAMS, DEFAULT_SAFETY_SETTINGS, {}, df_rag_documents, GEMINI_EMBEDDING_MODEL_NAME)
+    for q in ["What are EB best practices?", "Hello Agent!"]:
+        logging.info(f"\nQuery: {q}")
+        resp = await agent.process_query(q)
+        logging.info(f"Response: {resp}\n")
+        if _REAL_GENAI_LOADED and GEMINI_API_KEY: await asyncio.sleep(0.1)
 
 if __name__ == "__main__":
-    print(f"Script starting... Real GenAI Library Loaded: {_REAL_GENAI_LOADED}, API Key Set: {bool(GEMINI_API_KEY)}")
-    try: 
-        asyncio.run(main_test())
+    print(f"Script starting... Real GenAI: {_REAL_GENAI_LOADED}, API Key: {bool(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).")
-            print("If in Jupyter, you might need to 'await main_test()' in a cell after defining it.")
-        else: 
-            raise
-    except Exception as e_main:
-        print(f"Error during main_test execution: {e_main}")
-
+        if "asyncio.run() cannot be called" in str(e): print("Skip asyncio.run in existing loop.")
+        else: raise
+    except Exception as e_main: print(f"Test Error: {e_main}")