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 = [ 
        genai_types.SafetySetting(
            category=genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
            threshold=genai_types.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE, 
        ),
        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.")
    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..."]
} 
df_rag_documents = pd.DataFrame(rag_documents_data)

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


# --- Advanced RAG System ---
class AdvancedRAGSystem: # Truncated for brevity, assume correct from previous versions
    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 
        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() 
                self.embeddings_generated = True
                logging.info("RAG embeddings precomputed.")
            except Exception as e: logging.error(f"RAG precomputation error: {e}")
        else:
            logging.warning("RAG embeddings not precomputed (API key or genai.embed_content issue).")

    def _embed_fn(self, title: str, text: str) -> list[float]: 
        if not self.embeddings_generated: return [0.0] * 768
        try:
            return genai.embed_content(model=self.embedding_model_name, content=text, task_type="retrieval_document", title=title)["embedding"]
        except Exception as e:
            logging.error(f"Error in _embed_fn for '{title}': {e}")
            return [0.0] * 768
    def _precompute_embeddings(self): 
        if 'Embeddings' not in self.documents_df.columns:
            self.documents_df['Embeddings'] = pd.Series(dtype='object')
        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: 
        if not self.embeddings_generated or self.documents_df['Embeddings'].isnull().all():
            return "\n[RAG Context]\nEmbeddings not generated or all are null.\n"
        # Simplified retrieval logic for brevity
        try:
            query_embedding = np.array(genai.embed_content(model=self.embedding_model_name, content=query_text, task_type="retrieval_query")["embedding"])
            # Filter out rows with invalid embeddings before stacking
            valid_embeddings_df = self.documents_df.dropna(subset=['Embeddings'])
            valid_embeddings_df = valid_embeddings_df[valid_embeddings_df['Embeddings'].apply(lambda x: isinstance(x, list) and len(x) > 0)]

            if valid_embeddings_df.empty: return "\n[RAG Context]\nNo valid document embeddings for RAG.\n"
            
            document_embeddings = np.stack(valid_embeddings_df['Embeddings'].apply(np.array).values)
            if query_embedding.shape[0] != document_embeddings.shape[1]: return "\n[RAG Context]\nEmbedding dimension mismatch.\n"

            dot_products = np.dot(document_embeddings, query_embedding)
            idx = np.argsort(dot_products)[-min(top_k, len(valid_embeddings_df)):][::-1]
            
            relevant_passages = "".join([f"\n[RAG Context from: '{valid_embeddings_df.iloc[i]['Title']}']\n{valid_embeddings_df.iloc[i]['Text']}\n" for i in idx])
            return relevant_passages if relevant_passages else "\n[RAG Context]\nNo relevant passages found.\n"
        except Exception as e:
            logging.error(f"Error in RAG retrieve_relevant_info: {e}")
            return f"\n[RAG Context]\nError during RAG retrieval: {e}\n"


