Update app.py

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 from langchain_community.chat_models import ChatOpenAI
 from langchain.memory import ConversationBufferMemory, SimpleMemory
 from langchain.agents import initialize_agent, AgentType
+from langchain_community.callbacks import ClearMLCallbackHandler
+from langchain_core.callbacks import StdOutCallbackHandler
 from dotenv import load_dotenv
 from dotenv import load_dotenv, find_dotenv
 import os
@@ -22,7 +24,24 @@ def initialize_components():
     load_dotenv()
     
     OPENAI_API_KEY="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"
+    CLEARML_API_ACCESS_KEY="NYZZ07E2ZEW08V4DUGYY2PA7O6JX5F"
+    CLEARML_API_SECRET_KEY="MkfQrIOuKNFRWHfCz32cN-UVm_19M7_vgxAwRn8twnvHYJ1xeqD9T2GZcIX9RwnD8mw"
+    SERPAPI_API_KEY="619f2302253fbe56448bcf82565caf2a3263d845944682533f10b09a0d1650e6"
+    
     # openai.api_key=OPENAI_API_KEY
+
+    # 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]
     
     llm = ChatOpenAI(
         model_name="gpt-3.5-turbo",
@@ -38,7 +57,7 @@ def initialize_components():
     query_memory = SimpleMemory()
 
     # Initialize planning agent with both memories
-    planning_agent.initialize_planning_agent(llm, chat_memory, query_memory)
+    planning_agent.initialize_planning_agent(llm, chat_memory, query_memory, callback)
 
     logger.info("Components initialized successfully")
 
@@ -56,6 +75,7 @@ def process_query(query, history):
         
         # Execute query through planning agent
         response = planning_agent.execute(query)
+        clearml_callback.flush_tracker(langchain_asset=planning_agent, name="Planning agent", finish=True)
         
         # Add current interaction to chat memory
         if chat_memory and hasattr(chat_memory, 'chat_memory'):