The Agent Gateway is an agentic orchestration framework that offers native support for Snowflake tools.
Instead of requiring users or developers to choose between RAG with Cortex Search or Text2SQL with Cortex Analyst, the Agent Gateway orchestrates the request to the appropriate tool.
The Agent Gateway can be configured to work with 3 types of tools:
- Cortex Search Tool: For unstructured data analysis, which requires a standard RAG access pattern.
- Cortex Analyst Tool: For supporting structured data analysis, which requires a Text2SQL access pattern.
- Python Tool: For supporting custom user operations (i.e. sending API requests to third party services), which requires calling arbitrary python.
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 Agent Gateway. For a walkthrough of how to configure and run a system with all 3 types of tools, see the Quickstart notebook.
This library is optimized for client-side orchestration. If you prefer a managed service, we recommend the Snowflake Chat API.
In a new virtual environment with Python 3.10 or 3.11, install the latest version of this framework.
pip install orchestration-framework@git+https://github.com/Snowflake-Labs/orchestration-framework.git
Note For Mac Users: Mac users have reported SSL Certificate issues when using the Cortex REST API. This is related to python virtual environments not having access to local certificates. One potential solution to avoid SSL Certificate issues is to use Finder to locate the "Install Certificates.command" file in your relevant Python directory and run that file before initializing the agent. See this thread for more info.
Agents require the underlying Cortex Search, Cortex Analyst, or Python tools to be configured by the user.
To follow the Quickstart notebook in this repo, you can generate the Cortex Search and Cortex Analyst demo services as follows:
from agent_gateway.tools.utils import generate_demo_services
from snowflake.snowpark import Session
session = Session.builder.create()
generate_demo_services(session)
Tools must be configured with relevant metadata for the Agent Gateway to route requests to the appropriate service.
NOTE: For best results, use specific and mutually exclusive language in your metadata descriptions to make it easy for the agent to delegate work to the right tools.
from agent_gateway.tools import CortexSearchTool, CortexAnalystTool, PythonTool
# Cortex Search Service Config
search_config = {
"service_name": "SEC_SEARCH_SERVICE",
"service_topic": "Snowflake's business,product offerings,and performance",
"data_description": "Snowflake annual reports",
"retrieval_columns": ["CHUNK"],
"snowflake_connection": session,
}
annual_reports = CortexSearchTool(**search_config)
# Cortex Analyst Config
analyst_config = {
"semantic_model": "sp500_semantic_model.yaml",
"stage": "ANALYST",
"service_topic": "S&P500 company and stock metrics",
"data_description": "a table with stock and financial metrics about S&P500 companies ",
"snowflake_connection": session,
}
sp500 = CortexAnalystTool(**analyst_config)
def get_html(url):
response = requests.get(url)
return response.text
python_scraper_config = {
"tool_description": "reads the html from a given URL or website",
"output_description": "html of a webpage",
"python_func": get_html
}
web_crawler = PythonTool(**python_scraper_config)
from agent_gateway import Agent
# Config + Initialize Agent
snowflake_tools = [annual_reports, sp500, web_crawler]
snowflake_agent = Agent(snowflake_connection=session, tools=snowflake_tools)
# Structured Data Question (Text2SQL)
answer = snowflake_agent("What is market cap of company X?")
print(answer)
# Unstructured Data Question (RAG)
answer = snowflake_agent("What are the strategic plans for company X")
print(answer)
# Web Search Question (Python Tool)
answer = snowflake_agent(
"Summarize this article: http://localhost:8080/dummyproductannouncements/interview.html"
)
print(answer)
- This library is optimized for client-side orchestration. If you prefer a managed service that does the orchestration inside of Snowflake, we recommend using the Snowflake Chat API.
- Yes, the Agent Gateway can run in SPCS and Snowflake notebooks backed by a container runtime. To install the library directly from GitHub, you must enable a network rule with an external access integration. Here is an example configuration:
CREATE NETWORK RULE agent_network_rule
MODE = EGRESS
TYPE = HOST_PORT
VALUE_LIST = ('github.com');
CREATE EXTERNAL ACCESS INTEGRATION agent_network_int
ALLOWED_NETWORK_RULES = (agent_network_rule)
ENABLED = true;
- Yes, see the demo app for an example Streamlit app that uses the Agent Gateway for orchestration across Cortex Search, Cortex Analyst, and Python tools. Note, running the gateway is not yet supported in Stremlit in Snowflake.
- The Agent Gateway and its tools take an authenticated snowpark connection. Just create your session object with your standard connection parameters.
If I have multiple Cortex Search Services, can I use multiple Cortex Search tools with this framework?
- Yes, you can connect multiple tools of the same type to the Agent Gateway.
search_one = CortexSearchTool(**search_one_config)
search_two = CortexSearchTool(**search_two_config)
snowflake_agent = Agent(snowflake_connection=session, tools=[search_one, search_two])
- Yes. The Cortex Analyst and Cortex Search tools take in a snowpark session as an input. This allows users to use different sessions / accounts in the same gateway agent.
- The Agent Gateway logger is set to INFO level by default. This allows users to view which tools are being used to answer the user's question. For more detailed logging and visibility into intermediary results of the tool calls, set the LOGGING_LEVEL=DEBUG.
- Tools are implemented asynchronously. To validate your configuration, you can run each tool in isolation as follows:
tool_result = await my_cortex_search_tool("This is a sample cortex search question")
- This framework utilizes a dedicated planner LLM to generate a sequence of tool calls that can be executed in parallel. While the orchestration is done on the client-side, Snowflake compute is leveraged for plan generation and tooling execution. We leverage the LLM Compiler architecture from Berekeley AI Research Lab. Kim, S., Moon, S., Tabrizi, R., Lee, N., Mahoney, M. W., Keutzer, K., and Gholami, A. An LLM Compiler for Parallel Function Calling, 2024.
- You can add issues to the GitHub or email Alejandro Herrera ([email protected])