Update eb_agent_module.py

eb_agent_module.py

@@ -11,7 +11,6 @@ import textwrap
 class _DummyGenAIClientModels: # Represents the dummy model service client
     async def generate_content_async(self, model=None, contents=None, generation_config=None, safety_settings=None, stream=False, tools=None, tool_config=None):
         print(f"Dummy _DummyGenAI.Client.models.generate_content_async called for model: {model}")
-        # Simplified dummy response structure
         class DummyPart: text = "# Dummy response from _DummyGenAI async"
         class DummyContent: parts = [DummyPart()]
         class DummyCandidate: content = DummyContent(); finish_reason = "_DUMMY_STOP"; safety_ratings = []; token_count = 0; index = 0
@@ -45,7 +44,7 @@ class _DummyGenAIGenerativeModel:
         return DummyResponse()
 
 class _ActualDummyGenAI: # type: ignore # Renamed the main dummy class
-    Client = _DummyGenAIClient # Assign inner class
+    Client = _DummyGenAIClient 
 
     @staticmethod
     def configure(api_key):
@@ -58,41 +57,56 @@ class _ActualDummyGenAI: # type: ignore # Renamed the main dummy class
 
     @staticmethod
     def embed_content(model, content, task_type, title=None):
-        # This print is crucial for debugging which embed_content is called
         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):
+            print(f"Dummy _ActualDummyGenAI.types.GenerationConfig created with: {kwargs}")
+            return dict(kwargs)
+
+        @staticmethod
+        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
+        class BlockedPromptException(Exception): pass
+        class StopCandidateException(Exception): pass
 
-class _ActualDummyGenAITypes: # type: ignore # Renamed the main dummy types class
-    @staticmethod
-    def GenerationConfig(**kwargs):
-        print(f"Dummy _ActualDummyGenAITypes.GenerationConfig created with: {kwargs}")
-        return dict(kwargs)
-
-    @staticmethod
-    def SafetySetting(category, threshold):
-        print(f"Dummy _ActualDummyGenAITypes.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: # This should match the structure of the real FinishReason enum if possible
-        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"
 
 # --- Attempt to import the real library ---
 _REAL_GENAI_LOADED = False
+genai_types = None # Initialize genai_types
+
 try:
-    from google import generativeai as genai # This is the real 'genai'
-    from google.generativeai import types as genai_types # This is the real 'genai_types'
+    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
     _REAL_GENAI_LOADED = True
-    logging.info("Successfully imported 'google.generativeai' library.")
+    logging.info("Successfully imported 'google.genai' and accessed 'genai.types'.")
 except ImportError:
-    genai = _ActualDummyGenAI() # Alias to our dummy genai class instance if import fails
-    genai_types = _ActualDummyGenAITypes() # Alias to our dummy genai_types class instance
-    logging.warning("Google Generative AI library not found. Using dummy implementations.")
+    # 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
+    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
+    logging.warning("'google.genai' imported, but 'genai.types' not found. Falling back to dummy implementations.")
 
 
 # --- Configuration ---
@@ -105,7 +119,7 @@ GENERATION_CONFIG_PARAMS = {
 }
 
 # Default safety settings list for Gemini
-# Ensure genai_types used here is the one defined (real or dummy alias)
+# 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),
@@ -113,8 +127,8 @@ try:
         genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
         genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
     ]
-except Exception as e_safety: # Catch broader exception if dummy types are not perfect
-    logging.warning(f"Could not define DEFAULT_SAFETY_SETTINGS using genai_types: {e_safety}. Using placeholder list of dicts.")
+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"},
@@ -126,16 +140,16 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 
 if GEMINI_API_KEY and _REAL_GENAI_LOADED:
     try:
-        genai.configure(api_key=GEMINI_API_KEY) # genai is now consistently real or dummy
+        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.")
+    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.")
 elif not _REAL_GENAI_LOADED:
-    logging.info("Operating in DUMMY mode because 'google-generativeai' library was not found.")
-    if GEMINI_API_KEY: # Call dummy configure if key is present but library is dummy
-        genai.configure(api_key=GEMINI_API_KEY)
+    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
 
 
 # --- RAG Documents Definition (Example) ---
@@ -164,28 +178,24 @@ class AdvancedRAGSystem:
         self.embedding_model_name = embedding_model_name
         self.documents_df = documents_df.copy()
         self.embeddings_generated = False
-        # Use _REAL_GENAI_LOADED to determine if real client is available
         self.real_client_available_for_rag = _REAL_GENAI_LOADED and bool(GEMINI_API_KEY)
 
         if self.real_client_available_for_rag:
             try:
                 self._precompute_embeddings()
                 self.embeddings_generated = True
-                # This log should only appear if real genai.embed_content was used without printing dummy message
                 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)
         else:
