Cortex Agents

USE ROLE ACCOUNTADMIN;

CREATE ROLE IF NOT EXISTS nimble_role;
GRANT DATABASE ROLE SNOWFLAKE.CORTEX_USER TO ROLE nimble_role;

CREATE DATABASE IF NOT EXISTS NIMBLE_INTEGRATION;
CREATE SCHEMA   IF NOT EXISTS NIMBLE_INTEGRATION.TOOLS;
GRANT USAGE ON DATABASE NIMBLE_INTEGRATION       TO ROLE nimble_role;
GRANT USAGE ON SCHEMA   NIMBLE_INTEGRATION.TOOLS TO ROLE nimble_role;

CREATE WAREHOUSE IF NOT EXISTS NIMBLE_AGENT_WH
  WAREHOUSE_SIZE = XSMALL
  AUTO_SUSPEND   = 60
  AUTO_RESUME    = TRUE;
GRANT USAGE ON WAREHOUSE NIMBLE_AGENT_WH TO ROLE nimble_role;

CREATE OR REPLACE NETWORK RULE nimble_api_rule
  MODE       = EGRESS
  TYPE       = HOST_PORT
  VALUE_LIST = ('sdk.nimbleway.com:443');

CREATE OR REPLACE SECRET nimble_api_key
  TYPE          = GENERIC_STRING
  SECRET_STRING = 'YOUR_NIMBLE_API_KEY';

CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION nimble_eai
  ALLOWED_NETWORK_RULES          = (nimble_api_rule)
  ALLOWED_AUTHENTICATION_SECRETS = (nimble_api_key)
  ENABLED                        = TRUE;

GRANT USAGE ON INTEGRATION nimble_eai TO ROLE nimble_role;

CREATE OR REPLACE FUNCTION NIMBLE_INTEGRATION.TOOLS.NIMBLE_SEARCH(
  query        STRING,
  max_results  INTEGER DEFAULT 10
)
RETURNS VARIANT
LANGUAGE PYTHON
RUNTIME_VERSION = 3.11
PACKAGES        = ('requests')
HANDLER         = 'run'
EXTERNAL_ACCESS_INTEGRATIONS = (nimble_eai)
SECRETS         = ('api_key' = nimble_api_key)
AS $$
import _snowflake, requests

def run(query, max_results):
    key = _snowflake.get_generic_secret_string('api_key')
    r = requests.post(
        'https://sdk.nimbleway.com/v2/search',
        headers={'Authorization': f'Bearer {key}',
                 'Content-Type':  'application/json'},
        json={'query': query, 'max_results': max_results},
        timeout=60,
    )
    r.raise_for_status()
    return r.json()
$$;

GRANT USAGE ON FUNCTION NIMBLE_INTEGRATION.TOOLS.NIMBLE_SEARCH(STRING, INTEGER)
  TO ROLE nimble_role;

CREATE OR REPLACE FUNCTION NIMBLE_INTEGRATION.TOOLS.NIMBLE_EXTRACT(
  url     STRING,
  render  BOOLEAN DEFAULT TRUE
)
RETURNS VARIANT
LANGUAGE PYTHON
RUNTIME_VERSION = 3.11
PACKAGES        = ('requests')
HANDLER         = 'run'
EXTERNAL_ACCESS_INTEGRATIONS = (nimble_eai)
SECRETS         = ('api_key' = nimble_api_key)
AS $$
import _snowflake, requests

def run(url, render):
    key = _snowflake.get_generic_secret_string('api_key')
    r = requests.post(
        'https://sdk.nimbleway.com/v2/extract',
        headers={'Authorization': f'Bearer {key}',
                 'Content-Type':  'application/json'},
        json={'url': url, 'render': render, 'formats': ['markdown']},
        timeout=90,
    )
    r.raise_for_status()
    return r.json()
$$;

GRANT USAGE ON FUNCTION NIMBLE_INTEGRATION.TOOLS.NIMBLE_EXTRACT(STRING, BOOLEAN)
  TO ROLE nimble_role;

CREATE OR REPLACE AGENT nimble_web_research_agent
  COMMENT = 'Agent with Nimble web search and content retrieval'
  FROM SPECIFICATION
  $$
  models:
    orchestration: auto

  orchestration:
    budget:
      seconds: 60
      tokens:  32000

  instructions:
    system: |
      You are a research assistant with access to Nimble's web data toolkit.
      Always cite source URLs.
    orchestration: |
      - Use NIMBLE_SEARCH for queries requiring web search.
      - Use NIMBLE_EXTRACT when the user provides specific URLs.

  tools:
    - tool_spec:
        type: generic
        name: nimble_search
        description: Search the web using Nimble.
        input_schema:
          type: object
          properties:
            query:       { type: string,  description: The search query. }
            max_results: { type: integer, description: "Number of results (default 10)." }
          required: [query]

    - tool_spec:
        type: generic
        name: nimble_extract
        description: Retrieve parsed content for a single URL.
        input_schema:
          type: object
          properties:
            url:    { type: string,  description: The URL to extract. }
            render: { type: boolean, description: "JS-render the page (slower; needed for SPAs)." }
          required: [url]

  tool_resources:
    nimble_search:
      type: function
      execution_environment:
        type:      warehouse
        warehouse: NIMBLE_AGENT_WH
      identifier: NIMBLE_INTEGRATION.TOOLS.NIMBLE_SEARCH
    nimble_extract:
      type: function
      execution_environment:
        type:      warehouse
        warehouse: NIMBLE_AGENT_WH
      identifier: NIMBLE_INTEGRATION.TOOLS.NIMBLE_EXTRACT
  $$;