Update eb_agent_module.py

eb_agent_module.py

@@ -7,145 +7,105 @@ import logging
 import numpy as np
 import textwrap
 
-# Attempt to import Google Generative AI and related types
+# --- Define Dummy Classes with unique names first ---
+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
+        class DummyResponse: candidates = [DummyCandidate()]; text = DummyCandidate.content.parts[0].text; prompt_feedback = None
+        return DummyResponse()
+
+    def generate_content(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 called for model: {model}")
+        class DummyPart: text = "# Dummy response from _DummyGenAI sync"
+        class DummyContent: parts = [DummyPart()]
+        class DummyCandidate: content = DummyContent(); finish_reason = "_DUMMY_STOP"; safety_ratings = []; token_count = 0; index = 0
+        class DummyResponse: candidates = [DummyCandidate()]; text = DummyCandidate.content.parts[0].text; prompt_feedback = None
+        return DummyResponse()
+
+class _DummyGenAIClient: # Dummy Client
+    def __init__(self, api_key=None):
+        self.api_key = api_key
+        self.models = _DummyGenAIClientModels()
+        print(f"Dummy _DummyGenAI.Client initialized {'with api_key' if api_key else '(global API key expected)'}.")
+
+class _DummyGenAIGenerativeModel:
+    def __init__(self, model_name_in, generation_config_in, safety_settings_in, system_instruction_in):
+        self.model_name = model_name_in
+        print(f"Dummy _DummyGenAIGenerativeModel initialized for {model_name_in}")
+    async def generate_content_async(self, contents, stream=False):
+        print(f"Dummy _DummyGenAIGenerativeModel.generate_content_async called for {self.model_name}")
+        class DummyPart: text = f"# Dummy response from dummy _DummyGenAIGenerativeModel ({self.model_name})"
+        class DummyContent: parts = [DummyPart()]
+        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()
+
+class _ActualDummyGenAI: # type: ignore # Renamed the main dummy class
+    Client = _DummyGenAIClient # Assign inner class
+
+    @staticmethod
+    def configure(api_key):
+        print(f"Dummy _ActualDummyGenAI.configure called with API key: {'SET' if api_key else 'NOT SET'}")
+
+    @staticmethod
+    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):
+        # 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}
+
+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
 try:
-    from google import generativeai as genai # Renamed for clarity to avoid conflict
-    from google.generativeai import types as genai_types
-    # from google.generativeai import GenerationConfig # For direct use if needed
-    # from google.generativeai.types import HarmCategory, HarmBlockThreshold, SafetySetting # For direct use
-
+    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'
+    _REAL_GENAI_LOADED = True
+    logging.info("Successfully imported 'google.generativeai' library.")
 except ImportError:
