diff --git a/plot_agent/agno_plot_agent/agent.py b/plot_agent/agno_plot_agent/agent.py index 11a0b80..1da420a 100644 --- a/plot_agent/agno_plot_agent/agent.py +++ b/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,7 +67,7 @@ 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()], @@ -76,11 +75,11 @@ class SQLAndPlotWorkflow(Workflow): 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]: diff --git a/plot_agent/agno_plot_agent/api.py b/plot_agent/agno_plot_agent/api.py index bc2d486..b23002d 100644 --- a/plot_agent/agno_plot_agent/api.py +++ b/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") diff --git a/plot_agent/agno_plot_agent/clickhouse_tools.py b/plot_agent/agno_plot_agent/clickhouse_tools.py index 18a04cc..5e1d21c 100644 --- a/plot_agent/agno_plot_agent/clickhouse_tools.py +++ b/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") diff --git a/plot_agent/agno_plot_agent/plot_tools.py b/plot_agent/agno_plot_agent/plot_tools.py index ce1035c..0291fe4 100644 --- a/plot_agent/agno_plot_agent/plot_tools.py +++ b/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")