Update eb_agent_module.py

eb_agent_module.py

@@ -1,4 +1,3 @@
-# eb_agent_module.py
 import pandas as pd
 import json
 import os
@@ -7,695 +6,103 @@ import logging
 import numpy as np
 import textwrap
 
-# --- 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}")
-        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()
-
-    def embed_content(self, model=None, contents=None, config=None): # Added dummy embed_content
-        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): # 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")
-        print(f"Dummy _DummyGenAI.Client initialized {'with api_key in client_options' if api_key_present_in_options else '(global API key expected by dummy)'}.")
-
-
-class _DummyGenAIGenerativeModel: # This dummy might be less used if client.models is preferred
-    def __init__(self, model_name_in, generation_config=None, safety_settings=None, system_instruction=None):
-        self.model_name = model_name_in
-        print(f"Dummy _DummyGenAIGenerativeModel initialized for {model_name_in}")
-    async def generate_content_async(self, contents, stream=False):
-        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()
-    # This embed_content on the dummy GenerativeModel might not be used if AdvancedRAGSystem uses client.models.embed_content
-    def embed_content(self, content, task_type=None, title=None):
-        print(f"Dummy _DummyGenAIGenerativeModel.embed_content called for model {self.model_name} (task: {task_type})")
-        return {"embedding": [0.1] * 768}
-
-
-class _ActualDummyGenAI: # type: ignore
-    Client = _DummyGenAIClient 
-
-    @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)
-    
-    class types:
-        @staticmethod
-        def GenerationConfig(**kwargs):
-            print(f"Dummy _ActualDummyGenAI.types.GenerationConfig created with: {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}")
-            # 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):
-            print(f"Dummy _ActualDummyGenAI.types.EmbedContentConfig created with task_type: {task_type}")
-            conf = {}
-            if task_type: conf["task_type"] = task_type
-            if output_dimensionality: conf["output_dimensionality"] = output_dimensionality
-            if title: conf["title"] = title
-            # 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"
-        class HarmBlockThreshold: BLOCK_NONE = "BLOCK_NONE"; BLOCK_LOW_AND_ABOVE = "BLOCK_LOW_AND_ABOVE"; BLOCK_MEDIUM_AND_ABOVE = "BLOCK_MEDIUM_AND_ABOVE"; BLOCK_ONLY_HIGH = "BLOCK_ONLY_HIGH"
-        class FinishReason: FINISH_REASON_UNSPECIFIED = "UNSPECIFIED"; STOP = "STOP"; MAX_TOKENS = "MAX_TOKENS"; SAFETY = "SAFETY"; RECITATION = "RECITATION"; OTHER = "OTHER"
-        class BlockedReason: BLOCKED_REASON_UNSPECIFIED = "BLOCKED_REASON_UNSPECIFIED"; SAFETY = "SAFETY"; OTHER = "OTHER"
-        class BlockedPromptException(Exception): pass
-        class StopCandidateException(Exception): pass
-
-
-# --- Attempt to import the real library ---
-_REAL_GENAI_LOADED = False
-genai_types = None 
-
 try:
-    # 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.generativeai' and accessed 'genai.types'.")
+    from google import genai
+    from google.genai import types as genai_types
 except ImportError:
-    genai = _ActualDummyGenAI() 
-    genai_types = genai.types 
-    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.generativeai' imported, but 'genai.types' not found. Falling back to dummy implementations.")
+    print("Google Generative AI library not found. Please install it: pip install google-generativeai")
+    # Dummy classes defined here for development/debugging
+    ...  # KEEP YOUR EXISTING DUMMY DEFINITIONS
 
+# 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"
 
-# --- Configuration ---
-GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', "")
-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
+client = genai.Client(api_key=GEMINI_API_KEY)
 
