Update eb_agent_module.py

eb_agent_module.py

@@ -26,12 +26,18 @@ class _DummyGenAIClientModels: # Represents the dummy model service client
         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'}")
+        task_type_from_config = "N/A"
+        if isinstance(config, dict) and config.get("task_type"):
+            task_type_from_config = config["task_type"]
+        elif hasattr(config, "task_type"): # Handle if config is an object with a task_type attribute
+            task_type_from_config = config.task_type
+
+        print(f"Dummy _DummyGenAI.Client.models.embed_content called for model: {model}, task_type (from config): {task_type_from_config}")
         return {"embedding": [0.2] * 768} # Different values for dummy distinction
 
 
 class _DummyGenAIClient: # Dummy Client
-    def __init__(self, client_options=None): # Added client_options for signature consistency
+    def __init__(self, client_options=None): # client_options is kept for dummy's internal logic if any
         self.client_options = client_options
         self.models = _DummyGenAIClientModels()
         api_key_present_in_options = client_options and client_options.get("api_key")
@@ -71,12 +77,13 @@ class _ActualDummyGenAI: # type: ignore
         @staticmethod
         def GenerationConfig(**kwargs):
             print(f"Dummy _ActualDummyGenAI.types.GenerationConfig created with: {kwargs}")
-            return dict(kwargs)
+            return dict(kwargs) # Real lib returns an object, but dict is fine for dummy
 
         @staticmethod
         def SafetySetting(category, threshold):
             print(f"Dummy _ActualDummyGenAI.types.SafetySetting created: category={category}, threshold={threshold}")
-            return {"category": category, "threshold": threshold}
+            # Real lib returns an object, dict is fine for dummy, but ensure it matches expected structure if used
+            return {"category": category, "threshold": threshold} 
         
         @staticmethod # Added dummy EmbedContentConfig
         def EmbedContentConfig(task_type=None, output_dimensionality=None, title=None):
@@ -84,8 +91,16 @@ class _ActualDummyGenAI: # type: ignore
             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
+            if title: conf["title"] = title
+            # The real library returns a types.EmbedContentRequest, which has these as attributes.
+            # For the dummy, returning a dict that the dummy embed_content can understand is okay.
+            # Or, make it return a simple object:
+            class DummyEmbedConfig:
+                def __init__(self):
+                    self.task_type = task_type
+                    self.output_dimensionality = output_dimensionality
+                    self.title = title
+            return DummyEmbedConfig()
 
 
         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"
@@ -101,25 +116,27 @@ _REAL_GENAI_LOADED = False
 genai_types = None 
 
 try:
-    from google import genai 
+    # Attempt to import the actual google.generativeai library
+    import google.generativeai as real_genai 
+    genai = real_genai # Use the real library
     genai_types = genai.types 
     _REAL_GENAI_LOADED = True
-    logging.info("Successfully imported 'google.genai' and accessed 'genai.types'.")
+    logging.info("Successfully imported 'google.generativeai' and accessed 'genai.types'.")
 except ImportError:
     genai = _ActualDummyGenAI() 
     genai_types = genai.types 
-    logging.warning("Google AI library ('google.genai') not found. Using dummy implementations for 'genai' and 'genai_types'.")
-except AttributeError: # If 'genai' imported but 'genai.types' is missing
+    logging.warning("Google AI library ('google.generativeai') not found. Using dummy implementations for 'genai' and 'genai_types'.")
+except AttributeError: # If 'genai' imported but 'genai.types' is missing (less likely with modern SDK)
     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.")
+    logging.warning("'google.generativeai' 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" 
+LLM_MODEL_NAME = "gemini-1.5-flash" # Corrected to a generally available model like 1.5 flash
+GEMINI_EMBEDDING_MODEL_NAME = "text-embedding-004" # Corrected to a generally available embedding model
 
 GENERATION_CONFIG_PARAMS = {
     "temperature": 0.3, "top_p": 1.0, "top_k": 32, "max_output_tokens": 8192,
@@ -129,11 +146,15 @@ 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),
-        # ... other settings
+        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: 
     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"}] # Simplified
+    DEFAULT_SAFETY_SETTINGS = [ # Simplified for dummy if types fail
+        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
+        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}
+    ] 
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(filename)s:%(lineno)d - %(message)s')
 
@@ -145,14 +166,14 @@ if _REAL_GENAI_LOADED:
         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 for 'google.genai'.")
