Update agent/planning_agent.py

agent/planning_agent.py

@@ -24,6 +24,7 @@ llm = None
 chat_memory = None
 query_memory = None
 agent = None
+clearml_callback = None
 
 def initialize_planning_agent(llm_instance, chat_memory_instance, query_memory_instance, clearml_instance):
     global llm, chat_memory, query_memory, agent, clearml_callback
@@ -145,7 +146,7 @@ def summarize_chat_history():
     except Exception as e:
         logger.error(f"Error summarizing chat history: {str(e)}")
 
-def route_query(query):
+def route_query(query, callback):
     # Summarize chat history before routing
     summarize_chat_history()
     # Get original query from memory if needed
@@ -185,10 +186,11 @@ def compose_response(response):
 def execute(query):
     try:
         # Store original query
-        query_memory.memories['original_query'] = query
-        return agent.run(
-            f"Process this user query: {query}"
-        )
+        query_memory.memories['original_query'] = query 
+        response = agent.run(
+            f"Process this user query: {query}")
+        clearml_callback.flush_tracker(langchain_asset=planning_agent, name="Planning agent", finish=True)
+        return response
     except Exception as e:
         logger.error(f"Error in planning agent: {str(e)}")
         return f"Error in planning agent: {str(e)}"