feat: sql tool support

Quickstart.ipynb

@@ -51,7 +51,7 @@
    "outputs": [],
    "source": [
     "from agent_gateway import Agent\n",
-    "from agent_gateway.tools import CortexSearchTool, CortexAnalystTool, PythonTool\n",
+    "from agent_gateway.tools import CortexSearchTool, CortexAnalystTool, PythonTool, SQLTool\n",
     "from snowflake.snowpark import Session\n",
     "import os\n",
     "from dotenv import load_dotenv\n",
@@ -87,7 +87,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -106,6 +106,7 @@
     "    \"service_topic\": \"S&P500 company and stock metrics\",\n",
     "    \"data_description\": \"a table with stock and financial metrics about S&P500 companies \",\n",
     "    \"snowflake_connection\": snowpark,\n",
+    "    \"max_results\": 5,\n",
     "}"
    ]
   },
@@ -127,7 +128,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -146,6 +147,81 @@
     "}"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# SQL Tool"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The SQL Tool allows users to pre-define sql metrics and custom pipelines that the agent will be able to utilize to answer specialized questions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sql_query = \"\"\"WITH CompanyMetrics AS (\n",
+    "    SELECT \n",
+    "        LONGNAME,\n",
+    "        SECTOR,\n",
+    "        INDUSTRY,\n",
+    "        CURRENTPRICE,\n",
+    "        MARKETCAP,\n",
+    "        EBITDA,\n",
+    "        CASE \n",
+    "            WHEN MARKETCAP > 0 AND EBITDA IS NOT NULL THEN (EBITDA * 100.0) / MARKETCAP\n",
+    "            ELSE NULL\n",
+    "        END AS EBITDA_Margin\n",
+    "    FROM CUBE_TESTING.PUBLIC.SP500\n",
+    "),\n",
+    "AverageMetrics AS (\n",
+    "    SELECT \n",
+    "        AVG(EBITDA_Margin) AS Average_EBITDA_Margin\n",
+    "    FROM CompanyMetrics\n",
+    "),\n",
+    "NormalizedMetrics AS (\n",
+    "    SELECT \n",
+    "        cm.LONGNAME,\n",
+    "        cm.SECTOR,\n",
+    "        cm.INDUSTRY,\n",
+    "        cm.CURRENTPRICE,\n",
+    "        cm.MARKETCAP,\n",
+    "        cm.EBITDA,\n",
+    "        cm.EBITDA_Margin,\n",
+    "        CASE \n",
+    "            WHEN am.Average_EBITDA_Margin > 0 THEN cm.EBITDA_Margin / am.Average_EBITDA_Margin\n",
+    "            ELSE NULL\n",
+    "        END AS Normalized_EBITDA_Margin\n",
+    "    FROM CompanyMetrics cm\n",
+    "    CROSS JOIN AverageMetrics am\n",
+    ")\n",
+    "SELECT \n",
+    "    LONGNAME,\n",
+    "    SECTOR,\n",
+    "    INDUSTRY,\n",
+    "    CURRENTPRICE,\n",
+    "    MARKETCAP,\n",
+    "    EBITDA,\n",
+    "    EBITDA_Margin,\n",
+    "    Normalized_EBITDA_Margin\n",
+    "FROM NormalizedMetrics;\"\"\"\n",
+    "\n",
+    "sql_tool_config = {\n",
+    "    \"name\": \"margin_eval\",\n",
+    "    \"connection\": snowpark,\n",
+    "    \"sql_query\": sql_query,\n",
+    "    \"tool_description\": \"Calculates the normalized EBITDA Margin as a % relative to the SP500 average\",\n",
+    "    \"output_description\": \"EBITDA Margin %\",\n",
+    "}"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -162,26 +238,29 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "2025-03-13 17:24:15,091 - AgentGatewayLogger - INFO - Cortex Search Tool successfully initialized\n",
-      "2025-03-13 17:24:15,092 - AgentGatewayLogger - INFO - Cortex Analyst Tool successfully initialized\n",
-      "2025-03-13 17:24:15,093 - AgentGatewayLogger - INFO - Python Tool successfully initialized\n",
-      "2025-03-13 17:24:15,094 - AgentGatewayLogger - INFO - Cortex gateway successfully initialized\n"
+      "2025-03-25 13:48:11,284 - AgentGatewayLogger - INFO - Cortex Search Tool successfully initialized\n",
+      "2025-03-25 13:48:11,541 - AgentGatewayLogger - INFO - Cortex Analyst Tool successfully initialized\n",
+      "2025-03-25 13:48:11,544 - AgentGatewayLogger - INFO - Python Tool successfully initialized\n",
+      "2025-03-25 13:48:11,546 - AgentGatewayLogger - INFO - SQL Tool successfully initialized\n",
+      "2025-03-25 13:48:12,057 - AgentGatewayLogger - INFO - Cortex gateway successfully initialized\n"
      ]
     }
    ],
    "source": [
     "annual_reports = CortexSearchTool(**search_config)\n",
     "sp500 = CortexAnalystTool(**analyst_config)\n",
     "web_crawler = PythonTool(**python_crawler_config)\n",
+    "margin_eval = SQLTool(**sql_tool_config)\n",
     "\n",
-    "snowflake_tools = [annual_reports, sp500, web_crawler]\n",
+    "\n",
+    "snowflake_tools = [annual_reports, sp500, web_crawler, margin_eval]\n",
     "agent = Agent(snowflake_connection=snowpark, tools=snowflake_tools, max_retries=3)"
    ]
   },