-    if GEMINI_API_KEY: genai.configure(api_key=GEMINI_API_KEY)
+        logging.warning("REAL 'google.generativeai' loaded, but GEMINI_API_KEY not set. API calls might fail or use other auth.")
+elif not _REAL_GENAI_LOADED: # This means we are in dummy mode
+    logging.info("Operating in DUMMY mode for 'google.generativeai'.")
+    if GEMINI_API_KEY: genai.configure(api_key=GEMINI_API_KEY) # Call dummy configure
 
 
 # --- RAG Documents Definition (Example) ---
-rag_documents_data = { 'Title': ["EB Practices", "Tech Talent"], 'Text': ["Stories...", "Projects..."] }
+rag_documents_data = { 'Title': ["EB Practices", "Tech Talent"], 'Text': ["Stories about best practices...", "Projects showcasing tech talent..."] }
 df_rag_documents = pd.DataFrame(rag_documents_data)
 
 # --- Schema Representation ---
@@ -168,89 +189,146 @@ 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_for_api = embedding_model_name # Store raw name
+        self.embedding_model_name_for_api = embedding_model_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.embedding_service_client_models = None # Will hold client.models or its dummy equivalent
 
+        # Determine if we can use the real client for RAG
+        # This depends on the library being loaded AND an API key being present
         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
+                # REAL LIBRARY: genai.Client() does not take client_options.
+                # genai.configure() should have already set the API key.
+                rag_client = genai.Client() # NO client_options here for real client
+                self.embedding_service_client_models = rag_client.models # Access the 'models' service from the client instance
+                logging.info(f"RAG: REAL embedding service (genai.Client().models) initialized for '{self.embedding_model_name_for_api}'.")
+                self._precompute_embeddings() # Try to precompute
+                self.embeddings_generated = True # Mark as generated if precomputation started
             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 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()
+                logging.error(f"RAG: Error initializing REAL embedding service (genai.Client().models): {e}", exc_info=True)
+                self.embedding_service_client_models = None 
+                self.real_client_available_for_rag = False # Cannot use real client if init fails
+        
+        # Fallback to dummy if real client is not available or failed
+        if not self.real_client_available_for_rag:
+            logging.warning(f"RAG: Not using REAL embedding service. Real GenAI Loaded: {_REAL_GENAI_LOADED}, API Key Set: {bool(GEMINI_API_KEY)}, Real Client Init Failed: {self.embedding_service_client_models is None and _REAL_GENAI_LOADED and bool(GEMINI_API_KEY)}.")
+            # Ensure we use the DUMMY genai instance if _REAL_GENAI_LOADED is false
+            dummy_genai_instance = _ActualDummyGenAI() if not _REAL_GENAI_LOADED else genai
+            self.embedding_service_client_models = dummy_genai_instance.Client().models # Gets dummy service
+            logging.info(f"RAG: Using DUMMY embedding service for '{self.embedding_model_name_for_api}'.")
+            self._precompute_embeddings() # Precompute with dummy
+
 
     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
+        if not self.embedding_service_client_models:
+            logging.error(f"RAG _embed_fn: Embedding service (client.models) not available for model '{self.embedding_model_name_for_api}'.")
+            return [0.0] * 768 # Default embedding dimension
         try:
             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)
+            # For the REAL library, embed_content takes `content` (not contents) and `task_type` directly.
+            # The `config` parameter is for more advanced settings not used here.
+            # The DUMMY `embed_content` was adapted to take a `config` object/dict.
             
-            # 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
-            )
+            if _REAL_GENAI_LOADED and self.real_client_available_for_rag: # Check if we are using the real service
+                response = self.embedding_service_client_models.embed_content(
+                    model=self.embedding_model_name_for_api,
+                    content=contents_to_embed, # Real API uses 'content'
+                    task_type=task_type        # Real API takes task_type directly
+                )
+            else: # DUMMY mode or real library failed, use dummy logic
+                # The dummy embed_content expects a config object/dict
+                embed_config_for_dummy = genai_types.EmbedContentConfig(task_type=task_type)
+                response = self.embedding_service_client_models.embed_content(
+                    model=self.embedding_model_name_for_api,
+                    contents=contents_to_embed, # Dummy API was expecting 'contents'
+                    config=embed_config_for_dummy # Dummy API was expecting 'config'
+                )
             return response["embedding"]
         except Exception as e:
-            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)
+            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}, real_client_for_rag: {self.real_client_available_for_rag}): {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
+        # Ensure text to embed is string, handle NaN gracefully
+        mask = (self.documents_df['Text'].astype(str).str.strip() != '') | \
+               (self.documents_df['Title'].astype(str).str.strip() != '')
         
+        if not mask.any(): 
+            logging.warning("No content with Text or Title found for RAG embeddings.")
+            return
+        
+        logging.info(f"Attempting to precompute embeddings for {mask.sum()} documents.")
         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
+            text_to_embed = str(row.get('Text', '')) if pd.notna(row.get('Text')) and str(row.get('Text','')).strip() else str(row.get('Title', ''))
+            if not text_to_embed.strip(): # Double check if after selection it's still empty
+                logging.debug(f"Skipping row {index} due to empty text_to_embed after selection.")
+                self.documents_df.loc[index, 'Embeddings'] = [0.0] * 768 # Store default for empty
+                continue
+
+            # Corrected task type string to match API expectations (e.g., TASK_TYPE_RETRIEVAL_DOCUMENT)
+            # For google-generativeai, it's often just "RETRIEVAL_DOCUMENT"
+            task_type_for_embedding = "RETRIEVAL_DOCUMENT" 
+            if hasattr(genai_types, 'TaskType') and hasattr(genai_types.TaskType, 'RETRIEVAL_DOCUMENT'):
+                 task_type_for_embedding = genai_types.TaskType.RETRIEVAL_DOCUMENT # Use enum if available
+
+            self.documents_df.loc[index, 'Embeddings'] = self._embed_fn(text_to_embed, task_type=task_type_for_embedding)
         
-        logging.info(f"Applied RAG embedding function to {mask.sum()} rows (embedding_service active: {self.embedding_service is not None}).")
+        self.embeddings_generated = True # Mark as generated after attempting
+        logging.info(f"Finished RAG embedding precomputation for {mask.sum()} rows (embedding_service_client_models active: {self.embedding_service_client_models is not None}).")
 
 
     def retrieve_relevant_info(self, query_text: str, top_k: int = 2) -> str:
-        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"
+        if not self.embeddings_generated: # Check if embeddings were generated
+             logging.warning("RAG: Embeddings not generated, attempting to generate them now for retrieval.")
+             self._precompute_embeddings() # Attempt to generate if not already done
+             if not self.embeddings_generated: # If still not generated (e.g. no docs)
+                  return "\n[RAG Context]\nEmbeddings could not be generated. No documents to search.\n"
+
+        if not self.real_client_available_for_rag or not self.embedding_service_client_models:
+            # If in full dummy mode and service is available, log the call for testing purposes
+            if not _REAL_GENAI_LOADED and self.embedding_service_client_models:
+                task_type_for_query = "RETRIEVAL_QUERY"
+                if hasattr(genai_types, 'TaskType') and hasattr(genai_types.TaskType, 'RETRIEVAL_QUERY'):
+                    task_type_for_query = genai_types.TaskType.RETRIEVAL_QUERY
+                self._embed_fn(query_text, task_type=task_type_for_query) # Call for dummy log
+            logging.warning(f"Skipping real RAG retrieval. Real client available for RAG: {self.real_client_available_for_rag}, Embedding service OK: {self.embedding_service_client_models is not None}")
+            return "\n[RAG Context]\nReal RAG retrieval skipped (client/service issue or dummy mode).\n"
 
         try:
-            query_embedding = np.array(self._embed_fn(query_text, task_type="RETRIEVAL_QUERY")) # Corrected task type string
+            task_type_for_query = "RETRIEVAL_QUERY"
+            if hasattr(genai_types, 'TaskType') and hasattr(genai_types.TaskType, 'RETRIEVAL_QUERY'):
+                task_type_for_query = genai_types.TaskType.RETRIEVAL_QUERY
+            query_embedding_list = self._embed_fn(query_text, task_type=task_type_for_query)
+            query_embedding = np.array(query_embedding_list)
             
+            if not np.any(query_embedding): # Check if query embedding is all zeros (error)
+                logging.warning("RAG: Query embedding resulted in all zeros. Cannot retrieve.")
+                return "\n[RAG Context]\nFailed to generate a valid embedding for the query.\n"
+
             valid_df = self.documents_df.dropna(subset=['Embeddings'])
+            # Filter out embeddings that are not lists/arrays or are empty or 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"
+            if valid_df.empty: return "\n[RAG Context]\nNo valid document embeddings available after filtering.\n"
             
             doc_embeddings = np.stack(valid_df['Embeddings'].apply(np.array).values)
-            if query_embedding.shape[0] != doc_embeddings.shape[1]: return "\n[RAG Context]\nEmbedding dimension mismatch.\n"
+            if query_embedding.shape[0] != doc_embeddings.shape[1]:
+                logging.error(f"RAG: Embedding dimension mismatch. Query: {query_embedding.shape[0]}, Docs: {doc_embeddings.shape[1]}")
+                return "\n[RAG Context]\nEmbedding dimension mismatch between query and documents.\n"
             
             dot_products = np.dot(doc_embeddings, query_embedding)
             num_to_retrieve = min(top_k, len(valid_df))
             if num_to_retrieve == 0: return "\n[RAG Context]\nNo relevant passages found (num_to_retrieve is 0).\n"
             
-            idx = np.argsort(dot_products)[-num_to_retrieve:][::-1]
+            # Get indices of top_k largest dot products
+            idx = np.argsort(dot_products)[-num_to_retrieve:][::-1] 
             passages = "".join([f"\n[RAG Context from: '{valid_df.iloc[i]['Title']}']\n{valid_df.iloc[i]['Text']}\n" for i in idx if i < len(valid_df)])
             return passages if passages else "\n[RAG Context]\nNo relevant passages found after search.\n"
         except Exception as e:
@@ -268,40 +346,45 @@ 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_instance = None # Stores the genai.Client() instance
+        self.model_service_from_client = None # Stores client.models
 
-        if _REAL_GENAI_LOADED and GEMINI_API_KEY:
+        # Determine if we can use the real client for LLM
+        self.use_real_llm_service = _REAL_GENAI_LOADED and bool(GEMINI_API_KEY)
+
+        if self.use_real_llm_service:
             try:
-                # 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
+                # REAL LIBRARY: genai.Client() does not take client_options.
+                # genai.configure() should have already set the API key.
+                self.client_instance = genai.Client() # NO client_options here
+                self.model_service_from_client = self.client_instance.models # Access 'models' service
                 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)