-            logging.warning(f"RAG embeddings not precomputed. Real GenAI loaded: {_REAL_GENAI_LOADED}, API Key set: {bool(GEMINI_API_KEY)}.")
-            # If in dummy mode, call dummy precompute to see its log
-            if not _REAL_GENAI_LOADED:
-                self._precompute_embeddings() # This will call dummy genai.embed_content
+            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
 
     def _embed_fn(self, title: str, text: str) -> list[float]:
-        # genai here is now consistently the real or the aliased dummy
         try:
             content_to_embed = text if text else title
             if not content_to_embed: return [0.0] * 768
-            # The call to genai.embed_content will print its own message if it's the dummy
+            # 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"]
         except Exception as e:
             logging.error(f"Error in _embed_fn for '{title}' (real_genai_loaded: {_REAL_GENAI_LOADED}): {e}", exc_info=True)
@@ -195,26 +205,22 @@ class AdvancedRAGSystem:
         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
-        # This will call _embed_fn, which calls the current 'genai.embed_content' (real or dummy)
         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}).")
 
 
     def retrieve_relevant_info(self, query_text: str, top_k: int = 2) -> str:
-        if not (_REAL_GENAI_LOADED and GEMINI_API_KEY): # Check if we can use real embeddings
-            # If not using real, and dummy is active, dummy embed_content will print.
-            # If real loaded but no key, this will also be skipped for actual API call.
-            if not _REAL_GENAI_LOADED: # If in dummy mode, call dummy embed_content to see log
-                 genai.embed_content(model=self.embedding_model_name, content=query_text, task_type="retrieval_query") # Call for log
+        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"
 
-        # At this point, _REAL_GENAI_LOADED and GEMINI_API_KEY are true
-        # So, genai.embed_content should be the real one.
         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"])
             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))] # Ensure not all zeros
+            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"
             
             doc_embeddings = np.stack(valid_df['Embeddings'].apply(np.array).values)
@@ -247,42 +253,32 @@ class PandasLLM:
 
         if _REAL_GENAI_LOADED and GEMINI_API_KEY:
             try:
-                self.client = genai.Client() # Should be the REAL genai.Client
+                self.client = genai.Client() # Real genai.Client
                 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)
-                # No explicit fallback to dummy here; _call_gemini_api_async will use the global dummy if self.model_service is None and _REAL_GENAI_LOADED is False
+                # If this fails, self.model_service remains None, _call_gemini_api_async might fallback if _REAL_GENAI_LOADED somehow becomes False
         else:
-            logging.warning(f"PandasLLM: Not using REAL genai.Client. RealGenAILoaded: {_REAL_GENAI_LOADED}, APIKeySet: {bool(GEMINI_API_KEY)}. Will use DUMMY if library not loaded.")
+            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
-                self.model_service = self.client.models # Uses _DummyGenAIClientModels
-                logging.info("PandasLLM: Initialized with DUMMY genai.Client().models because real library failed to load.")
+                self.model_service = self.client.models 
+                logging.info("PandasLLM: Initialized with DUMMY genai.Client().models (real library failed to load).")
 
 
     async def _call_gemini_api_async(self, prompt_text: str, history: list = None) -> str:
-        # Determine if we should use the real service or expect dummy behavior
         use_real_service = _REAL_GENAI_LOADED and GEMINI_API_KEY and self.model_service is not None
         
-        # If not using real service, and we are in dummy mode (library not loaded),
-        # self.model_service should be the dummy one.
-        # If real library loaded but no key, self.model_service might be None or real (but calls would fail).
-
         active_model_service = self.model_service
-        if not use_real_service and not _REAL_GENAI_LOADED:
-            # Ensure we have a dummy service if we are in full dummy mode and self.model_service wasn't set
-            # This case should ideally be covered by PandasLLM.__init__
-            if active_model_service is None:
-                logging.debug("PandasLLM._call_gemini_api_async: active_model_service is None in dummy mode, attempting to get dummy service.")
-                dummy_client_instance = _ActualDummyGenAI.Client() # Get a fresh dummy client models service
-                active_model_service = dummy_client_instance.models
-
+        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__
+                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
 
         if not active_model_service:
-             logging.error("PandasLLM: Model service not available (real or dummy). Cannot call API.")
-             return "# Error: Gemini model service not available."
-
+             logging.error(f"PandasLLM: Model service not available (use_real_service: {use_real_service}, _REAL_GENAI_LOADED: {_REAL_GENAI_LOADED}). Cannot call API.")
+             return "# Error: Gemini model service not available for API call."
 
         gemini_history = []
         if history:
@@ -300,17 +296,15 @@ class PandasLLM:
 
         api_generation_config = None
         if self.generation_config_dict:
-            try: # genai_types is now consistently real or dummy alias
+            try: 
                 api_generation_config = genai_types.GenerationConfig(**self.generation_config_dict)
             except Exception as e_cfg:
