Update eb_agent_module.py

eb_agent_module.py

@@ -433,15 +433,23 @@ class PandasLLM:
 
 # --- Employer Branding Agent ---
 class EmployerBrandingAgent:
-    def __init__(self, llm_model_name: str, gc_dict: dict, ss_list: list, all_dfs: dict, rag_df: pd.DataFrame, emb_m_name: str, dp=True, fs=True):
-        self.pandas_llm = PandasLLM(llm_model_name, gc_dict, ss_list, dp, fs)
-        self.rag_system = AdvancedRAGSystem(rag_df, emb_m_name)
-        self.all_dataframes = all_dfs if all_dfs else {}
+    def __init__(self, llm_model_name: str, 
+                 generation_config_dict: dict, # Changed from gc_dict
+                 safety_settings_list: list,   # Changed from ss_list
+                 all_dataframes: dict,         # Changed from all_dfs
+                 rag_documents_df: pd.DataFrame, # Changed from rag_df
+                 embedding_model_name: str,    # Changed from emb_m_name
+                 data_privacy=True,           # Changed from dp
+                 force_sandbox=True):         # Changed from fs
+        
+        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}).")
 
-    def _build_prompt(self, user_query: str, role="EB Analyst", task_hint=None, cot=True) -> str:
+    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}'. Goal: insights from DataFrames & RAG.\n"
         if self.pandas_llm.data_privacy: prompt += "PRIVACY: Summarize/aggregate PII.\n"
         if self.pandas_llm.force_sandbox:
@@ -461,7 +469,7 @@ class EmployerBrandingAgent:
             else: prompt += "--- TEXT RESPONSE THOUGHT PROCESS ---\n1.Goal? 2.Data? 3.Insights (DFs+RAG). 4.Structure response.\n"
         return prompt
 
-    async def process_query(self, user_query: str, role="EB Analyst", task_hint=None, cot=True) -> str:
+    async def process_query(self, user_query: str, role="EB Analyst", task_hint=None, cot=True) -> str: # Simplified role
         hist_for_llm = self.chat_history[:]
         self.chat_history.append({"role": "user", "content": user_query})
         prompt = self._build_prompt(user_query, role, task_hint, cot)
@@ -471,13 +479,29 @@ class EmployerBrandingAgent:
         if len(self.chat_history) > 10: self.chat_history = self.chat_history[-10:]; logging.info("Chat history truncated.")
         return response
 
-    def update_dataframes(self, new_dfs: dict): self.all_dataframes = new_dfs if new_dfs else {}; self.schemas_representation = get_all_schemas_representation(self.all_dataframes); logging.info("Agent DFs updated.")
-    def clear_chat_history(self): self.chat_history = []; logging.info("Agent chat history cleared.")
+    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("Agent DFs updated.")
+    
+    def clear_chat_history(self): 
+        self.chat_history = []
+        logging.info("Agent chat history cleared.")
 
 # --- Example Usage (Conceptual) ---
 async def main_test():
     logging.info(f"Test (Real GenAI: {_REAL_GENAI_LOADED}, API Key: {bool(GEMINI_API_KEY)})")
-    agent = EmployerBrandingAgent(LLM_MODEL_NAME, GENERATION_CONFIG_PARAMS, DEFAULT_SAFETY_SETTINGS, {}, df_rag_documents, GEMINI_EMBEDDING_MODEL_NAME)
+    # Example: Provide all parameters to EmployerBrandingAgent constructor
+    agent = EmployerBrandingAgent(
+        llm_model_name=LLM_MODEL_NAME,
+        generation_config_dict=GENERATION_CONFIG_PARAMS,
+        safety_settings_list=DEFAULT_SAFETY_SETTINGS,
+        all_dataframes={}, # Empty dict for example
+        rag_documents_df=df_rag_documents,
+        embedding_model_name=GEMINI_EMBEDDING_MODEL_NAME,
+        data_privacy=True,
+        force_sandbox=True
+    )
     for q in ["What are EB best practices?", "Hello Agent!"]:
         logging.info(f"\nQuery: {q}")
         resp = await agent.process_query(q)