-                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: 
-                self.client = genai.Client() 
-                self.model_service = self.client.models 
-                logging.info("PandasLLM: Initialized with DUMMY genai.Client().models (real library failed to load).")
+                logging.error(f"Failed to initialize REAL PandasLLM with genai.Client().models: {e}", exc_info=True)
+                self.client_instance = None 
+                self.model_service_from_client = None
+                self.use_real_llm_service = False # Fallback if init fails
+        
+        # Fallback to dummy if real client is not available or failed
+        if not self.use_real_llm_service:
+            logging.warning(f"PandasLLM: Not using REAL genai.Client. RealGenAILoaded: {_REAL_GENAI_LOADED}, APIKeySet: {bool(GEMINI_API_KEY)}, Real Client Init Failed: {self.model_service_from_client is None and _REAL_GENAI_LOADED and bool(GEMINI_API_KEY)}.")
+            # Ensure we use the DUMMY genai instance if _REAL_GENAI_LOADED is false
+            dummy_genai_instance = _ActualDummyGenAI() if not _REAL_GENAI_LOADED else genai
+            self.client_instance = dummy_genai_instance.Client() 
+            self.model_service_from_client = self.client_instance.models 
+            logging.info("PandasLLM: Initialized with DUMMY genai.Client().models.")
 
 
     async def _call_gemini_api_async(self, prompt_text: str, history: list = None) -> str:
-        use_real_service = _REAL_GENAI_LOADED and GEMINI_API_KEY and self.model_service is not None
+        # active_model_service will be self.model_service_from_client (either real or dummy)
+        active_model_service = self.model_service_from_client
         
-        active_model_service = self.model_service
-        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
+        is_actually_using_real_service = self.use_real_llm_service and active_model_service is not None and hasattr(active_model_service, 'generate_content_async') and not isinstance(active_model_service, _DummyGenAIClientModels)
+
 
         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}, self.model_service is None: {self.model_service is None}). Cannot call API.")
-             return "# Error: Gemini model service not available for API call."
+            logging.error(f"PandasLLM: Model service (client.models) not available. Cannot call API.")
+            return "# Error: Gemini model service not available for API call."
 
         gemini_history = []
         if history:
@@ -320,12 +403,13 @@ class PandasLLM:
         api_generation_config = None
         if self.generation_config_dict:
             try: 
+                # genai_types will point to real or dummy types
                 api_generation_config = genai_types.GenerationConfig(**self.generation_config_dict)
             except Exception as e_cfg:
