diff --git a/.github/workflows/integration_tests.yml b/.github/workflows/integration_tests.yml index ee0faf6..88152a5 100644 --- a/.github/workflows/integration_tests.yml +++ b/.github/workflows/integration_tests.yml @@ -47,6 +47,14 @@ jobs: SNOWFLAKE_TEST_WAREHOUSE: ${{ secrets.SNOWFLAKE_TEST_WAREHOUSE }} SNOWFLAKE_TEST_SCHEMA: ${{ secrets.SNOWFLAKE_TEST_SCHEMA }} + - name: Print compiled and run SQL on failure + if: failure() + run: | + echo "=== Compiled cortex_agent_test.sql ===" + cat integration_tests/target/compiled/dbt_cortex_agent_integration_tests/models/cortex_agent_test.sql || echo "(file not found)" + echo "=== Run cortex_agent_test.sql ===" + cat integration_tests/target/run/dbt_cortex_agent_integration_tests/models/cortex_agent_test.sql || echo "(file not found)" + - name: Run dbt tests working-directory: integration_tests run: dbt test diff --git a/.gitignore b/.gitignore index 2cb680e..2d35c02 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ .vscode/ target/ -logs/ \ No newline at end of file +logs/ +.venv/ \ No newline at end of file diff --git a/README.md b/README.md index 7df28ae..4709496 100644 --- a/README.md +++ b/README.md @@ -37,7 +37,7 @@ Create a model file whose **body is the raw Snowflake agent YAML specification** ) }} models: - orchestration: claude-4-sonnet + orchestration: auto orchestration: budget: seconds: 30 @@ -79,6 +79,9 @@ SHOW AGENTS IN SCHEMA my_db.my_schema; | `comment` | string | No | Agent description visible in Snowflake | | `profile` | string (JSON) | No | `{"display_name": "...", "avatar": "...", "color": "..."}` | | `agent_grants` | list | No | Role names to grant `USAGE` on the agent, e.g. `['my_role']` | +| `create_feedback_table` | bool | No | Whether to create the feedback table and procedure. Defaults to `true`. Set to `false` to skip. See [Feedback Tool](#feedback-tool). | +| `feedback_schema` | string | No | Schema for the feedback table and `AGENT_SUBMIT_FEEDBACK` procedure. Accepts `'SCHEMA'` or `'DB.SCHEMA'`. Defaults to the agent's own database and schema. See [Feedback Tool](#feedback-tool). | +| `feedback_table` | string | No | Fully-qualified table name override for user feedback. Defaults to `{feedback_schema}.AGENT_FEEDBACK`. Ignored when `create_feedback_table` is `false`. See [Feedback Tool](#feedback-tool). | ## How It Works @@ -95,13 +98,101 @@ Every `dbt run` issues `CREATE OR REPLACE AGENT`, so re-runs are safe and fully ## Supported Tools in Specification -As of the current Snowflake documentation: +Built-in tool types (add to the `tools:` section of your spec body): - `cortex_analyst_text_to_sql` — text-to-SQL via a Cortex Analyst semantic model - `cortex_search` — semantic search over unstructured content +- `generic` — any Snowflake stored procedure or UDF (see [Feedback Tool](#feedback-tool) for an example) Refer to the [Snowflake CREATE AGENT docs](https://docs.snowflake.com/en/sql-reference/sql/create-agent) for the full and up-to-date YAML specification reference. +## Feedback Tool + +All agents share a single feedback table and a single `AGENT_SUBMIT_FEEDBACK` procedure — no config required. On each `dbt run` the materialization provisions (idempotently): + +1. A **feedback table** named `AGENT_FEEDBACK` in the agent's database and schema, with columns: `feedback_id`, `agent_name`, `user_name`, `rating`, `comment`, `conversation_history`, `created_at` +2. A **stored procedure** named `AGENT_SUBMIT_FEEDBACK` in the same location + +By default both land in the agent's own schema. Use `feedback_schema` to place them in a shared schema instead: + +```sql +{{ config( + materialized='cortex_agent', + feedback_schema='SHARED_SCHEMA' +) }} +``` + +Or with an explicit database: + +```sql +{{ config( + materialized='cortex_agent', + feedback_schema='MY_DB.SHARED_SCHEMA' +) }} +``` + +To override the table name independently of the schema, use `feedback_table`: + +```sql +{{ config( + materialized='cortex_agent', + feedback_schema='SHARED_SCHEMA', + feedback_table='MY_DB.SHARED_SCHEMA.AGENT_FEEDBACK' +) }} +``` + +To disable feedback provisioning entirely, set `create_feedback_table: false`: + +```sql +{{ config( + materialized='cortex_agent', + create_feedback_table=false +) }} +``` + +Add the tool entry to your spec body so the agent can call it. Update the `identifier` to match your `feedback_schema` if you set one: + +```sql +{{ config( + materialized='cortex_agent', + feedback_schema='SHARED_SCHEMA' +) }} + +tools: + - tool_spec: + type: generic + name: AGENT_SUBMIT_FEEDBACK + description: 'Records user feedback. Call when the user says "feedback" or explicitly rates or comments on a response. Always pass the last 10 conversation messages. Example: user says "feedback: 1. Output wrong, expected $100" → call with rating=1, user_comment="Output wrong, expected $100". Example: user says "feedback: 5. Completely Correct" → call with rating=5, user_comment="Completely Correct".' + input_schema: + type: object + properties: + agent_name: + type: string + description: 'Name of this agent. Always pass "{{ this.identifier | upper }}".' + rating: + type: number + description: 'The user rating from 1 (worst) to 5 (best).' + user_comment: + type: string + description: 'Optional free-text feedback.' + conversation_history: + type: string + description: 'Last 10 messages from the conversation as a JSON string, e.g. [{"role":"user","content":"..."}].' + required: [agent_name, rating, conversation_history] + +tool_resources: + AGENT_SUBMIT_FEEDBACK: + execution_environment: + query_timeout: 300 + type: warehouse + warehouse: '' + identifier: '{{ this.database }}.SHARED_SCHEMA.AGENT_SUBMIT_FEEDBACK' + name: 'AGENT_SUBMIT_FEEDBACK(VARCHAR, NUMBER, VARCHAR, VARCHAR)' + type: procedure +``` + +The procedure is recreated on every `dbt run`, so changes to the feedback table schema are picked up automatically. The feedback table uses `CREATE TABLE IF NOT EXISTS`, so existing data is never dropped. + ## License Apache 2.0 diff --git a/integration_tests/dbt_project.yml b/integration_tests/dbt_project.yml index 7e84879..5356837 100644 --- a/integration_tests/dbt_project.yml +++ b/integration_tests/dbt_project.yml @@ -17,8 +17,7 @@ models: dbt_cortex_agent_integration_tests: +schema: INTEGRATION_TESTS -# Override test_cortex_search_service to point at a real Cortex Search Service. -# Snowflake Cortex Search is not a dbt object so it cannot be ref()'d. -# Passed via: dbt run --vars '{"test_cortex_search_service": "DB.SCHEMA.MY_SERVICE"}' vars: - test_cortex_search_service: "DB.SCHEMA.SEARCH_SERVICE" + # Points at the Cortex Search Service created by base_table's post_hook. + # The schema suffix matches dbt's custom-schema behaviour (target.schema ~ '_INTEGRATION_TESTS'). + test_cortex_search_service: "{{ target.database }}.{{ target.schema }}_INTEGRATION_TESTS.TEST_SEARCH_SERVICE" diff --git a/integration_tests/models/base_table.sql b/integration_tests/models/base_table.sql index 6e8aee2..0bc7bc5 100644 --- a/integration_tests/models/base_table.sql +++ b/integration_tests/models/base_table.sql @@ -1,7 +1,13 @@ -{{- config(materialized='table') -}} +{{- config( + materialized='table', + alias='BASE_TABLE', + post_hook="create or replace cortex search service {{ this.database }}.{{ this.schema }}.TEST_SEARCH_SERVICE on description attributes id warehouse = {{ target.warehouse }} target_lag = '7 days' as (select id, description from {{ this }})" +) -}} --- Minimal base table used by test_semantic_view in the integration test suite. -select - 1 as id, - 'test row' as description, - 100.00 as amount +-- Minimal base table used by test_semantic_view and TEST_SEARCH_SERVICE in the integration test suite. +select * from values + (1, 'Revenue figures show strong growth in the enterprise segment last quarter.', 100.00), + (2, 'Chargeback rates increased by 2% among small business clients in Q3.', 200.00), + (3, 'Top client by revenue is Acme Corp with 1.2M in annual recurring revenue.', 300.00), + (4, 'Integration test row for search indexing — safe to ignore in production.', 0.00) + as t(id, description, amount) diff --git a/integration_tests/models/cortex_agent_test.sql b/integration_tests/models/cortex_agent_test.sql index c9d1049..df9114a 100644 --- a/integration_tests/models/cortex_agent_test.sql +++ b/integration_tests/models/cortex_agent_test.sql @@ -10,7 +10,7 @@ }} models: - orchestration: claude-4-sonnet + orchestration: auto orchestration: budget: @@ -52,11 +52,45 @@ tools: required: - query + - tool_spec: + type: generic + name: AGENT_SUBMIT_FEEDBACK + description: 'Records user feedback about agent responses. Call when the user says "feedback" or explicitly rates or comments on a response. Always include the last 10 conversation messages. Example: user says "feedback: 1. Output wrong, expected $100" → call with rating=1, user_comment="Output wrong, expected $100". Example: user says "feedback: 5. Completely Correct" → call with rating=5, user_comment="Completely Correct".' + input_schema: + type: object + properties: + agent_name: + type: string + description: 'Name of this agent. Always pass "{{ this.identifier | upper }}".' + rating: + type: number + description: 'The user rating from 1 (worst) to 5 (best).' + user_comment: + type: string + description: 'Optional free-text feedback from the user.' + conversation_history: + type: string + description: 'Last 10 messages from the conversation as a JSON string, e.g. [{"role":"user","content":"..."}].' + required: + - agent_name + - rating + - conversation_history + +skills: [] + tool_resources: analyst_tool: - semantic_view: '{{ ref("test_semantic_view") }}' + semantic_view: '{{ ref('test_semantic_view') }}' search_tool: - name: '{{ var("test_cortex_search_service", "DB.SCHEMA.SEARCH_SERVICE") }}' + name: '{{ var("test_cortex_search_service") }}' max_results: 10 - title_column: title + title_column: description id_column: id + AGENT_SUBMIT_FEEDBACK: + execution_environment: + query_timeout: 300 + type: warehouse + warehouse: '' + identifier: '{{ this.database }}.{{ this.schema }}.AGENT_SUBMIT_FEEDBACK' + name: 'AGENT_SUBMIT_FEEDBACK(VARCHAR, NUMBER, VARCHAR, VARCHAR)' + type: procedure diff --git a/integration_tests/models/test_semantic_view.sql b/integration_tests/models/test_semantic_view.sql index fb36923..da33313 100644 --- a/integration_tests/models/test_semantic_view.sql +++ b/integration_tests/models/test_semantic_view.sql @@ -1,4 +1,4 @@ -{{ config(materialized='semantic_view') }} +{{ config(materialized='semantic_view', alias='TEST_SEMANTIC_VIEW') }} -- Minimal semantic view used as a tool_resource in cortex_agent_test. -- Requires the dbt_semantic_view package. diff --git a/integration_tests/tests/cortex_agent_test_has_feedback_proc.sql b/integration_tests/tests/cortex_agent_test_has_feedback_proc.sql new file mode 100644 index 0000000..09e43e1 --- /dev/null +++ b/integration_tests/tests/cortex_agent_test_has_feedback_proc.sql @@ -0,0 +1,10 @@ +-- Fails if the shared AGENT_SUBMIT_FEEDBACK stored procedure was not created by the materialization. +-- Returns 0 rows on success (standard dbt test contract). + +select 'AGENT_SUBMIT_FEEDBACK procedure not found' as error +where not exists ( + select 1 + from {{ target.database }}.information_schema.procedures + where procedure_schema = '{{ target.schema }}_INTEGRATION_TESTS' + and procedure_name = 'AGENT_SUBMIT_FEEDBACK' +) diff --git a/integration_tests/tests/cortex_agent_test_spec_content.sql b/integration_tests/tests/cortex_agent_test_spec_content.sql index 5b0677e..0d88686 100644 --- a/integration_tests/tests/cortex_agent_test_spec_content.sql +++ b/integration_tests/tests/cortex_agent_test_spec_content.sql @@ -13,7 +13,7 @@ with ddl as ( ), assertions as ( - select 'orchestration model missing' as error from ddl where not contains(content, 'claude-4-sonnet') + select 'orchestration model missing' as error from ddl where not contains(content, 'auto') union all select 'token budget missing' as error from ddl where not contains(content, '32000') union all diff --git a/macros/materializations/cortex_agent.sql b/macros/materializations/cortex_agent.sql index 2e97849..38a9e13 100644 --- a/macros/materializations/cortex_agent.sql +++ b/macros/materializations/cortex_agent.sql @@ -3,15 +3,27 @@ -- The materialization wraps it in CREATE OR REPLACE AGENT ... FROM SPECIFICATION $$ ... $$. -- -- Config options: --- comment (string, optional) : agent-level comment visible in Snowflake --- profile (string, optional) : JSON object with display_name, avatar, and color --- agent_grants (list, optional) : list of role names to grant USAGE on the agent +-- comment (string, optional) : agent-level comment visible in Snowflake +-- profile (string, optional) : JSON object with display_name, avatar, and color +-- agent_grants (list, optional) : list of role names to grant USAGE on the agent +-- create_feedback_table (bool, optional) : whether to create the feedback table and procedure. +-- Defaults to true. Set to false to skip. +-- feedback_schema (string, optional) : schema for the feedback table and AGENT_SUBMIT_FEEDBACK +-- procedure. Accepts 'SCHEMA' or 'DB.SCHEMA'. Defaults to +-- the agent's own database and schema. +-- feedback_table (string, optional) : fully-qualified override for the feedback table name. +-- Defaults to {feedback_schema}.AGENT_FEEDBACK. +-- Creates the table (if absent) and a stored procedure +-- named AGENT_SUBMIT_FEEDBACK on every dbt run. {% materialization cortex_agent, adapter='snowflake' %} - {%- set comment = config.get('comment', default=none) -%} - {%- set profile = config.get('profile', default=none) -%} - {%- set agent_grants = config.get('agent_grants', default=[]) -%} + {%- set comment = config.get('comment', default=none) -%} + {%- set profile = config.get('profile', default=none) -%} + {%- set agent_grants = config.get('agent_grants', default=[]) -%} + {%- set create_feedback_table = config.get('create_feedback_table', default=true) -%} + {%- set feedback_schema_config = config.get('feedback_schema', default=none) -%} + {%- set feedback_table = config.get('feedback_table', default=none) -%} {%- set target_relation = api.Relation.create( identifier=this.identifier, @@ -20,8 +32,36 @@ type='view' ) -%} + {%- if feedback_schema_config is not none -%} + {%- set _parts = feedback_schema_config.split('.') -%} + {%- if _parts | length == 2 -%} + {%- set feedback_db = _parts[0] -%} + {%- set feedback_schema = _parts[1] -%} + {%- else -%} + {%- set feedback_db = target_relation.database -%} + {%- set feedback_schema = feedback_schema_config -%} + {%- endif -%} + {%- else -%} + {%- set feedback_db = target_relation.database -%} + {%- set feedback_schema = target_relation.schema -%} + {%- endif -%} + + {%- if create_feedback_table and feedback_table is none -%} + {%- set feedback_table = feedback_db ~ '.' ~ feedback_schema ~ '.AGENT_FEEDBACK' -%} + {%- endif -%} + {{ run_hooks(pre_hooks) }} + {%- if create_feedback_table and feedback_table is not none %} + {% call statement('feedback_table') %} + {{ dbt_cortex_agent.snowflake__create_feedback_table(feedback_table) }} + {% endcall %} + + {% call statement('feedback_procedure') %} + {{ dbt_cortex_agent.snowflake__create_feedback_procedure(feedback_db, feedback_schema, feedback_table) }} + {% endcall %} + {%- endif %} + {% call statement('main') %} {{ dbt_cortex_agent.snowflake__create_cortex_agent(target_relation, sql, comment, profile) }} {% endcall %} diff --git a/macros/relations/cortex_agent/create_feedback_procedure.sql b/macros/relations/cortex_agent/create_feedback_procedure.sql new file mode 100644 index 0000000..a105b04 --- /dev/null +++ b/macros/relations/cortex_agent/create_feedback_procedure.sql @@ -0,0 +1,25 @@ +{% macro snowflake__create_feedback_procedure(feedback_db, feedback_schema, feedback_table) %} + + create or replace procedure + {{ feedback_db }}.{{ feedback_schema }}.AGENT_SUBMIT_FEEDBACK( + AGENT_NAME varchar, + RATING number, + USER_COMMENT varchar, + CONVERSATION_HISTORY varchar + ) + returns varchar + language sql + execute as caller + as + $$ + begin + let parsed_history variant := parse_json(:CONVERSATION_HISTORY); + insert into {{ feedback_table }} + (agent_name, user_name, rating, comment, conversation_history, created_at) + select + :AGENT_NAME, current_user(), :RATING, :USER_COMMENT, :parsed_history, current_timestamp(); + return 'Feedback submitted'; + end; + $$ + +{% endmacro %} diff --git a/macros/relations/cortex_agent/create_feedback_table.sql b/macros/relations/cortex_agent/create_feedback_table.sql new file mode 100644 index 0000000..6867fe7 --- /dev/null +++ b/macros/relations/cortex_agent/create_feedback_table.sql @@ -0,0 +1,13 @@ +{% macro snowflake__create_feedback_table(feedback_table) %} + + create table if not exists {{ feedback_table }} ( + feedback_id varchar default uuid_string(), + agent_name varchar, + user_name varchar, + rating number, + comment varchar, + conversation_history variant, + created_at timestamp_ntz default current_timestamp() + ) + +{% endmacro %}