-    print("Google Generative AI library not found. Please install it: pip install google-generativeai")
-    # Define dummy classes/functions if the import fails, to allow the rest of the script to be parsed
-    class genai: # type: ignore
-        @staticmethod
-        def configure(api_key):
-            print(f"Dummy genai.configure called with API key: {'SET' if api_key else 'NOT SET'}")
-
-        # Dummy Client and related structures
-        class Client:
-            def __init__(self, api_key=None): # api_key is optional for Client constructor
-                self.api_key = api_key
-                self.models = self._Models() # This is the service client for models
-                print(f"Dummy genai.Client initialized {'with api_key' if api_key else '(global API key expected)'}.")
-
-            class _Models: # Represents the model service client
-                async def generate_content_async(self, model=None, contents=None, generation_config=None, safety_settings=None, stream=False, tools=None, tool_config=None): # Matched real signature better
-                    print(f"Dummy genai.Client.models.generate_content_async called for model: {model} with config: {generation_config}, safety_settings: {safety_settings}, stream: {stream}")
-                    class DummyPart:
-                        def __init__(self, text): self.text = text
-                    class DummyContent:
-                        def __init__(self): self.parts = [DummyPart("# Dummy response from dummy client's async generate_content")]
-                    class DummyCandidate:
-                        def __init__(self):
-                            self.content = DummyContent()
-                            self.finish_reason = genai_types.FinishReason.STOP # Use dummy FinishReason
-                            self.safety_ratings = []
-                            self.token_count = 0
-                            self.index = 0
-                    class DummyResponse:
-                        def __init__(self):
-                            self.candidates = [DummyCandidate()]
-                            self.prompt_feedback = self._PromptFeedback()
-                            self.text = "# Dummy response text from dummy client's async generate_content"
-                        class _PromptFeedback:
-                            def __init__(self):
-                                self.block_reason = None
-                                self.safety_ratings = []
-                    return DummyResponse()
-
-                def generate_content(self, model=None, contents=None, generation_config=None, safety_settings=None, stream=False, tools=None, tool_config=None): # Matched real signature better
-                    print(f"Dummy genai.Client.models.generate_content called for model: {model} with config: {generation_config}, safety_settings: {safety_settings}, stream: {stream}")
-                    # Re-using the async dummy structure for simplicity
-                    class DummyPart:
-                        def __init__(self, text): self.text = text
-                    class DummyContent:
-                        def __init__(self): self.parts = [DummyPart("# Dummy response from dummy client's generate_content")]
-                    class DummyCandidate:
-                        def __init__(self):
-                            self.content = DummyContent()
-                            self.finish_reason = genai_types.FinishReason.STOP
-                            self.safety_ratings = []
-                            self.token_count = 0
-                            self.index = 0
-                    class DummyResponse:
-                        def __init__(self):
-                            self.candidates = [DummyCandidate()]
-                            self.prompt_feedback = self._PromptFeedback()
-                            self.text = "# Dummy response text from dummy client's generate_content"
-                        class _PromptFeedback:
-                            def __init__(self):
-                                self.block_reason = None
-                                self.safety_ratings = []
-                    return DummyResponse()
-
-        @staticmethod
-        def GenerativeModel(model_name, generation_config=None, safety_settings=None, system_instruction=None): # Kept for AdvancedRAGSystem if it uses it, or if user switches back
-            print(f"Dummy genai.GenerativeModel called for model: {model_name} (This might be unused if Client approach is preferred)")
-            # ... (rest of DummyGenerativeModel as before, for completeness) ...
-            class DummyGenerativeModel:
-                def __init__(self, model_name_in, generation_config_in, safety_settings_in, system_instruction_in):
-                    self.model_name = model_name_in
-                async def generate_content_async(self, contents, stream=False):
-                    class DummyPart:
-                        def __init__(self, text): self.text = text
-                    class DummyContent:
-                        def __init__(self): self.parts = [DummyPart(f"# Dummy response from dummy GenerativeModel ({self.model_name})")]
-                    class DummyCandidate:
-                        def __init__(self):
-                            self.content = DummyContent(); self.finish_reason = genai_types.FinishReason.STOP; self.safety_ratings = []
-                    class DummyResponse:
-                        def __init__(self):
-                            self.candidates = [DummyCandidate()]; self.prompt_feedback = None; self.text = f"# Dummy GM response"
-                    return DummyResponse()
-            return DummyGenerativeModel(model_name, generation_config, safety_settings, system_instruction)
-
-
-        @staticmethod
-        def embed_content(model, content, task_type, title=None):
-            print(f"Dummy genai.embed_content called for model: {model}, task_type: {task_type}, title: {title}")
-            return {"embedding": [0.1] * 768}
-
-    class genai_types: # type: ignore
-        @staticmethod
-        def GenerationConfig(**kwargs):
-            print(f"Dummy genai_types.GenerationConfig created with: {kwargs}")
-            return dict(kwargs)
-
-        @staticmethod
-        def SafetySetting(category, threshold):
-            print(f"Dummy SafetySetting created: category={category}, threshold={threshold}")
-            return {"category": category, "threshold": threshold}
-
-        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"
-        
-        # Dummy for BlockedReason if needed by response parsing
-        class BlockedReason:
-            BLOCKED_REASON_UNSPECIFIED = "BLOCKED_REASON_UNSPECIFIED"
-            SAFETY = "SAFETY"
-            OTHER = "OTHER"
+    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.")
+
 
 # --- Configuration ---
 GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', "")
-# User-specified model names:
-# LLM_MODEL_NAME = "gemini-2.0-flash" # Original
 LLM_MODEL_NAME = "gemini-2.0-flash" 
 GEMINI_EMBEDDING_MODEL_NAME = "gemini-embedding-exp-03-07"
 
