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 PROCEDURE NIMBLE_INTEGRATION.TOOLS.NIMBLE_SEARCH(
  query        STRING,
  num_results  INTEGER DEFAULT 10
)
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION = 3.11
PACKAGES        = ('requests')
HANDLER         = 'run'
EXTERNAL_ACCESS_INTEGRATIONS = (nimble_eai)
SECRETS         = ('api_key' = nimble_api_key)
AS $$
import json, _snowflake, requests

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

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

CREATE OR REPLACE PROCEDURE NIMBLE_INTEGRATION.TOOLS.NIMBLE_EXTRACT(urls STRING)
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION = 3.11
PACKAGES        = ('requests')
HANDLER         = 'run'
EXTERNAL_ACCESS_INTEGRATIONS = (nimble_eai)
SECRETS         = ('api_key' = nimble_api_key)
AS $$
import json, _snowflake, requests

def run(session, urls):
    key = _snowflake.get_generic_secret_string('api_key')
    out = []
    for u in [x.strip() for x in urls.split(',') if x.strip()]:
        r = requests.post(
            'https://sdk.nimbleway.com/v1/extract',
            headers={'Authorization': f'Bearer {key}',
                     'Content-Type':  'application/json'},
            json={'url': u, 'parse': True},
            timeout=60,
        )
        r.raise_for_status()
        out.append(r.json())
    return json.dumps(out)
$$;

GRANT USAGE ON PROCEDURE NIMBLE_INTEGRATION.TOOLS.NIMBLE_EXTRACT(STRING)
  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. }
            num_results: { type: integer, description: "Number of results (default 10)." }
          required: [query]

    - tool_spec:
        type: generic
        name: nimble_extract
        description: Retrieve parsed content for specific URLs.
        input_schema:
          type: object
          properties:
            urls: { type: string, description: Comma-separated URLs. }
          required: [urls]

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