-                logging.error(f"Error creating GenerationConfig (real_loaded: {_REAL_GENAI_LOADED}): {e_cfg}. Using dict fallback.")
-                api_generation_config = self.generation_config_dict
+                logging.error(f"Error creating GenerationConfig (real_loaded: {_REAL_GENAI_LOADED}, using_real_service_now: {is_actually_using_real_service}): {e_cfg}. Using dict fallback.")
+                api_generation_config = self.generation_config_dict # Fallback for safety
 
-        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")
+        logging.info(f"\n--- Calling Gemini API (model: {model_id_for_api}, RealModeActive: {is_actually_using_real_service}) ---\nConfig: {type(api_generation_config)}\nSafety: {bool(self.safety_settings_list)}\nContent (last part text): {contents_for_api[-1]['parts'][0]['text'][:100]}...\n")
         
         try:
             response = await active_model_service.generate_content_async(
@@ -335,55 +419,62 @@ class PandasLLM:
                 safety_settings=self.safety_settings_list
             )
 
+            # Handle blocked prompts (common with real API)
             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
-                block_reason_str = str(block_reason_val.name if hasattr(block_reason_val, 'name') else block_reason_val)
+                # Convert enum to string if it's an enum
+                block_reason_str = str(block_reason_val.name if hasattr(block_reason_val, 'name') and not isinstance(block_reason_val, str) else block_reason_val)
                 logging.warning(f"Prompt blocked by API. Reason: {block_reason_str}.")
                 return f"# Error: Prompt blocked by API. Reason: {block_reason_str}."
 
             llm_output = ""
-            if hasattr(response, 'text') and isinstance(response.text, str):
+            # The real API response often has .text directly for simple cases, or via candidates
+            if hasattr(response, 'text') and isinstance(response.text, str) and response.text:
                 llm_output = response.text
             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'))
                 
+                # Check finish reason if output is empty
                 if not llm_output and candidate.finish_reason:
                     finish_reason_val = candidate.finish_reason
+                    # Convert enum to string if it's an enum
                     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:
-                             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)}"
+                    if finish_reason_str == "SAFETY" or (hasattr(genai_types, 'FinishReason') and finish_reason_val == genai_types.FinishReason.SAFETY) : 
+                        safety_messages = []
+                        if hasattr(candidate, 'safety_ratings') and 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]}")
+            else: # If no .text and no candidates, or candidates structure is unexpected
+                logging.error(f"Unexpected API response structure (no .text or valid candidates): {str(response)[:500]}")
                 return f"# Error: Unexpected API response structure: {str(response)[:200]}"
             
             return llm_output
         
-        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':
+        # Specific exceptions from the google.generativeai library
+        except (getattr(genai_types, 'BlockedPromptException', Exception)) as bpe: # Use getattr for safe access
+            if _REAL_GENAI_LOADED and type(bpe).__name__ == 'BlockedPromptException': # Check type name if real lib
                  logging.error(f"Prompt blocked (BlockedPromptException): {bpe}", exc_info=True)
                  return f"# Error: Prompt blocked. Details: {bpe}"
+            # If it's not the specific exception from the real library, or if in dummy mode, re-raise if it's a general Exception
             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:
+        except (getattr(genai_types, 'StopCandidateException', Exception)) as sce:
             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}"
+                 logging.error(f"Candidate stopped (StopCandidateException): {sce}", exc_info=True)
+                 return f"# Error: Content generation stopped. Details: {sce}"
             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)
+            logging.error(f"Error calling Gemini API (RealModeActive: {is_actually_using_real_service}): {e}", exc_info=True)
             return f"# Error during API call: {type(e).__name__} - {str(e)[:100]}."
 
 
@@ -392,97 +483,139 @@ class PandasLLM:
 
         if self.force_sandbox:
             code_to_execute = ""
+            # Improved code extraction to handle potential leading/trailing newlines or spaces
             if "```python" in llm_response_text:
                 try:
-                    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:
-                        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"): code_to_execute = code_to_execute[1:]
-                except IndexError: code_to_execute = ""
+                    # Split by ```python, take the part after it
+                    parts = llm_response_text.split("```python", 1)
+                    if len(parts) > 1:
+                        # From the second part, split by ```, take the part before it
+                        code_block_parts = parts[1].split("```", 1)
+                        code_to_execute = code_block_parts[0].strip() # .strip() to remove leading/trailing whitespace/newlines
+                except IndexError:
+                    code_to_execute = "" # Should not happen with the check above but as a safeguard
 
             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 for sandbox. Raw LLM response: {llm_response_text[:200]}")
+                # If no code was extracted, but the response wasn't an error, it might be a textual answer.
                 if not code_to_execute.strip() and not llm_response_text.startswith("# Error:"):
                     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