@@ -240,6 +319,36 @@
     "agent(\"What is the market cap of Apple?\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2025-03-25 10:53:27,340 - AgentGatewayLogger - INFO - running margin_eval task\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'output': 'The normalized EBITDA margin for Microsoft Corporation is approximately 0.3999 or 39.99%.',\n",
+       " 'sources': [{'tool_type': 'SQL',\n",
+       "   'tool_name': 'margin_eval',\n",
+       "   'metadata': [None]}]}"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "agent(\"What is MSFT's normalized EBITDA margin?\")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 6,
@@ -393,6 +502,62 @@
     "    \"What is the market cap of each of the cloud providers mentioned in Snowflake's annual report?\"\n",
     ")"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Agent Observability"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "install framework with requisite dependencies with `pip install orchestration-framework[trulens]` and initialize as follows:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from agent_gateway import TruAgent\n",
+    "from trulens.connectors.snowflake import SnowflakeConnector\n",
+    "\n",
+    "tru_conn = SnowflakeConnector(**connection_parameters)\n",
+    "\n",
+    "agent = TruAgent(\n",
+    "    app_name=\"observable\",\n",
+    "    app_version=\"v0\",\n",
+    "    trulens_snowflake_connection=tru_conn,\n",
+    "    snowflake_connection=snowpark,\n",
+    "    tools=snowflake_tools,\n",
+    "    max_retries=3,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Run the dashboard to view traces"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trulens.core import TruSession\n",
+    "from trulens.dashboard import run_dashboard\n",
+    "\n",
+    "session = TruSession(connector=tru_conn)\n",
+    "\n",
+    "run_dashboard(session, port=8084)"
+   ]
   }
  ],
  "metadata": {

README.md

@@ -14,6 +14,7 @@ access pattern.
 Text2SQL access pattern.
 - **Python Tool**: For supporting custom user operations (i.e. sending API requests to
 third party services), which requires calling arbitrary python.
+- **SQL Tool**: For supporting custom SQL pipelines built by users.
 
 The Agent Gateway supports multi-step and multi-tool workflows. Users have the
 flexibility to create multiple Cortex Search and Cortex Analyst tools for use with the
@@ -65,7 +66,7 @@ tools.
 ##### Cortex Search Tool Configuration
 
 ```python
-from agent_gateway.tools import CortexSearchTool, CortexAnalystTool, PythonTool
+from agent_gateway.tools import CortexSearchTool, CortexAnalystTool, PythonTool, SQLTool
 
 # Cortex Search Service Config
 search_config = {
@@ -109,6 +110,20 @@ python_scraper_config = {
 web_crawler = PythonTool(**python_scraper_config)
 ```
 
+##### SQL Tool Configuration
+
+```python
+sql_query = '''SELECT * FROM MY_EVENTS_TABLE '''
+sql_tool_config = {
+    "tool_description": "analyzes custom user metrics",
+    "output_description": "key user metrics",
+    "sql": sql_query,
+    "connection":session
+    }
+
+custom_metrics = SQLTool(**sql_tool_config)
+```
+
 ## Agent Configuration + Usage
 
 ````python