# --- PandasLLM Class (Gemini-Powered) ---
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 = 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}]})
        
        api_config_object = None
        try:
            api_config_object = genai_types.GenerateContentConfig(
                **self.generation_config_dict, 
                safety_settings=self.safety_settings_list 
            )
        except Exception as e_cfg:
            logging.error(f"Error creating GenerateContentConfig object: {e_cfg}.")
            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}) with config: {api_config_object} ---\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,
                config=api_config_object 
            )

            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: 
                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:
            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:
                        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:
                # If LLM returns an error or no code, pass that through directly.
                # The user will see the LLM's error message or its non-code response.
                logging.warning(f"LLM response is an error or not code: {llm_response_text}")
                return llm_response_text 

            logging.info(f"\n--- Code to Execute: ---\n{code_to_execute}\n----------------------\n")
            from io import StringIO
            import sys
            old_stdout = sys.stdout; sys.stdout = captured_output = StringIO()
            # Ensure dataframes_dict is correctly populated for exec_globals
            exec_globals = {'pd': pd, 'np': np} 
            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}' in dataframes_dict is not a DataFrame. Skipping for exec_globals.")

            try:
                exec(code_to_execute, exec_globals, {})
                final_output_str = captured_output.getvalue()
                # Check if the output is just whitespace or truly empty
                if not final_output_str.strip(): # If only whitespace or empty
                    # This is where the "no print output" message originates.
                    # We can now add a more informative message if the code itself ran without error.
                    logging.info("Code executed successfully, but no explicit print() output was generated by the LLM's code.")
                    return "# Code executed successfully, but it did not produce any printed output. Please ensure the LLM's Python code includes print() statements for the desired results."
                return final_output_str
            except Exception as e:
                logging.error(f"Sandbox Execution Error: {e}\nCode was:\n{code_to_execute}", exc_info=False)
                return f"# Sandbox Execution Error: {type(e).__name__}: {e}\n# --- Code that caused error: ---\n{textwrap.indent(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, 
                 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 
        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 = f"You are a helpful and expert '{role}'. Your primary goal is to assist with analyzing LinkedIn-related data using Pandas DataFrames.\n"
        prompt += "You will be provided with schemas for available Pandas DataFrames and a user query.\n"

        if self.pandas_llm.data_privacy:
            prompt += "IMPORTANT: Be mindful of data privacy. Do not output raw Personally Identifiable Information (PII) like names or specific user details unless explicitly asked and absolutely necessary for the query. Summarize or aggregate data where possible.\n"

        if self.pandas_llm.force_sandbox:
            prompt += "Your main task is to GENERATE PYTHON CODE using the Pandas library to answer the user query based on the provided DataFrames. Output ONLY the Python code block.\n"
            prompt += "The available DataFrames are already loaded and can be accessed by their dictionary keys prefixed with 'df_' (e.g., df_follower_stats, df_posts) within the execution environment.\n"
            prompt += "Example of accessing a DataFrame: `df_follower_stats['country']`.\n"
            prompt += "CRITICAL INSTRUCTION: Your Python code MUST include `print()` statements for ANY results, DataFrames, or values that should be displayed as the answer to the user's query. The output of these `print()` statements will be the final answer shown to the user.\n"
            prompt += "If you define a function to perform the analysis, you MUST call this function with the appropriate DataFrame(s) and `print()` its returned value. Do not just define functions without executing them and printing their results.\n"
            prompt += "If the query is simple and the result is a single value or a small piece of information, compute it and `print()` it directly.\n"
            prompt += "For example, if asked for 'total followers', your code should end with something like `print(total_followers)` or `print(df_result.to_string())`.\n"

            prompt += "If a column contains lists (e.g., 'skills' in a hypothetical 'df_employees'), you might need to use methods like `.explode()` or `.apply(pd.Series)` or `.apply(lambda x: ...)` for analysis.\n"
            prompt += "If the query is ambiguous or requires clarification, ask for it instead of making assumptions. If the query cannot be answered with the given data, state that clearly in a comment within the code block (e.g. `# Cannot answer: data not available`).\n"
            prompt += "If the query is not about data analysis or code generation (e.g. 'hello', 'how are you?'), respond politely and briefly in a comment, do not attempt to generate code (e.g. `# Hello there! How can I help you with data analysis today?`).\n"
            prompt += "Structure your code clearly. Add comments (#) to explain each step of your logic.\n"
        else: 
            prompt += "Your task is to analyze the data and provide a comprehensive textual answer to the user query. You can explain your reasoning step-by-step.\n"

        prompt += "\n--- AVAILABLE DATA AND SCHEMAS ---\n"
        prompt += self.schemas_representation
        
        rag_context = self.rag_system.retrieve_relevant_info(user_query)
        if rag_context and "[RAG Context]" in rag_context and "No specific pre-defined context found" not in rag_context and "No highly relevant passages found" not in rag_context and "Embeddings not generated" not in rag_context:
            prompt += f"\n--- ADDITIONAL CONTEXT (from internal knowledge base, consider this information) ---\n{rag_context}\n"

        prompt += f"\n--- USER QUERY ---\n{user_query}\n"

        if task_decomposition_hint:
            prompt += f"\n--- GUIDANCE FOR ANALYSIS (Task Decomposition) ---\n{task_decomposition_hint}\n"
        
        if cot_hint:
            if self.pandas_llm.force_sandbox:
                prompt += "\n--- INSTRUCTIONS FOR PYTHON CODE GENERATION (Chain of Thought & Output) ---\n"
                prompt += "1. Understand the query: What specific information is requested?\n"
                prompt += "2. Identify relevant DataFrame(s) and column(s) from the schemas provided.\n"
                prompt += "3. Plan the steps: Outline the Pandas operations needed (filtering, grouping, aggregation, merging, etc.) as comments in your code.\n"
                prompt += "4. Write the code: Implement the steps using Pandas. Remember to use `df_name_of_dataframe` (e.g. `df_follower_stats`).\n"
                prompt += "5. CRITICAL - Ensure output: Call any functions you define and use `print()` for ALL results that should be displayed. For DataFrames, you can print the DataFrame directly (e.g., `print(my_result_df)`), or `print(df.to_string())` if it might be large. For single values, `print(my_value)`.\n"
                prompt += "6. Review: Check for correctness, efficiency, and adherence to the prompt (especially the CRITICAL `print()` requirement for the final answer).\n"
                prompt += "7. Generate ONLY the Python code block starting with ```python and ending with ```. No explanations outside the code block's comments.\n"
            else: 
                prompt += "\n--- INSTRUCTIONS FOR RESPONSE (Chain of Thought) ---\n"
                prompt += "Please provide a step-by-step explanation of your analysis before giving the final answer.\n"
        
        return prompt

    async def process_query(self, user_query: str, role="Employer Branding Analyst", task_decomposition_hint=None, cot_hint=True) -> str:
        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"Full prompt to LLM (last 300 chars of user query part for brevity in log): ... {full_prompt[-500:]}") # Log end of prompt

        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})
        
        MAX_HISTORY_TURNS = 5 
        if len(self.chat_history) > MAX_HISTORY_TURNS * 2:
            self.chat_history = self.chat_history[-(MAX_HISTORY_TURNS * 2):]
            
        return response_text

    def update_dataframes(self, new_dataframes: dict): 
        self.all_dataframes = new_dataframes
        self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
        logging.info("EmployerBrandingAgent DataFrames updated.")
    def clear_chat_history(self): 
        self.chat_history = []
        logging.info("EmployerBrandingAgent chat history cleared.")