+                return llm_response_text # Return original LLM response (error or non-code)
 
-            logging.info(f"\n--- Code to Execute: ---\n{code_to_execute}\n----------------------\n")
+            logging.info(f"\n--- Code to Execute (Sandbox): ---\n{code_to_execute}\n----------------------\n")
             from io import StringIO; import sys
             old_stdout, sys.stdout = sys.stdout, StringIO()
+            # Prepare a safe global scope for exec
             exec_globals = {'pd': pd, 'np': np}
-            if dataframes_dict:
+            if dataframes_dict: # Add DataFrames to the execution scope
                 for name, df_instance in dataframes_dict.items():
                     if isinstance(df_instance, pd.DataFrame): exec_globals[f"df_{name}"] = df_instance
             try:
-                exec(code_to_execute, exec_globals, {})
+                exec(code_to_execute, exec_globals, {}) # Execute in the prepared scope
                 final_output_str = sys.stdout.getvalue()
-                if not final_output_str.strip():
+                if not final_output_str.strip(): # If code ran but produced no stdout
+                    # Check if the code was just comments or whitespace
                     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 by sandbox."
-                    return "# Code executed by sandbox, but no print() output. Ensure print() for results."
+                        return "# LLM generated only comments or empty code. No output from sandbox."
+                    return "# Code executed by sandbox, but no print() output. Ensure print() is used for results."
                 return final_output_str
             except Exception as e:
                 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
+            finally: sys.stdout = old_stdout # Restore stdout
+        else: # Not forcing sandbox, return raw LLM response
+            return llm_response_text
 
 # --- Employer Branding Agent ---
 class EmployerBrandingAgent:
     def __init__(self, llm_model_name: str, 
-                 generation_config_dict: dict, # Changed from gc_dict
-                 safety_settings_list: list,   # Changed from ss_list
-                 all_dataframes: dict,         # Changed from all_dfs
-                 rag_documents_df: pd.DataFrame, # Changed from rag_df
-                 embedding_model_name: str,    # Changed from emb_m_name
-                 data_privacy=True,           # Changed from dp
-                 force_sandbox=True):         # Changed from fs
+                 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 {}
         self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
         self.chat_history = [] 
-        logging.info(f"EmployerBrandingAgent Initialized (Real GenAI Loaded: {_REAL_GENAI_LOADED}).")
+        logging.info(f"EmployerBrandingAgent Initialized (Real GenAI Loaded: {_REAL_GENAI_LOADED}, RAG Real Client: {self.rag_system.real_client_available_for_rag}, LLM Real Service: {self.pandas_llm.use_real_llm_service}).")
 
     def _build_prompt(self, user_query: str, role="EB Analyst", task_hint=None, cot=True) -> str: # Simplified role for brevity in example
-        prompt = f"You are '{role}'. Goal: insights from DataFrames & RAG.\n"
-        if self.pandas_llm.data_privacy: prompt += "PRIVACY: Summarize/aggregate PII.\n"
+        prompt = f"You are '{role}'. Your primary goal is to provide insights from the provided DataFrames and RAG context.\n"
+        if self.pandas_llm.data_privacy: prompt += "IMPORTANT PRIVACY NOTE: If dealing with Personally Identifiable Information (PII), you must summarize or aggregate it. Do not output raw PII.\n"
+        
         if self.pandas_llm.force_sandbox:
-            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"
+            prompt += "RESPONSE FORMAT: Generate Python code to analyze the data and `print()` the textual insights or answers. Your entire response MUST be a single Python code block enclosed in ```python ... ```. Do NOT provide any text outside this block.\n"
+            prompt += "CODE GUIDELINES: Access DataFrames using their assigned 'df_name' (e.g., df_sales, df_employees). Comment your code clearly. Handle potential issues like ambiguity or missing data by printing informative messages.\n"
+            prompt += "CRITICAL: Your Python code's `print()` statements should output the final synthesized insights, not just raw DataFrames (unless specifically asked to show a DataFrame sample).\n"
+        else: 
+            prompt += "RESPONSE FORMAT: Provide comprehensive textual insights. Explain your reasoning and steps clearly.\n"
+        
+        prompt += f"--- AVAILABLE DATA SCHEMAS ---\n{self.schemas_representation if self.schemas_representation.strip() and self.schemas_representation.strip() != 'No DataFrames provided.' else 'No DataFrames are currently loaded or available.'}\n"
         
+        # Retrieve RAG context. The RAG system itself handles logging if it's skipped.
         rag_context = self.rag_system.retrieve_relevant_info(user_query)
