Update eb_agent_module.py

eb_agent_module.py

@@ -295,13 +295,16 @@ class PandasLLM:
             return "# Error: Gemini model service not available for API call."
 
         gemini_history = []
-        if history:
+        if history: # history is agent.chat_history, which will now have 'content'
             for entry in history:
-                role = "model" if entry.get("role") == "assistant" else entry.get("role", "user")
-                gemini_history.append({"role": role, "parts": [{"text": entry.get("content", "")}]})
+                # Standardize role for API: 'model' for assistant responses
+                role_for_api = "model" if entry.get("role") == "assistant" else entry.get("role", "user")
+                # Get text from 'content' key
+                text_content = entry.get("content", "") 
+                gemini_history.append({"role": role_for_api, "parts": [{"text": text_content}]})
         
-        current_content = [{"role": "user", "parts": [{"text": prompt_text}]}]
-        contents_for_api = gemini_history + current_content
+        current_prompt_content = [{"role": "user", "parts": [{"text": prompt_text}]}]
+        contents_for_api = gemini_history + current_prompt_content
         
         # Prepare model ID (e.g., "models/gemini-2.0-flash")
         model_id_for_api = self.llm_model_name
@@ -327,10 +330,8 @@ class PandasLLM:
                 safety_settings=self.safety_settings_list
             )
 
-            # ... (Response parsing logic remains largely the same as before) ...
             if hasattr(response, 'prompt_feedback') and response.prompt_feedback and \
                hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason:
-                # ... block reason handling ...
                 block_reason_val = response.prompt_feedback.block_reason
                 block_reason_str = str(block_reason_val.name if hasattr(block_reason_val, 'name') else block_reason_val)
                 logging.warning(f"Prompt blocked by API. Reason: {block_reason_str}.")
@@ -345,14 +346,18 @@ class PandasLLM:
                     llm_output = "".join(part.text for part in candidate.content.parts if hasattr(part, 'text'))
                 
                 if not llm_output and candidate.finish_reason:
-                    # ... finish reason handling ...
                     finish_reason_val = candidate.finish_reason
                     finish_reason_str = str(finish_reason_val.name if hasattr(finish_reason_val, 'name') else finish_reason_val)
 
-                    if finish_reason_str == "SAFETY": # or candidate.finish_reason == genai_types.FinishReason.SAFETY:
-                         # ... safety message handling ...
-                         logging.warning(f"Content generation stopped due to safety. Finish reason: {finish_reason_str}.")
-                         return f"# Error: Content generation stopped by API due to safety. Finish Reason: {finish_reason_str}."
+                    if finish_reason_str == "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}."
@@ -434,7 +439,7 @@ class EmployerBrandingAgent:
         self.rag_system = AdvancedRAGSystem(rag_documents_df, embedding_model_name)
         self.all_dataframes = all_dataframes if all_dataframes else {}
         self.schemas_representation = get_all_schemas_representation(self.all_dataframes)
-        self.chat_history = []
+        self.chat_history = [] # This will store entries like {"role": "user/assistant", "content": "text"}
         logging.info("EmployerBrandingAgent Initialized (using Client API approach).")
 
     def _build_prompt(self, user_query: str, role="Employer Branding Analyst & Strategist", task_decomposition_hint=None, cot_hint=True) -> str:
@@ -480,12 +485,21 @@ class EmployerBrandingAgent:
         return prompt
 
     async def process_query(self, user_query: str, role="Employer Branding Analyst & Strategist", task_decomposition_hint=None, cot_hint=True) -> str:
-        current_turn_history_for_llm = self.chat_history[:]
-        self.chat_history.append({"role": "user", "parts": [{"text": user_query}]})
+        current_turn_history_for_llm = self.chat_history[:] # History before this turn
+        
+        # Append new user message to chat_history using 'content' key
+        self.chat_history.append({"role": "user", "content": user_query})
+        
         full_prompt = self._build_prompt(user_query, role, task_decomposition_hint, cot_hint)
         logging.info(f"Built prompt for query: {user_query[:100]}...")
+        
+        # Pass history (which now uses 'content') to pandas_llm.query
         response_text = await self.pandas_llm.query(full_prompt, self.all_dataframes, history=current_turn_history_for_llm)
-        self.chat_history.append({"role": "model", "parts": [{"text": response_text}]})
+        
+        # Append new assistant message to chat_history using 'content' key
+        # Standardize role to 'assistant' for model responses in history
+        self.chat_history.append({"role": "assistant", "content": response_text})
+        
         MAX_HISTORY_TURNS = 5
         if len(self.chat_history) > MAX_HISTORY_TURNS * 2:
             self.chat_history = self.chat_history[-(MAX_HISTORY_TURNS * 2):]