-GENERATION_CONFIG_PARAMS = {
-    "temperature": 0.3, "top_p": 1.0, "top_k": 32, "max_output_tokens": 8192,
-}
-
-try:
-    DEFAULT_SAFETY_SETTINGS = [
-        genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
-        genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
-        genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
-        genai_types.SafetySetting(category=genai_types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=genai_types.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE),
-    ]
-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 = [ # 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')
-
-if _REAL_GENAI_LOADED:
-    if GEMINI_API_KEY:
-        try:
-            genai.configure(api_key=GEMINI_API_KEY) 
-            logging.info(f"Gemini API key configured globally using REAL genai.configure.")
-        except Exception as e:
-            logging.error(f"Failed to configure REAL Gemini API globally: {e}", exc_info=True)
-    else:
-        logging.warning("REAL 'google.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 about best practices...", "Projects showcasing tech talent..."] }
-df_rag_documents = pd.DataFrame(rag_documents_data)
-
-# --- Schema Representation ---
-def get_schema_representation(df_name: str, df: pd.DataFrame) -> str:
-    if not isinstance(df, pd.DataFrame): return f"Schema for item '{df_name}': Not a DataFrame.\n"
-    if df.empty: return f"Schema for DataFrame 'df_{df_name}': Empty.\n"
-    return f"DataFrame 'df_{df_name}': Cols: {df.columns.tolist()}, Shape: {df.shape}\nSample:\n{textwrap.indent(df.head(1).to_string(), '    ')}\n"
-
-def get_all_schemas_representation(dataframes_dict: dict) -> str:
-    if not dataframes_dict: return "No DataFrames provided.\n"
-    return "".join(get_schema_representation(name, df) for name, df in dataframes_dict.items())
-
-# --- 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 
-        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_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:
-                # 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 (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_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
-            
-            # 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.
-            
-            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}, 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')
-        # 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 = 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)
-        
-        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.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:
-            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 available after filtering.\n"
-            
-            doc_embeddings = np.stack(valid_df['Embeddings'].apply(np.array).values)
-            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"
-            
-            # 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:
-            logging.error(f"Error in RAG retrieve_relevant_info (real mode with embedding service): {e}", exc_info=True)
-            return f"\n[RAG Context]\nError during RAG retrieval (real mode): {type(e).__name__} - {e}\n"
-
-# --- PandasLLM Class (Gemini-Powered using genai.Client) ---
-class PandasLLM:
-    def __init__(self, llm_model_name: str,
-                 generation_config_dict: dict,
-                 safety_settings_list: list,
-                 data_privacy=True, force_sandbox=True):
-        self.llm_model_name = llm_model_name
-        self.generation_config_dict = generation_config_dict
-        self.safety_settings_list = safety_settings_list
-        self.data_privacy = data_privacy
-        self.force_sandbox = force_sandbox
-        self.client_instance = None # Stores the genai.Client() instance
-        self.model_service_from_client = None # Stores client.models
-
-        # 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:
-                # 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().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:
-        # active_model_service will be self.model_service_from_client (either real or dummy)
-        active_model_service = self.model_service_from_client
-        
-        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 (client.models) not available. Cannot call API.")
-            return "# Error: Gemini model service not available for API call."
-
-        gemini_history = []
-        if history:
-            for entry in history:
-                role_for_api = "model" if entry.get("role") == "assistant" else entry.get("role", "user")
-                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
-        
-        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}"
-
-        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}, 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}, 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(
-                model=model_id_for_api,
-                contents=contents_for_api,
-                generation_config=api_generation_config,
-                safety_settings=self.safety_settings_list
+        self.documents_df = documents_df
+        self.embedding_model_name = embedding_model_name
+        self.embeddings = self._embed_documents()
+
+    def _embed_documents(self):
+        embedded_docs = []
+        for text in self.documents_df['text']:
+            response = client.models.embed_content(
+                model=self.embedding_model_name,
+                contents=text,
+                config=genai_types.EmbedContentConfig(task_type="SEMANTIC_SIMILARITY")
             )
+            embedded_docs.append(np.array(response.embeddings.values))
+        return np.vstack(embedded_docs)
+
+    def retrieve_relevant_info(self, query: str, top_k=3) -> str:
+        query_embedding = client.models.embed_content(
+            model=self.embedding_model_name,
+            contents=query,
+            config=genai_types.EmbedContentConfig(task_type="SEMANTIC_SIMILARITY")
+        )
+        query_vector = np.array(query_embedding.embeddings.values)
+
+        scores = np.dot(self.embeddings, query_vector)
+        top_indices = np.argsort(scores)[-top_k:][::-1]
+        context = "\n\n".join(self.documents_df.iloc[i]['text'] for i in top_indices)
+        return context
 
-            # 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
-                # 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 = ""
-            # 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" 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: # 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
-        
-        # 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 (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}"
-            if not (_REAL_GENAI_LOADED and type(sce).__name__ == 'StopCandidateException'): raise
-        except Exception as e:
-            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]}."
-
-
-    async def query(self, prompt_with_query_and_context: str, dataframes_dict: dict, history: list = None) -> str:
-        llm_response_text = await self._call_gemini_api_async(prompt_with_query_and_context, history)
-
-        if self.force_sandbox:
-            code_to_execute = ""
-            # Improved code extraction to handle potential leading/trailing newlines or spaces
-            if "```python" in llm_response_text:
-                try:
-                    # 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 original LLM response (error or non-code)
+class EmployerBrandingAgent:
+    def __init__(self, all_dataframes: dict, rag_documents_df: pd.DataFrame):
+        self.all_dataframes = all_dataframes
+        self.schemas_representation = self._get_all_schemas_representation()
+        self.chat_history = []
+        self.rag_system = AdvancedRAGSystem(rag_documents_df, GEMINI_EMBEDDING_MODEL_NAME)
+        logging.info("EmployerBrandingAgent initialized with Gemini")
 