-        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"
+        # Check if RAG context is meaningful (not an error or "skipped" message)
+        non_meaningful_rag_keywords = ["[RAG Context]\nError", "[RAG Context]\nSkipped", "[RAG Context]\nNo valid", "[RAG Context]\nNo relevant", "could not be generated", "not available", "dimension mismatch"]
+        is_meaningful_rag = bool(rag_context.strip()) and not any(keyword in rag_context for keyword in non_meaningful_rag_keywords)
+        
+        prompt += f"--- RAG (Retrieval Augmented Generation) CONTEXT (Real RAG active: {self.rag_system.real_client_available_for_rag and self.rag_system.embedding_service_client_models is not None}) ---\n"
+        if is_meaningful_rag:
+            prompt += f"{rag_context}\n"
+        else:
+            prompt += f"No specific relevant RAG context found for this query, or RAG system reported an issue. Details from RAG: {rag_context.strip()[:150]}...\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"
+        if task_hint: prompt += f"--- ADDITIONAL GUIDANCE ---\n{task_hint}\n"
+        
+        if cot: # Chain of Thought guidance
+            if self.pandas_llm.force_sandbox: 
+                prompt += "--- PYTHON CODE THOUGHT PROCESS (Follow these steps internally before writing code) ---\n"
+                prompt += "1. Understand the Goal: What specific question does the user query ask? What insight is needed?\n"
+                prompt += "2. Identify Data Sources: Which DataFrame(s) are relevant? Is the RAG context useful?\n"
+                prompt += "3. Formulate a Plan: Outline the steps to take. E.g., filter df_X, aggregate df_Y, combine with RAG info.\n"
+                prompt += "4. Write Python Code: Implement the plan. Use pandas for DataFrame operations. Access RAG context as text.\n"
+                prompt += "5. CRITICAL - Print Results: Ensure `print()` statements clearly output the synthesized answer/insights. Do not just print raw data.\n"
+                prompt += "6. Review and Refine: Check code for correctness, clarity, and efficiency. Ensure it directly addresses the query.\n"
+                prompt += "7. Final Output: Ensure the entire response is ONLY the Python code block: ```python ... ```.\n"
+            else: 
+                prompt += "--- TEXT RESPONSE THOUGHT PROCESS (Follow these steps internally) ---\n"
+                prompt += "1. Understand the Goal: What is the user asking for?\n"
+                prompt += "2. Identify Data Sources: Which DataFrames and RAG context sections are relevant?\n"
+                prompt += "3. Synthesize Insights: Combine information from DataFrames and RAG to form a coherent answer.\n"
+                prompt += "4. Structure Response: Organize the answer logically and explain the findings clearly.\n"
         return prompt
 
     async def process_query(self, user_query: str, role="EB Analyst", task_hint=None, cot=True) -> str: # Simplified role
-        hist_for_llm = self.chat_history[:]
+        # Keep a temporary copy of history for this call, so current user_query isn't in it
+        hist_for_llm = self.chat_history[:] 
+        
+        # Add current user query to persistent history *before* calling LLM
         self.chat_history.append({"role": "user", "content": user_query})
+        
         prompt = self._build_prompt(user_query, role, task_hint, cot)
-        logging.info(f"Prompt for query: {user_query[:70]}... (Real GenAI: {_REAL_GENAI_LOADED})")
+        logging.info(f"Agent processing query: '{user_query[:70]}...' (Real GenAI Loaded: {_REAL_GENAI_LOADED}, LLM Real Service: {self.pandas_llm.use_real_llm_service})")
+        
         response = await self.pandas_llm.query(prompt, self.all_dataframes, history=hist_for_llm)
+        
+        # Add assistant's response to persistent history
         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.")
+        
+        # Truncate history if it gets too long
+        if len(self.chat_history) > 10: # Keep last 5 pairs (user/assistant)
+            self.chat_history = self.chat_history[-10:]
+            logging.info("Chat history truncated to the last 10 entries.")
         return response
 
     def update_dataframes(self, new_dataframes: dict): # Changed from new_dfs
         self.all_dataframes = new_dataframes if new_dataframes else {}
         self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
-        logging.info("Agent DFs updated.")
+        logging.info(f"Agent DataFrames updated. New schema: {self.schemas_representation[:200]}...")
     
     def clear_chat_history(self): 
         self.chat_history = []
@@ -490,28 +623,79 @@ class EmployerBrandingAgent:
 
 # --- Example Usage (Conceptual) ---
 async def main_test():