-# Base generation configuration for the LLM
 GENERATION_CONFIG_PARAMS = {
-    "temperature": 0.3,
-    "top_p": 1.0,
-    "top_k": 32,
-    "max_output_tokens": 8192,
+    "temperature": 0.3, "top_p": 1.0, "top_k": 32, "max_output_tokens": 8192,
 }
 
 # Default safety settings list for Gemini
+# Ensure genai_types used here is the one defined (real or dummy alias)
 try:
     DEFAULT_SAFETY_SETTINGS = [
         genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
@@ -153,26 +113,29 @@ 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 AttributeError as e:
-    logging.warning(f"Could not define DEFAULT_SAFETY_SETTINGS using real genai_types: {e}. Using placeholder list of dicts.")
+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.")
     DEFAULT_SAFETY_SETTINGS = [
         {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
         {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-        {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-        {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
     ]
 
+
 # Logging setup
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(filename)s:%(lineno)d - %(message)s')
 
-if GEMINI_API_KEY:
+if GEMINI_API_KEY and _REAL_GENAI_LOADED:
     try:
-        genai.configure(api_key=GEMINI_API_KEY)
-        logging.info(f"Gemini API key configured globally.")
+        genai.configure(api_key=GEMINI_API_KEY) # genai is now consistently real or dummy
+        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)
-else:
-    logging.warning("GEMINI_API_KEY environment variable not set. Agent will use dummy responses if real genai library is not fully mocked or if API calls fail.")
+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.")
+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)
 
 
 # --- RAG Documents Definition (Example) ---
@@ -201,44 +164,58 @@ class AdvancedRAGSystem:
         self.embedding_model_name = embedding_model_name
         self.documents_df = documents_df.copy()
         self.embeddings_generated = False
-        # Check if genai.embed_content is the real one or our dummy
-        self.client_available = hasattr(genai, 'embed_content') and not (hasattr(genai.embed_content, '__func__') and genai.embed_content.__func__.__qualname__.startswith('genai.embed_content'))
+        # 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 GEMINI_API_KEY and self.client_available:
+        if self.real_client_available_for_rag:
             try:
                 self._precompute_embeddings()
                 self.embeddings_generated = True
-                logging.info(f"RAG embeddings precomputed using '{self.embedding_model_name}'.")
-            except Exception as e: logging.error(f"RAG precomputation error: {e}", exc_info=True)
+                # 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. Key: {bool(GEMINI_API_KEY)}, embed_content_ok: {self.client_available}.")
+            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
 
     def _embed_fn(self, title: str, text: str) -> list[float]:
-        if not self.client_available: return [0.0] * 768 
+        # 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
             return genai.embed_content(model=self.embedding_model_name, content=content_to_embed, task_type="retrieval_document", title=title if title else None)["embedding"]
         except Exception as e:
-            logging.error(f"Error in _embed_fn for '{title}': {e}", exc_info=True)
+            logging.error(f"Error in _embed_fn for '{title}' (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
+        # 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.")
+        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 self.client_available: return "\n[RAG Context]\nEmbedding client not available.\n"
-        if not self.embeddings_generated or 'Embeddings' not in self.documents_df.columns or self.documents_df['Embeddings'].isnull().all():
-            return "\n[RAG Context]\nEmbeddings not ready for RAG.\n"
+        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
+            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:
             query_embedding = np.array(genai.embed_content(model=self.embedding_model_name, content=query_text, task_type="retrieval_query")["embedding"])
             valid_df = self.documents_df.dropna(subset=['Embeddings'])
-            valid_df = valid_df[valid_df['Embeddings'].apply(lambda x: isinstance(x, (list, np.ndarray)) and len(x) > 0)]
-            if valid_df.empty: return "\n[RAG Context]\nNo valid document embeddings.\n"
+            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
+            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)
             if query_embedding.shape[0] != doc_embeddings.shape[1]: return "\n[RAG Context]\nEmbedding dimension mismatch.\n"
@@ -251,8 +228,8 @@ 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: {e}", exc_info=True)
-            return f"\n[RAG Context]\nError during RAG retrieval: {type(e).__name__} - {e}\n"
+            logging.error(f"Error in RAG retrieve_relevant_info (real mode): {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) ---
 class PandasLLM:
@@ -261,69 +238,80 @@ class PandasLLM:
                  safety_settings_list: list,
                  data_privacy=True, force_sandbox=True):
         self.llm_model_name = llm_model_name
-        self.generation_config_dict = generation_config_dict # Will be passed to API call
-        self.safety_settings_list = safety_settings_list     # Will be passed to API call
+        self.generation_config_dict = generation_config_dict
+        self.safety_settings_list = safety_settings_list
         self.data_privacy = data_privacy
         self.force_sandbox = force_sandbox
         self.client = None
-        self.model_service = None # This will be client.models
+        self.model_service = None
+
+        if _REAL_GENAI_LOADED and GEMINI_API_KEY:
+            try:
+                self.client = genai.Client() # Should be the 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
+        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.")
+            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.")
 
-        # Check if genai.Client is the real one or our dummy
-        is_real_genai_client = hasattr(genai, 'Client') and not (hasattr(genai.Client, '__func__') and genai.Client.__func__.__qualname__.startswith('genai.Client'))
 
-        if not GEMINI_API_KEY and is_real_genai_client: # Real client but no API key
-            logging.warning(f"PandasLLM: GEMINI_API_KEY not set, but real 'genai.Client' seems available. API calls may fail if global config is not sufficient.")
-            # Proceed to initialize client; it might work if genai.configure() was successful without explicit key here
-            # or if the environment provides credentials in another way.
+    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
         
-        try:
-            self.client = genai.Client() # API key is usually set via genai.configure or environment
-            self.model_service = self.client.models
-            logging.info(f"PandasLLM: Initialized with genai.Client().models for '{self.llm_model_name}'.")
-        except Exception as e:
-            logging.error(f"Failed to initialize PandasLLM with genai.Client: {e}", exc_info=True)
-            # Fallback to dummy if real initialization fails, to prevent crashes
-            if not is_real_genai_client: # If it was already the dummy, re-initialize dummy
-                 self.client = genai.Client() 
-                 self.model_service = self.client.models
-                 logging.warning("PandasLLM: Falling back to DUMMY genai.Client due to real initialization error or it was already dummy.")
+        # 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 active_model_service:
+             logging.error("PandasLLM: Model service not available (real or dummy). Cannot call API.")
+             return "# Error: Gemini model service not available."
 
-    async def _call_gemini_api_async(self, prompt_text: str, history: list = None) -> str:
-        if not self.model_service:
-            logging.error("PandasLLM: Model service (client.models) not available. Cannot call API.")
-            return "# Error: Gemini model service not available for API call."
 
         gemini_history = []
-        if history: # history is agent.chat_history, which will now have 'content'
+        if history:
             for entry in history:
-                # Standardize role for API: 'model' for assistant responses
                 role_for_api = "model" if entry.get("role") == "assistant" else entry.get("role", "user")
-                # Get text from 'content' key
                 text_content = entry.get("content", "") 
                 gemini_history.append({"role": role_for_api, "parts": [{"text": text_content}]})
         
         current_prompt_content = [{"role": "user", "parts": [{"text": prompt_text}]}]
         contents_for_api = gemini_history + current_prompt_content
         
-        # Prepare model ID (e.g., "models/gemini-2.0-flash")
         model_id_for_api = self.llm_model_name
         if not model_id_for_api.startswith("models/"):
             model_id_for_api = f"models/{model_id_for_api}"
 
-        # Prepare generation config object
         api_generation_config = None
         if self.generation_config_dict:
-            try:
+            try: # genai_types is now consistently real or dummy alias
                 api_generation_config = genai_types.GenerationConfig(**self.generation_config_dict)
             except Exception as e_cfg:
-                logging.error(f"Error creating GenerationConfig object: {e_cfg}. Using dict as fallback.")
-                api_generation_config = self.generation_config_dict # Fallback to dict
+                logging.error(f"Error creating GenerationConfig object (real_loaded: {_REAL_GENAI_LOADED}): {e_cfg}. Using dict fallback.")
+                api_generation_config = self.generation_config_dict
 
-        logging.info(f"\n--- Calling Gemini API via Client (model: {model_id_for_api}) ---\nConfig: {api_generation_config}\nSafety: {bool(self.safety_settings_list)}\nContent (last part text): {contents_for_api[-1]['parts'][0]['text'][:100]}...\n")
+        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:
-            response = await self.model_service.generate_content_async(
+            # 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,
                 generation_config=api_generation_config,
@@ -340,14 +328,16 @@ class PandasLLM:
             llm_output = ""
             if hasattr(response, 'text') and isinstance(response.text, str):
                 llm_output = response.text
-            elif response.candidates:
+            elif response.candidates: # Standard way to get text from Gemini response
                 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
-                    finish_reason_str = str(finish_reason_val.name if hasattr(finish_reason_val, 'name') else finish_reason_val)
+                    # 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 = []
@@ -366,15 +356,16 @@ class PandasLLM:
                 return f"# Error: Unexpected API response structure: {str(response)[:200]}"
             
             return llm_output
-
-        except genai_types.BlockedPromptException as bpe:
+        
+        # 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:
+        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}"
         except Exception as e:
-            logging.error(f"Error calling Gemini API via Client: {e}", exc_info=True)
+            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]}."
 
 
@@ -395,10 +386,12 @@ class PandasLLM:
                 except IndexError: code_to_execute = ""
 
             if llm_response_text.startswith("# Error:") or not code_to_execute.strip():
-                logging.warning(f"LLM error or no code: {llm_response_text[:200]}")
+                # 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 "```" not in llm_response_text and len(llm_response_text.strip()) > 0:
-                        logging.info(f"LLM text output in sandbox mode: {llm_response_text[:200]}")
+                    # 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
+                        logging.info(f"LLM produced text output instead of Python code in sandbox mode. Passing through: {llm_response_text[:200]}")
                 return llm_response_text
 
             logging.info(f"\n--- Code to Execute: ---\n{code_to_execute}\n----------------------\n")
@@ -409,19 +402,19 @@ class PandasLLM:
             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.")
+                    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."
-                    return "# Code executed, but no print() output. Ensure print() for results."
+                        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 final_output_str
             except Exception as e:
-                logging.error(f"Sandbox Exec Error: {e}\nCode:\n{code_to_execute}", exc_info=True)
+                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 Exec Error: {type(e).__name__}: {e}\n# Code:\n{indented_code}"
+                return f"# Sandbox Execution Error: {type(e).__name__}: {e}\n# --- Code that caused error: ---\n{indented_code}"
             finally: sys.stdout = old_stdout
         else: return llm_response_text
 
@@ -439,8 +432,8 @@ class EmployerBrandingAgent:
         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 = [] # This will store entries like {"role": "user/assistant", "content": "text"}
-        logging.info("EmployerBrandingAgent Initialized (using Client API approach).")
+        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"
@@ -466,11 +459,16 @@ 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"]
+        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 ---\n{rag_context}\n"
-        else: prompt += "\n--- RAG CONTEXT ---\nNo specific RAG context found or RAG error.\n"
+        
+        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"
@@ -485,21 +483,12 @@ class EmployerBrandingAgent:
         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[:] # History before this turn
-        
-        # Append new user message to chat_history using 'content' key
+        current_turn_history_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]}...")
-        
-        # Pass history (which now uses 'content') to pandas_llm.query
+        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)
-        
-        # Append new assistant message to chat_history using 'content' key
-        # Standardize role to 'assistant' for model responses in history
         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):]
@@ -515,12 +504,16 @@ class EmployerBrandingAgent:
 
 # --- Example Usage (Conceptual) ---
 async def main_test():
-    logging.info("Starting main_test for EmployerBrandingAgent...")
+    # 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]})
     df_posts = pd.DataFrame({'post_id': [1], 'theme': ['Culture'], 'engagement_rate': [0.05]})
     test_dataframes = {"follower_stats": df_follower_stats, "posts": df_posts}
 
-    if not GEMINI_API_KEY: logging.warning("GEMINI_API_KEY not set. Testing with dummy functionality.")
+    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!"]
@@ -528,12 +521,24 @@ 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 GEMINI_API_KEY: await asyncio.sleep(1)
+        if _REAL_GENAI_LOADED and GEMINI_API_KEY: await asyncio.sleep(0.1) # Small delay for real API
 
 if __name__ == "__main__":
-    if GEMINI_API_KEY:
-        try: asyncio.run(main_test())
-        except RuntimeError as e:
-            if "asyncio.run() cannot be called from a running event loop" in str(e): print("Skip asyncio.run in existing loop.")
-            else: raise
-    else: print("GEMINI_API_KEY not set. Skipping main_test().")
+    # 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:
+        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}")
+