-                logging.error(f"Error creating GenerationConfig object (real_loaded: {_REAL_GENAI_LOADED}): {e_cfg}. Using dict fallback.")
+                logging.error(f"Error creating GenerationConfig (real_loaded: {_REAL_GENAI_LOADED}): {e_cfg}. Using dict fallback.")
                 api_generation_config = self.generation_config_dict
 
         logging.info(f"\n--- Calling Gemini API (model: {model_id_for_api}, RealMode: {use_real_service}) ---\nConfig: {api_generation_config}\nSafety: {bool(self.safety_settings_list)}\nContent (last part text): {contents_for_api[-1]['parts'][0]['text'][:100]}...\n")
         
         try:
-            # This call will use either the real model_service or the dummy one.
-            # The dummy service's methods have print statements.
             response = await active_model_service.generate_content_async(
                 model=model_id_for_api,
                 contents=contents_for_api,
@@ -328,17 +322,15 @@ class PandasLLM:
             llm_output = ""
             if hasattr(response, 'text') and isinstance(response.text, str):
                 llm_output = response.text
-            elif response.candidates: # Standard way to get text from Gemini response
+            elif response.candidates: 
                 candidate = response.candidates[0]
                 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 candidate.finish_reason:
                     finish_reason_val = candidate.finish_reason
-                    # Try to get enum name if available (for real API) or use string (for dummy)
                     finish_reason_str = str(finish_reason_val.name if hasattr(finish_reason_val, 'name') and not isinstance(finish_reason_val, str) else finish_reason_val)
 
-
                     if finish_reason_str == "SAFETY": 
                          safety_messages = []
                          if hasattr(candidate, 'safety_ratings') and candidate.safety_ratings:
@@ -357,13 +349,20 @@ class PandasLLM:
             
             return llm_output
         
-        # Specific exceptions for the real API, might not be raised by dummy
-        except genai_types.BlockedPromptException as bpe: # type: ignore
-            logging.error(f"Prompt blocked (BlockedPromptException): {bpe}", exc_info=True)
-            return f"# Error: Prompt blocked. Details: {bpe}"
-        except genai_types.StopCandidateException as sce: # type: ignore
-            logging.error(f"Candidate stopped (StopCandidateException): {sce}", exc_info=True)
-            return f"# Error: Content generation stopped. Details: {sce}"
+        # 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
+                 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
+
+        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
+                logging.error(f"Candidate stopped (StopCandidateException): {sce}", exc_info=True)
+                return f"# Error: Content generation stopped. Details: {sce}"
+            pass # Fallthrough
+
         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]}."
@@ -386,11 +385,9 @@ class PandasLLM:
                 except IndexError: code_to_execute = ""
 
             if llm_response_text.startswith("# Error:") or not code_to_execute.strip():
-                # Log if it's an error from LLM or if it's just non-code/comment response.
                 logging.warning(f"LLM response is an error, or no valid Python code block found for sandbox. Raw LLM response: {llm_response_text[:200]}")
                 if not code_to_execute.strip() and not llm_response_text.startswith("# Error:"):
-                    # If it's not an error and not code, it might be a natural language refusal or comment.
-                    if "```" not in llm_response_text and len(llm_response_text.strip()) > 0: # Heuristic for non-code text
+                    if "```" not in llm_response_text and len(llm_response_text.strip()) > 0: 
                         logging.info(f"LLM produced text output instead of Python code in sandbox mode. Passing through: {llm_response_text[:200]}")
                 return llm_response_text
 
@@ -459,7 +456,6 @@ class EmployerBrandingAgent:
         prompt += "\n--- AVAILABLE DATA AND SCHEMAS ---\n"
         prompt += self.schemas_representation if self.schemas_representation.strip() != "No DataFrames provided." else "No DataFrames loaded.\n"
         
-        # RAG retrieval will use the current state of 'genai' (real or dummy)
         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)
@@ -504,7 +500,6 @@ class EmployerBrandingAgent:
 
 # --- Example Usage (Conceptual) ---
 async def main_test():
-    # This test will reflect whether _REAL_GENAI_LOADED is true or false
     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]})
@@ -521,16 +516,10 @@ async def main_test():
         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) # Small delay for real API
+        if _REAL_GENAI_LOADED and GEMINI_API_KEY: await asyncio.sleep(0.1) 
 
 if __name__ == "__main__":
-    # Note: To test with real API, ensure GEMINI_API_KEY is set in your environment
-    # and 'google-generativeai' is installed.
-    # Otherwise, it will run in dummy mode.
-    
-    # Check mode before running test
     print(f"Script starting... Real GenAI Library Loaded: {_REAL_GENAI_LOADED}, API Key Set: {bool(GEMINI_API_KEY)}")
-
     try: 
         asyncio.run(main_test())
     except RuntimeError as e: