Enhance tracing on plot_agents

plot_agent/agno_plot_agent/agent.py

@@ -10,7 +10,7 @@
 from agno.utils.log import logger
 from dotenv import load_dotenv
 from traceloop.sdk import Traceloop
-from traceloop.sdk.decorators import workflow, agent, task
+from traceloop.sdk.decorators import workflow, task
 
 load_dotenv()
 
@@ -37,10 +37,9 @@ class VisualizationRequest(BaseModel):
     title: Optional[str] = Field(None, description="Plot title.")
     hue: Optional[str] = Field(None, description="Column for color grouping.")
 
-@agent(name="sql_and_plot_workflow")
-class SQLAndPlotWorkflow(Workflow):        
-    # SQL Agent that generates and executes Clickhouse queries
-    sql_agent: Agent = Agent(
+def create_sql_agent() -> Agent:
+    """Creates and returns a SQL agent for generating and executing Clickhouse queries."""
+    return Agent(
         model=OpenAIChat(id="openai-main/gpt-4o", api_key=os.getenv("LLM_GATEWAY_API_KEY"), base_url=os.getenv("LLM_GATEWAY_BASE_URL")),
         description="You are an expert in generating and executing Clickhouse SQL queries from user queries in English.",
         instructions=[
@@ -68,19 +67,19 @@ class SQLAndPlotWorkflow(Workflow):
             "- applied_configs: Map(LowCardinality(String), Map(LowCardinality(String), String)): Configuration settings applied to the request.",
             "- created_at: DateTime64(9) Delta(8), ZSTD(1): The timestamp when the request was made.",
             "Clickhouse has slighlty different syntax rules than MySQL or PostgreSQL. Please make sure to use the correct syntax for Clickhouse."
-            "Syntax rule: Use toIntervalXXX(N) (e.g., toIntervalDay(30)) instead of INTERVAL N UNIT (e.g., INTERVAL 30 DAY) for interval arithmetic in ClickHouse."
+            "Syntax rule: Use toIntervalXXX(N) (e.g., toIntervalDay(30)) instead of INTERVAL N UNIT (e.g., INTERVAL 30 DAY) for interval arithmetic in ClickHouse.",
             "Syntax rule: Do not end in a semicolon (;) in the query. Only end with a newline.",
         ],
         tools=[ClickHouseTools()],
         show_tool_calls=True,
         markdown=True,
         response_model=SQLQueryResult,
         structured_outputs=True,
-        # debug_mode=True,
     )
 
-    # Plot Agent that creates visualizations
-    plot_agent: Agent = Agent(
+def create_plot_agent() -> Agent:
+    """Creates and returns a Plot agent for creating data visualizations."""
+    return Agent(
         model=OpenAIChat(id="openai-main/gpt-4o", api_key=os.getenv("LLM_GATEWAY_API_KEY"), base_url=os.getenv("LLM_GATEWAY_BASE_URL")),
         description="You are an expert in creating data visualizations from SQL query results.",
         instructions=[
@@ -103,10 +102,13 @@ class SQLAndPlotWorkflow(Workflow):
         show_tool_calls=True,
         markdown=True,
         response_model=VisualizationRequest,
-        # structured_outputs=True,
-        # debug_mode=True,
     )
 
+class SQLAndPlotWorkflow(Workflow):
+    def __init__(self):
+        super().__init__()
+        self.sql_agent = create_sql_agent()
+        self.plot_agent = create_plot_agent()
 
     @workflow(name="plotting workflow")
     def run_workflow(self, query: str) -> Iterator[RunResponse]:

plot_agent/agno_plot_agent/api.py

@@ -7,6 +7,7 @@
 import uuid
 from agent import SQLAndPlotWorkflow, PlotResult
 from agno.utils.log import logger
+from traceloop.sdk.decorators import agent
 
 # Create plots directory if it doesn't exist
 PLOTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "plots")

plot_agent/agno_plot_agent/clickhouse_tools.py

@@ -5,7 +5,9 @@
 from agno.utils.log import logger
 import os
 from dotenv import load_dotenv
+from traceloop.sdk.decorators import tool
 
+@tool(name="clickhouse_tools")
 class ClickHouseTools(Toolkit):
     def __init__(self):
         super().__init__(name="clickhouse_tools")

plot_agent/agno_plot_agent/plot_tools.py

@@ -9,11 +9,13 @@
 import os
 import uuid
 from matplotlib.ticker import FuncFormatter
+from traceloop.sdk.decorators import tool
 
 # Create plots directory if it doesn't exist
 PLOTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "plots")
 os.makedirs(PLOTS_DIR, exist_ok=True)
 
+@tool(name="plot_tools")
 class PlotTools(Toolkit):
     def __init__(self):
         super().__init__(name="plot_tools")