Update eb_agent_module.py

eb_agent_module.py

@@ -103,10 +103,10 @@ GENERATION_CONFIG_PARAMS = {
 # 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
+    DEFAULT_SAFETY_SETTINGS = [ 
         genai_types.SafetySetting(
             category=genai_types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
-            threshold=genai_types.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE, # As per user example
+            threshold=genai_types.HarmBlockThreshold.BLOCK_LOW_AND_ABOVE, 
         ),
         genai_types.SafetySetting(
             category=genai_types.HarmCategory.HARM_CATEGORY_HARASSMENT,
@@ -123,7 +123,6 @@ try:
     ]
 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"},
@@ -149,43 +148,75 @@ else:
 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) ---
+# --- 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"
-    return f"Schema for DataFrame 'df_{df_name}': {df.columns.tolist()[:3]}...\nSample:\n{df.head(1).to_string()}\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:
-    return "".join(get_schema_representation(name, df) for name, df in dataframes_dict.items())
+    # 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 (truncated for brevity) ---
-class AdvancedRAGSystem:
+
+# --- 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 # Simplified
+        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() # Simplified
+                self._precompute_embeddings() 
                 self.embeddings_generated = True
+                logging.info("RAG embeddings precomputed.")
             except Exception as e: logging.error(f"RAG precomputation error: {e}")
-    def _embed_fn(self, title: str, text: str) -> list[float]: # Simplified
+        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
-        return genai.embed_content(model=self.embedding_model_name, content=text, task_type="retrieval_document", title=title)["embedding"]
-    def _precompute_embeddings(self): # Simplified
+        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: # Simplified
-        if not self.embeddings_generated: return "\n[RAG Context]\nEmbeddings not generated.\n"
+    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
-        return f"\n[RAG Context]\nRetrieved info for: {query_text} (Top {top_k})\n"
+        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, # Base config: temp, top_k, etc.
-                 safety_settings_list: list,   # List of SafetySetting objects/dicts
+                 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 
@@ -222,58 +253,37 @@ class PandasLLM:
                 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.
+            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}) ---\n")
+        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 # 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
+                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: # Standard structure
+            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'))
@@ -281,23 +291,19 @@ class PandasLLM:
                      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
+                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:]
@@ -305,30 +311,46 @@ class PandasLLM:
                     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}"
+                # 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")
-            # 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
+            # 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()
-                return final_output_str if final_output_str else "# Code executed, no print output."
+                # 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:
-                return f"# Sandbox Execution Error: {e}\nCode:\n{code_to_execute}"
-            finally: sys.stdout = old_stdout
+                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, # Base config (temp, top_k)
-                 safety_settings_list: list,   # List of SafetySetting objects/dicts
+                 generation_config_dict: dict, 
+                 safety_settings_list: list,   
                  all_dataframes: dict, 
                  rag_documents_df: pd.DataFrame, 
                  embedding_model_name: str,
@@ -337,7 +359,7 @@ class EmployerBrandingAgent:
         self.pandas_llm = PandasLLM(
             llm_model_name, 
             generation_config_dict, 
-            safety_settings_list, # Pass the list here
+            safety_settings_list, 
             data_privacy, 
             force_sandbox
         )
@@ -348,24 +370,76 @@ class EmployerBrandingAgent:
         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 = 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
-        prompt += f"User Query: {user_query}\n"
-        prompt += "Generate Python code using Pandas...\n"
+        
+        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:
-        # ... (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)
+        
+        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})
-        # Limit history
-        if len(self.chat_history) > 10: self.chat_history = self.chat_history[-10:]
+        
+        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): # Simplified
+    def update_dataframes(self, new_dataframes: dict): 
         self.all_dataframes = new_dataframes
         self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
-    def clear_chat_history(self): self.chat_history = []
+        logging.info("EmployerBrandingAgent DataFrames updated.")
+    def clear_chat_history(self): 
+        self.chat_history = []
+        logging.info("EmployerBrandingAgent chat history cleared.")
+