-            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: # 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, {}) # Execute in the prepared scope
-                final_output_str = sys.stdout.getvalue()
-                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 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 # Restore stdout
-        else: # Not forcing sandbox, return raw LLM response
-            return llm_response_text
+    def _get_all_schemas_representation(self):
+        schema_descriptions = []
+        for key, df in self.all_dataframes.items():
+            schema = f"DataFrame: df_{key}\nColumns: {', '.join(df.columns)}\n"
+            schema_descriptions.append(schema)
+        return "\n".join(schema_descriptions)
 
-# --- Employer Branding Agent ---
-class EmployerBrandingAgent:
-    def __init__(self, llm_model_name: str, 
-                 generation_config_dict: dict, 
-                 safety_settings_list: list, 
-                 all_dataframes: dict, 
-                 rag_documents_df: pd.DataFrame, 
-                 embedding_model_name: str, 
-                 data_privacy=True, 
-                 force_sandbox=True):
-        
-        self.pandas_llm = PandasLLM(llm_model_name, generation_config_dict, safety_settings_list, data_privacy, force_sandbox)
-        self.rag_system = AdvancedRAGSystem(rag_documents_df, embedding_model_name)
-        self.all_dataframes = all_dataframes if all_dataframes else {}
-        self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
-        self.chat_history = [] 
-        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) -> str:
+        prompt = f"You are an expert Employer Branding Analyst. Analyze the query based on the following DataFrames.\n"
+        prompt += self.schemas_representation
 
-    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}'. 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 += "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)
-        # 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"--- 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"
+        if rag_context:
+            prompt += f"\n\nAdditional Context:\n{rag_context}"
+
+        prompt += f"\n\nUser Query:\n{user_query}"
         return prompt
 
-    async def process_query(self, user_query: str, role="EB Analyst", task_hint=None, cot=True) -> str: # Simplified role
-        # 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
+    async def process_query(self, user_query: str) -> str:
         self.chat_history.append({"role": "user", "content": user_query})
-        
-        prompt = self._build_prompt(user_query, role, task_hint, cot)
-        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})
-        
-        # 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(f"Agent DataFrames updated. New schema: {self.schemas_representation[:200]}...")
-    
-    def clear_chat_history(self): 
-        self.chat_history = []
-        logging.info("Agent chat history cleared.")
-
-# --- Example Usage (Conceptual) ---
-async def main_test():
-    # 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={'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 # Set to True to test code generation, False for direct text
-    )
-    
-    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"--- 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 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')
+        prompt = self._build_prompt(user_query)
+
+        response = client.models.generate_content(
+            model=LLM_MODEL_NAME,
+            contents=[prompt],
+            config=genai_types.GenerateContentConfig(
+                safety_settings=[
+                    genai_types.SafetySetting(
+                        category=genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+                        threshold=genai_types.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
+                    )
+                ]
+            )
+        )
 
-    try: 
-        asyncio.run(main_test())
-    except RuntimeError as e:
-        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())
+        answer = response.text.strip()
+        self.chat_history.append({"role": "assistant", "content": answer})
+        return answer
 
-        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)
+    def update_dataframes(self, new_dataframes: dict):
+        self.all_dataframes = new_dataframes
+        self.schemas_representation = self._get_all_schemas_representation()
+        logging.info("EmployerBrandingAgent DataFrames updated.")
 
+    def clear_chat_history(self):
+        self.chat_history = []
+        logging.info("EmployerBrandingAgent chat history cleared.")