-    logging.info(f"Test (Real GenAI: {_REAL_GENAI_LOADED}, API Key: {bool(GEMINI_API_KEY)})")
-    # Example: Provide all parameters to EmployerBrandingAgent constructor
+    # Configure logging for the test
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(filename)s:%(lineno)d - %(message)s')
+    
+    print(f"--- main_test() ---")
+    logging.info(f"Test starting with Real GenAI Loaded: {_REAL_GENAI_LOADED}, API Key Set: {bool(GEMINI_API_KEY)}")
+    
+    # Example DataFrames for testing
+    sample_employee_data = {
+        'EmployeeID': [1, 2, 3, 4, 5],
+        'Name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
+        'Department': ['HR', 'Tech', 'Tech', 'Sales', 'HR'],
+        'Projects': [['Onboarding Revamp'], ['AI Chatbot', 'Platform Upgrade'], ['AI Chatbot'], ['Client Outreach'], ['Benefits Portal']]
+    }
+    df_employees = pd.DataFrame(sample_employee_data)
+
+    sample_rag_docs_data = {
+        'Title': ["Company Culture Handbook", "Tech Team Achievements Q1", "Remote Work Policy"],
+        'Text': [
+            "Our company values collaboration, innovation, and inclusivity. We encourage cross-departmental projects.",
+            "The tech team successfully launched the AI Chatbot project, significantly improving customer engagement. The Platform Upgrade is on track.",
+            "We offer flexible remote work options. All employees are expected to maintain high levels of communication."
+        ]
+    }
+    df_test_rag_documents = pd.DataFrame(sample_rag_docs_data)
+
     agent = EmployerBrandingAgent(
         llm_model_name=LLM_MODEL_NAME,
         generation_config_dict=GENERATION_CONFIG_PARAMS,
         safety_settings_list=DEFAULT_SAFETY_SETTINGS,
-        all_dataframes={}, # Empty dict for example
-        rag_documents_df=df_rag_documents,
+        all_dataframes={'employees': df_employees}, # Provide a sample DataFrame
+        rag_documents_df=df_test_rag_documents,   # Provide sample RAG documents
         embedding_model_name=GEMINI_EMBEDDING_MODEL_NAME,
         data_privacy=True,
-        force_sandbox=True
+        force_sandbox=True # Set to True to test code generation, False for direct text
     )
-    for q in ["What are EB best practices?", "Hello Agent!"]:
-        logging.info(f"\nQuery: {q}")
+    
+    queries = [
+        "What are the key aspects of our company culture according to the handbook?",
+        "Which employees are working on the 'AI Chatbot' project? Summarize the project's impact.",
+        "How many employees are in the Tech department?"
+    ]
+    
+    for q_idx, q in enumerate(queries):
+        logging.info(f"\n--- Query {q_idx+1}: {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)
+        logging.info(f"--- Agent Response for Query {q_idx+1} ---:\n{resp}\n---------------------\n")
+        # Add a small delay if using real API to avoid rate limits, though for flash models it's usually fine.
+        if _REAL_GENAI_LOADED and GEMINI_API_KEY: await asyncio.sleep(0.2) 
 
 if __name__ == "__main__":
-    print(f"Script starting... Real GenAI: {_REAL_GENAI_LOADED}, API Key: {bool(GEMINI_API_KEY)}")
-    try: asyncio.run(main_test())
+    print(f"Script starting... Real GenAI: {_REAL_GENAI_LOADED}, API Key Set: {bool(GEMINI_API_KEY)}")
+    # Setup basic logging if not already configured by a higher-level application
+    if not logging.getLogger().hasHandlers():
+        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s')
+
+    try: 
+        asyncio.run(main_test())
     except RuntimeError as e:
-        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}")
+        if "asyncio.run() cannot be called when another asyncio event loop is running" in str(e): 
+            print("Skipping asyncio.run(main_test()) as an event loop is already running (e.g., in Jupyter).")
+            # If in Jupyter or similar, you might need to run like this:
+            # loop = asyncio.get_event_loop()
+            # if loop.is_running():
+            #     print("Event loop is running, creating task for main_test()")
+            #     asyncio.create_task(main_test())
+            # else:
+            #     print("Event loop not running, using asyncio.run()")
+            #     asyncio.run(main_test())
+
+        else: 
+            print(f"Main Test Runtime Error: {e}")
+            raise
+    except Exception as e_main: 
+        print(f"Main Test Exception: {e_main}")
+        logging.error("Exception in __main__ execution of main_test:", exc_info=True)
+