eb_agent_module.py

# eb_agent_module.py
import pandas as pd
import json
import os
import asyncio
import logging
import numpy as np
import textwrap

# Attempt to import Google Generative AI and related types
try:
    from google import genai
    from google.genai import types as genai_types
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): pass
        
        @staticmethod
        def Client(api_key=None): 
            class DummyModels:
                @staticmethod
                def generate_content(model=None, contents=None, config=None, safety_settings=None): # Added config, kept safety_settings for older dummy
                    print(f"Dummy genai.Client.models.generate_content called for model: {model} with config: {config}, safety_settings: {safety_settings}")
                    class DummyPart:
                        def __init__(self, text): self.text = text
                    class DummyContent:
                        def __init__(self): self.parts = [DummyPart("# Dummy response from dummy client")]
                    class DummyCandidate:
                        def __init__(self):
                            self.content = DummyContent()
                            self.finish_reason = "DUMMY"
                            self.safety_ratings = [] # Ensure this attribute exists
                    class DummyResponse:
                        def __init__(self):
                            self.candidates = [DummyCandidate()]
                            self.prompt_feedback = None # Ensure this attribute exists
                        @property
                        def text(self): 
                            if self.candidates and self.candidates[0].content and self.candidates[0].content.parts:
                                return "".join(p.text for p in self.candidates[0].content.parts)
                            return ""
                    return DummyResponse()

            class DummyClient:
                def __init__(self): self.models = DummyModels()
            
            if api_key: return DummyClient()
            return None

        @staticmethod
        def GenerativeModel(model_name): 
             print(f"Dummy genai.GenerativeModel called for model: {model_name}")
             return None
        
        @staticmethod
        def embed_content(model, content, task_type, title=None): 
            print(f"Dummy genai.embed_content called for model: {model}")
            return {"embedding": [0.1] * 768} 

    class genai_types: # type: ignore
        @staticmethod
        def GenerateContentConfig(**kwargs): # The dummy now just returns the kwargs
            print(f"Dummy genai_types.GenerateContentConfig called with: {kwargs}")
            return kwargs 
        
        # Dummy SafetySetting to allow instantiation if real genai_types is missing
        @staticmethod
        def SafetySetting(category, threshold):
            print(f"Dummy SafetySetting created: category={category}, threshold={threshold}")
            return {"category": category, "threshold": threshold} # Return a dict for dummy

        class BlockReason: 
            SAFETY = "SAFETY"
        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"


# --- 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"

# Base generation configuration for the LLM (without safety settings here)
GENERATION_CONFIG_PARAMS = {
    "temperature": 0.2,
    "top_p": 1.0,
    "top_k": 32,
    "max_output_tokens": 4096, 
}

# Default safety settings list for Gemini
# This is now a list of SafetySetting objects (or dicts if using dummy)
try:
    DEFAULT_SAFETY_SETTINGS = [ # Renamed from DEFAULT_SAFETY_SETTINGS_LIST for consistency with app.py import
        genai_types.SafetySetting(
            category=genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
            threshold=genai_types.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE, # As per user example
        ),
        genai_types.SafetySetting(
            category=genai_types.HarmCategory.HARM_CATEGORY_HARASSMENT,
            threshold=genai_types.HarmBlockThreshold.BLOCK_NONE, 
        ),
        genai_types.SafetySetting(
            category=genai_types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
            threshold=genai_types.HarmBlockThreshold.BLOCK_NONE,
        ),
        genai_types.SafetySetting(
            category=genai_types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
            threshold=genai_types.HarmBlockThreshold.BLOCK_NONE,
        ),
    ]
