Update agent/planning_agent.py

agent/planning_agent.py

@@ -25,26 +25,14 @@ chat_memory = None
 query_memory = None
 agent = None
 
-def initialize_planning_agent(llm_instance, chat_memory_instance, query_memory_instance):
-    global llm, chat_memory, query_memory, agent
+def initialize_planning_agent(llm_instance, chat_memory_instance, query_memory_instance, clearml_instance):
+    global llm, chat_memory, query_memory, agent, clearml_callback
     
     llm = llm_instance
     chat_memory = chat_memory_instance
     query_memory = query_memory_instance
-
-    # Setup and use the ClearML Callback
-    clearml_callback = ClearMLCallbackHandler(
-    task_type="inference",
-    project_name="langchain_callback_demo",
-    task_name="llm",
-    tags=["test"],
-    # Change the following parameters based on the amount of detail you want tracked
-    visualize=True,
-    complexity_metrics=True,
-    stream_logs=True,)
+    clearml_callback = clearml_instance
     
-    callbacks = [StdOutCallbackHandler(), clearml_callback]
-
     # Initialize agents
     router_agent.initialize_router_agent(llm, chat_memory)
     product_review_agent.initialize_product_review_agent(llm, chat_memory)
@@ -119,7 +107,7 @@ def initialize_planning_agent(llm_instance, chat_memory_instance, query_memory_i
         tools,
         llm,
         agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-        callbacks=callbacks,
+        callbacks=clearml_callback,
         verbose=True,
         memory=chat_memory,
         system_message=system_prompt,