Update agent/planning_agent.py

agent/planning_agent.py

@@ -1,7 +1,10 @@
 from langchain_openai import ChatOpenAI
 from langchain.agents import initialize_agent, AgentType
+from langchain.agents import AgentType, initialize_agent, load_tools
 from langchain.tools import Tool
 from langchain.memory import ConversationBufferMemory, SimpleMemory
+from langchain_community.callbacks import ClearMLCallbackHandler
+from langchain_core.callbacks import StdOutCallbackHandler
 import agent.router_agent as router_agent
 import agent.product_review_agent as product_review_agent
 import agent.generic_agent as generic_agent
@@ -22,14 +25,26 @@ chat_memory = None
 query_memory = None
 agent = None
 
-def initialize_planning_agent(llm_instance, chat_memory_instance, query_memory_instance, callback_fn):
-    global llm, chat_memory, query_memory, agent, callbacks
+def initialize_planning_agent(llm_instance, chat_memory_instance, query_memory_instance):
+    global llm, chat_memory, query_memory, agent
     
     llm = llm_instance
     chat_memory = chat_memory_instance
     query_memory = query_memory_instance
-    callbacks = callback_fn
+
+    # 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,)
     
+    callbacks = [StdOutCallbackHandler(), clearml_callback]
+
     # Initialize agents
     router_agent.initialize_router_agent(llm, chat_memory)
     product_review_agent.initialize_product_review_agent(llm, chat_memory)