except AttributeError as e: 
    logging.warning(f"Could not define DEFAULT_SAFETY_SETTINGS using real genai_types: {e}. Using placeholder list of dicts.")
    # Fallback to list of dicts if genai_types.SafetySetting or HarmCategory/HarmBlockThreshold are dummies that don't work as expected
    DEFAULT_SAFETY_SETTINGS = [
        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_LOW_AND_ABOVE"},
        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
        {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
        {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
    ]


# Logging setup
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(message)s')

if GEMINI_API_KEY:
    try:
        genai.configure(api_key=GEMINI_API_KEY)
        logging.info(f"Gemini API key configured globally...")
    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.")


# --- RAG Documents Definition ---
rag_documents_data = {
    'Title': ["Employer Branding Best Practices 2024", "Attracting Tech Talent"],
    'Text': ["Focus on authentic employee stories...", "Tech candidates value challenging projects..."]
} # Truncated for brevity
df_rag_documents = pd.DataFrame(rag_documents_data)

# --- Schema Representation (truncated for brevity) ---
def get_schema_representation(df_name: str, df: pd.DataFrame) -> str:
    if df.empty: return f"Schema for DataFrame '{df_name}': Empty.\n"
    return f"Schema for DataFrame 'df_{df_name}': {df.columns.tolist()[:3]}...\nSample:\n{df.head(1).to_string()}\n"
def get_all_schemas_representation(dataframes_dict: dict) -> str:
    return "".join(get_schema_representation(name, df) for name, df in dataframes_dict.items())

# --- Advanced RAG System (truncated for brevity) ---
class AdvancedRAGSystem:
    def __init__(self, documents_df: pd.DataFrame, embedding_model_name: str):
        self.embedding_model_name = embedding_model_name
        self.documents_df = documents_df.copy()
        self.embeddings_generated = False # Simplified
        if GEMINI_API_KEY and hasattr(genai, 'embed_content') and not (hasattr(genai.embed_content, '__func__') and genai.embed_content.__func__.__qualname__.startswith('genai.embed_content')):
            try:
                self._precompute_embeddings() # Simplified
                self.embeddings_generated = True
            except Exception as e: logging.error(f"RAG precomputation error: {e}")
    def _embed_fn(self, title: str, text: str) -> list[float]: # Simplified
        if not self.embeddings_generated: return [0.0] * 768
        return genai.embed_content(model=self.embedding_model_name, content=text, task_type="retrieval_document", title=title)["embedding"]
    def _precompute_embeddings(self): # Simplified
        self.documents_df['Embeddings'] = self.documents_df.apply(lambda row: self._embed_fn(row['Title'], row['Text']), axis=1)
    def retrieve_relevant_info(self, query_text: str, top_k: int = 1) -> str: # Simplified
        if not self.embeddings_generated: return "\n[RAG Context]\nEmbeddings not generated.\n"
        # Simplified retrieval logic for brevity
        return f"\n[RAG Context]\nRetrieved info for: {query_text} (Top {top_k})\n"


# --- PandasLLM Class (Gemini-Powered) ---
class PandasLLM:
    def __init__(self, llm_model_name: str, 
                 generation_config_dict: dict, # Base config: temp, top_k, etc.
                 safety_settings_list: list,   # List of SafetySetting objects/dicts
                 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 = None
        self.generative_model_service = None

        if not GEMINI_API_KEY:
            logging.warning("PandasLLM: GEMINI_API_KEY not set.")
        else:
            try:
                self.client = genai.Client(api_key=GEMINI_API_KEY) 
                if self.client and hasattr(self.client, 'models') and hasattr(self.client.models, 'generate_content'):
                    self.generative_model_service = self.client.models
                    logging.info(f"PandasLLM: Using client.models for '{self.llm_model_name}'.")
                elif self.client and hasattr(self.client, 'generate_content'):
                    self.generative_model_service = self.client
                    logging.info(f"PandasLLM: Using client.generate_content for '{self.llm_model_name}'.")
                else:
                    logging.warning(f"PandasLLM: genai.Client suitable 'generate_content' not found.")
            except Exception as e:
                logging.error(f"Failed to initialize PandasLLM with genai.Client: {e}", exc_info=True)

    async def _call_gemini_api_async(self, prompt_text: str, history: list = None) -> str:
        if not self.generative_model_service:
            return "# Error: Gemini client/service not available."

        contents_for_api = []
        if history:
            for entry in history:
                role = "model" if entry.get("role") == "assistant" else entry.get("role", "user")
                contents_for_api.append({"role": role, "parts": [{"text": entry.get("content", "")}]})
        contents_for_api.append({"role": "user", "parts": [{"text": prompt_text}]})
        
        # Prepare the full configuration object for the API call
        api_config_object = None
        try:
            # **self.generation_config_dict provides temperature, top_p, etc.
            # safety_settings takes the list of SafetySetting objects/dicts
            api_config_object = genai_types.GenerateContentConfig(
                **self.generation_config_dict, 
                safety_settings=self.safety_settings_list 
            )
            logging.debug(f"Constructed GenerateContentConfig object: {api_config_object}")
        except Exception as e_cfg:
            logging.error(f"Error creating GenerateContentConfig object: {e_cfg}. API call may fail or use defaults.")
            # Fallback: try to pass the raw dicts if GenerateContentConfig class itself fails (e.g. dummy issues)
            # This is less ideal as the API might strictly expect the object.
            api_config_object = {**self.generation_config_dict, "safety_settings": self.safety_settings_list}


        logging.info(f"\n--- Calling Gemini API via Client (model: {self.llm_model_name}) ---\n")
        
        try:
            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}"

            response = await asyncio.to_thread(
                self.generative_model_service.generate_content,
                model=model_id_for_api, 
                contents=contents_for_api,
                generation_config=api_config_object # Use 'generation_config' as it's common, but user example used 'config'.
                                                    # If 'client.models.generate_content' specifically needs 'config', change this.
                                                    # For now, assuming 'generation_config' is more standard for the object.
                                                    # UPDATE based on user's example: it should be 'config'
                                                    # config=api_config_object 
            )
            # Re-checking user's example: client.models.generate_content(..., config=types.GenerateContentConfig(...))
            # So, the parameter name should indeed be 'config'.

            response = await asyncio.to_thread(
                self.generative_model_service.generate_content,
                model=model_id_for_api,
                contents=contents_for_api,
                config=api_config_object # CORRECTED to 'config' based on user example
            )


            if hasattr(response, 'prompt_feedback') and response.prompt_feedback and response.prompt_feedback.block_reason:
                return f"# Error: Prompt blocked by API: {response.prompt_feedback.block_reason}."

            llm_output = ""
            if hasattr(response, 'text') and response.text:
                llm_output = response.text
            elif hasattr(response, 'candidates') and response.candidates: # Standard structure
                candidate = response.candidates[0]
                if hasattr(candidate, 'content') and candidate.content and hasattr(candidate.content, 'parts') and candidate.content.parts:
                    llm_output = "".join(part.text for part in candidate.content.parts if hasattr(part, 'text'))
                if not llm_output and hasattr(candidate, 'finish_reason'):
                     return f"# Error: Empty response. Finish reason: {candidate.finish_reason}."
            else:
                return f"# Error: Unexpected API response structure: {str(response)[:200]}"

            return llm_output

        except Exception as e:
            logging.error(f"Error calling Gemini API via Client: {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:
            # ... (sandbox execution logic - truncated for brevity, assumed correct from previous versions)
            code_to_execute = ""
            if "```python" in llm_response_text:
                try:
                    code_to_execute = llm_response_text.split("```python\n", 1)[1].split("\n```", 1)[0]
                except IndexError: # Try alternative split
                    try:
                        code_to_execute = llm_response_text.split("```python", 1)[1].split("```", 1)[0]
                        if code_to_execute.startswith("\n"): code_to_execute = code_to_execute[1:]
                        if code_to_execute.endswith("\n"): code_to_execute = code_to_execute[:-1]
                    except IndexError: code_to_execute = ""
            
            if llm_response_text.startswith("# Error:") or not code_to_execute:
                return f"# LLM Error or No Code: {llm_response_text}"

            logging.info(f"\n--- Code to Execute: ---\n{code_to_execute}\n----------------------\n")
            # Sandbox execution (simplified for brevity)
            from io import StringIO
            import sys
            old_stdout = sys.stdout; sys.stdout = captured_output = StringIO()
            exec_globals = {'pd': pd, 'np': np} # Simplified builtins for brevity
            for name, df in dataframes_dict.items(): exec_globals[f"df_{name}"] = df
            try:
                exec(code_to_execute, exec_globals, {})
                final_output_str = captured_output.getvalue()
                return final_output_str if final_output_str else "# Code executed, no print output."
            except Exception as e:
                return f"# Sandbox Execution Error: {e}\nCode:\n{code_to_execute}"
            finally: sys.stdout = old_stdout
        else: 
            return llm_response_text

# --- Employer Branding Agent ---
class EmployerBrandingAgent:
    def __init__(self, llm_model_name: str, 
                 generation_config_dict: dict, # Base config (temp, top_k)
                 safety_settings_list: list,   # List of SafetySetting objects/dicts
                 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, # Pass the list here
            data_privacy, 
            force_sandbox
        )
        self.rag_system = AdvancedRAGSystem(rag_documents_df, embedding_model_name)
        self.all_dataframes = all_dataframes 
        self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
        self.chat_history = [] 
        logging.info("EmployerBrandingAgent Initialized with updated safety settings handling.")

    def _build_prompt(self, user_query: str, role="Employer Branding Analyst", task_decomposition_hint=None, cot_hint=True) -> str:
        # ... (prompt building logic - truncated for brevity, assumed correct from previous versions)
        prompt = f"You are a helpful '{role}'...\n"
        prompt += self.schemas_representation
        prompt += f"User Query: {user_query}\n"
        prompt += "Generate Python code using Pandas...\n"
        return prompt

    async def process_query(self, user_query: str, role="Employer Branding Analyst", task_decomposition_hint=None, cot_hint=True) -> str:
        # ... (process query logic - truncated for brevity, assumed correct from previous versions)
        self.chat_history.append({"role": "user", "content": user_query})
        full_prompt = self._build_prompt(user_query, role, task_decomposition_hint, cot_hint)
        response_text = await self.pandas_llm.query(full_prompt, self.all_dataframes, history=self.chat_history[:-1]) 
        self.chat_history.append({"role": "assistant", "content": response_text})
        # Limit history
        if len(self.chat_history) > 10: self.chat_history = self.chat_history[-10:]
        return response_text

    def update_dataframes(self, new_dataframes: dict): # Simplified
        self.all_dataframes = new_dataframes
        self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
    def clear_chat_history(self): self.chat_history = []