SNOW-1449796 Snowflake Cortex Example (snowflakedb#5)

sfc-gh-ralfaroviquez · sfc-gh-osalazarlizano · sfc-gh-cgorrie · web-flow · commit 8bdd146ef610 · 2024-07-04T12:09:22.000-06:00
* cortex initial commit

* cortex initial commit missing file

* draft for cortex, missing db integrations and several improvements

* minor improvements, using real dataset

* using cortex from ml library, adding chat, adding test sample

* testing attempts

* fix test

* Update dashboard.py

* tests fixed, added instructions in the readme

* trying to fix ci runtest flow

* adding cortex lib to ci env

* adding snowflake-cortex project to readme index

* Update README.md

Co-authored-by: Cam Gorrie &lt;cam.gorrie@snowflake.com&gt;

* Update snowflake-cortex/README.md

Co-authored-by: Cam Gorrie &lt;cam.gorrie@snowflake.com&gt;

* Update snowflake-cortex/README.md

Co-authored-by: Cam Gorrie &lt;cam.gorrie@snowflake.com&gt;

* addressing comments

* changing name to dataset columns and db, fixing some wrong commands, deleting version from manifest

* repeated dot deletion

* changing name from dashboard to ui

* added logic for ungiven permissions

* fixing tests and adding more

* pushing to fix path error

* attempt to fix the path, deleting init.py

---------

Co-authored-by: Oscar Salazar &lt;oscar.salazarlizano@snowflake.com&gt;
Co-authored-by: Cam Gorrie &lt;cam.gorrie@snowflake.com&gt;
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -104,8 +104,8 @@ jobs:
               python -m pip install pytest
       - name: Run tests
         run: |
-          args=${{ steps.tests_to_run.outputs.pytestArgs }}
-          pythonpath=${{ steps.tests_to_run.outputs.pytestPaths }}
+          args="${{ steps.tests_to_run.outputs.pytestArgs }}"
+          pythonpath="${{ steps.tests_to_run.outputs.pytestPaths }}"
           if [ -z "${args}" ] || [ -z "${pythonpath}" ]; then
             echo “Nothing to test”
           else
diff --git a/README.md b/README.md
@@ -27,6 +27,7 @@ Some applications require other account-level setup before they can be properly
 | [Object-level References](./object-level-references/) | A simple dashboard to show how to interact with the object-level references and bindings. |
 | [Reference Usage](./reference-usage/) | How to share a provider table with a native application whose data is replicated to any consumer in the data cloud. |
 | [SPCS Three-tier](./spcs-three-tier/) | A simple three-tiered web app that can be deployed in Snowpark Container Services. It queries the TPC-H 100 data set and returns the top sales clerks. |
+| [Snowflake Cortex](./snowflake-cortex/) | A simple example on how to implement the Cortex Complete and to make it interact with user data. |
 | [Tasks and Streams](./tasks-streams/) | How to execute a task and visualize changes using streams within a native application. |
 
 ## Contributing
diff --git a/object-level-references/app/manifest.yml b/object-level-references/app/manifest.yml
@@ -9,7 +9,7 @@ version:
 artifacts:
   setup_script: setup_script.sql
   readme: README.md
-  default_streamlit: core.dashboard
+  default_streamlit: core.ui
   extension_code: true 
   
 references:
diff --git a/shared_python_ci_env.yml b/shared_python_ci_env.yml
@@ -12,4 +12,5 @@ dependencies:
       - snowflake-cli-labs>=2.0.0
       - pytest
       - streamlit>=1.26.0
+      - snowflake-ml-python
 
diff --git a/snowflake-cortex/.gitignore b/snowflake-cortex/.gitignore
@@ -0,0 +1,4 @@
+snowflake.local.yml
+output/
+**/__pycache__/
+**/.pytest_cache/
diff --git a/snowflake-cortex/README.md b/snowflake-cortex/README.md
@@ -0,0 +1,29 @@
+# Snowflake Cortex
+
+This simple Native App shows how to use Cortex Complete and make it interact with user data.
+
+For this use case, the dataset used is rather small: only 10 entries. This is because the language model used in the Cortex function restricts input data size. You can change the language model used to a bigger / different one, according to your needs.
+For more information about it please visit **[this page](https://docs.snowflake.com/en/user-guide/snowflake-cortex/llm-functions#cost-considerations)**.
+
+## Data preparation
+
+To run this example first execute this command, that is going to create a database and table with information about songs charts:
+```sh
+snow sql -f 'prepare/provider_data.sql'
+```
+## App execution
+
+Then run `snow app run` on your terminal.
+
+## App and Data Deletion
+To delete the database and the app run
+
+```sh
+snow sql -q 'DROP DATABASE SONGS_CORTEX_DB;'
+snow app teardown
+```
+
+## Further reading
+
+For more information about the different ways to use snowflake AI capabilities visit this page:
+**[Snowflake AI and ML documentation](https://docs.snowflake.com/en/guides-overview-ai-features)**
diff --git a/snowflake-cortex/app/README.md b/snowflake-cortex/app/README.md
@@ -0,0 +1,3 @@
+# Snowflake Native App - Snowflake Cortex
+
+This is a sample Snowflake Native App that shows the use of Snowflake Cortex inside Native Apps.
diff --git a/snowflake-cortex/app/manifest.yml b/snowflake-cortex/app/manifest.yml
@@ -0,0 +1,12 @@
+# For more information on creating manifest, go to https://docs.snowflake.com/en/developer-guide/native-apps/creating-manifest
+manifest_version: 1
+
+artifacts:
+  setup_script: setup_script.sql
+  readme: README.md
+  default_streamlit: core.ui
+  extension_code: true 
+
+privileges:
+  - IMPORTED PRIVILEGES ON SNOWFLAKE DB:
+      description: "Imported privileges to use cortex DB"
diff --git a/snowflake-cortex/app/setup_script.sql b/snowflake-cortex/app/setup_script.sql
@@ -0,0 +1,21 @@
+-- This is the setup script that runs while installing a Snowflake Native App in a consumer account.
+-- For more information on how to create setup file, visit https://docs.snowflake.com/en/developer-guide/native-apps/creating-setup-script
+
+-- A general guideline to building this script looks like:
+-- 1. Create application roles
+CREATE APPLICATION ROLE IF NOT EXISTS app_public;
+
+-- 2. Create a versioned schema to hold those UDFs/Stored Procedures
+CREATE OR ALTER VERSIONED SCHEMA core;
+GRANT USAGE ON SCHEMA core TO APPLICATION ROLE app_public;
+
+-- 3. Create a streamlit object using the code you wrote in you wrote in src/module-ui, as shown below. 
+-- The `from` value is derived from the stage path described in snowflake.yml
+CREATE STREAMLIT core.ui
+     FROM '/streamlit/'
+     MAIN_FILE = 'dashboard.py';
+
+-- 4. Grant appropriate privileges over these objects to your application roles. 
+GRANT USAGE ON STREAMLIT core.ui TO APPLICATION ROLE app_public;
+
+-- A detailed explanation can be found at https://docs.snowflake.com/en/developer-guide/native-apps/adding-streamlit 
diff --git a/snowflake-cortex/local_test_env.yml b/snowflake-cortex/local_test_env.yml
@@ -0,0 +1,15 @@
+# This file is used to install packages for local testing
+name: snowflake-cortex-testing
+channels:
+  - snowflake
+dependencies:
+  - python=3.8
+  - pip
+  - pip:
+      - snowflake-native-apps-permission-stub
+      - snowflake-snowpark-python>=1.15.0
+      - snowflake-cli-labs>=2.0.0
+      - pytest
+      - streamlit>=1.26.0
+      - snowflake-ml-python
+
diff --git a/snowflake-cortex/prepare/provider_data.sql b/snowflake-cortex/prepare/provider_data.sql
@@ -0,0 +1,25 @@
+USE ROLE ACCOUNTADMIN;
+CREATE DATABASE IF NOT EXISTS SONGS_CORTEX_DB;
+CREATE SCHEMA IF NOT EXISTS SONGS_CORTEX_DB.SONGS_CORTEX_SCHEMA;
+CREATE OR REPLACE TABLE SONGS_CORTEX_DB.SONGS_CORTEX_SCHEMA.SONGS_PROVIDER_DATA(
+    track_name VARCHAR,
+    artists_name VARCHAR,
+    artist_count INTEGER,
+    released_year INTEGER,
+    released_month INTEGER,
+    released_day INTEGER,
+    in_song_charts INTEGER,
+    song_streams INTEGER,
+    danceability_percentage INTEGER
+);
+INSERT INTO SONGS_CORTEX_DB.SONGS_CORTEX_SCHEMA.SONGS_PROVIDER_DATA VALUES
+('Seven \(feat. Latto\) \(Explicit Ver.\)','Latto, Jung Kook',2,2023,7,14,147,141381703,80),
+('LALA','Myke Towers',1,2023,3,23,48,133716286,71),
+('vampire','Olivia Rodrigo',1,2023,6,30,113,140003974,51),
+('Cruel Summer','Taylor Swift',1,2019,8,23,100,800840817,55),
+('WHERE SHE GOES','Bad Bunny',1,2023,5,18,50,303236322,65),
+('Sprinter','Dave, Central Cee',2,2023,6,1,91,183706234,92),
+('Ella Baila Sola','Eslabon Armado, Peso Pluma',2,2023,3,16,50,725980112,67),
+('Columbia','Quevedo',1,2023,7,7,43,58149378,67),
+('fukumean','Gunna',1,2023,5,15,83,95217315,85),
+('La Bebe - Remix','Peso Pluma, Yng Lvcas',2,2023,3,17,44,553634067,81);
diff --git a/snowflake-cortex/pytest.ini b/snowflake-cortex/pytest.ini
@@ -0,0 +1,2 @@
+[pytest]
+pythonpath=src/ui
diff --git a/snowflake-cortex/scripts/shared_content.sql b/snowflake-cortex/scripts/shared_content.sql
@@ -0,0 +1,16 @@
+-- Grant usage to a external database from the application.
+create schema if not exists {{ package_name }}.PACKAGE_SHARED;
+use schema {{ package_name }}.PACKAGE_SHARED;
+
+grant reference_usage on database SONGS_CORTEX_DB
+    to share in application package {{ package_name }};
+
+-- Create a view that references the provider table.
+-- The view is going to be shared by the package to the application.
+create view if not exists PACKAGE_SHARED.PROVIDER_SONGS_VIEW
+  as select * from SONGS_CORTEX_DB.SONGS_CORTEX_SCHEMA.SONGS_PROVIDER_DATA;
+
+grant usage on schema PACKAGE_SHARED
+  to share in application package {{ package_name }};
+grant select on view PACKAGE_SHARED.PROVIDER_SONGS_VIEW
+  to share in application package {{ package_name }};
diff --git a/snowflake-cortex/snowflake.yml b/snowflake-cortex/snowflake.yml
@@ -0,0 +1,12 @@
+definition_version: 1
+native_app:
+  name: snowflake_cortex
+  source_stage: app_src.stage
+  artifacts:
+    - src: app/*
+      dest: ./
+    - src: src/ui/*
+      dest: streamlit/
+  package:
+    scripts:
+      - scripts/shared_content.sql
diff --git a/snowflake-cortex/src/ui/cortexCaller.py b/snowflake-cortex/src/ui/cortexCaller.py
@@ -0,0 +1,14 @@
+from snowflake.cortex import Complete
+
+class CortexCaller:
+
+   def call_complete(_self,input_with_table):
+      response = Complete('llama2-70b-chat', input_with_table)
+      return response
+
+   def call_cortex(_self, df, input: str):
+      sample_df = df.to_string()
+      input_with_table = f"""The following data is a table that contains songs information it is wrapped in this tags <dataset>{sample_df}</dataset> {input}"""
+      response = _self.call_complete(input_with_table)
+      return response
+   
diff --git a/snowflake-cortex/src/ui/dashboard.py b/snowflake-cortex/src/ui/dashboard.py
@@ -0,0 +1,47 @@
+import streamlit as st
+from snowflake.snowpark import Session
+from cortexCaller import CortexCaller
+import snowflake.permissions as permission
+
+class Dashboard:
+   
+   def __init__(self, session: Session, cortex_caller: CortexCaller = CortexCaller()) -> None:
+      self.session = session
+      self.cortex_caller = cortex_caller
+
+   def setup(self):
+      st.header("Privileges setup")
+      st.caption("""
+         Follow the instructions below to set up your application.
+         Once you have completed the steps, you will be able to continue to the main example.
+      """)
+      privilege = 'IMPORTED PRIVILEGES ON SNOWFLAKE DB'
+      if not permission.get_held_account_privileges([privilege]):
+         st.button(f"Grant import privileges on snowflake DB ↗", on_click=permission.request_account_privileges, args=[[privilege]], key='IMPORTED PRIVILEGES ON SNOWFLAKE DB')
+      else:
+         st.session_state.privileges_granted = True
+         st.rerun()
+
+   def run_streamlit(self):
+      if not permission.get_held_account_privileges(['IMPORTED PRIVILEGES ON SNOWFLAKE DB']):
+         del st.session_state.privileges_granted
+         st.rerun()
+      
+      st.header('Snowflake Cortex Example')
+      st.subheader("Simple app showing how cortex can answer questions using the following song's ranking table")
+      table_df = self.session.table("PACKAGE_SHARED.PROVIDER_SONGS_VIEW").to_pandas()
+      st.write(table_df)
+
+      cortex_chat = st.container(height=200)
+      cortex_input = st.chat_input("What do you want to ask to cortex about the table from above?", key="cortex_input")
+      if cortex_input and not cortex_input.isspace():
+         cortex_chat.chat_message("user").write(cortex_input)
+         cortex_fn = self.cortex_caller.call_cortex(table_df, cortex_input)
+         cortex_chat.chat_message("assistant").write(cortex_fn)
+
+
+if __name__ == '__main__':
+   if 'privileges_granted' not in st.session_state:
+      Dashboard(Session.builder.getOrCreate()).setup()
+   else:
+      Dashboard(Session.builder.getOrCreate()).run_streamlit()
diff --git a/snowflake-cortex/src/ui/environment.yml b/snowflake-cortex/src/ui/environment.yml
@@ -0,0 +1,9 @@
+# This file is used to install packages used by the Streamlit App.
+# For more details, refer to https://docs.snowflake.com/en/developer-guide/streamlit/create-streamlit-sql#label-streamlit-install-packages-manual
+
+channels:
+- snowflake
+dependencies:
+- streamlit=1.31.0
+- snowflake-ml-python
+- snowflake-native-apps-permission
diff --git a/snowflake-cortex/test/test_cortex.py b/snowflake-cortex/test/test_cortex.py
@@ -0,0 +1,20 @@
+import unittest
+from unittest.mock import patch, MagicMock
+import pandas as pd
+import pytest as pytest
+from snowflake import cortex
+from snowflake.cortex import Complete
+
+class TestDataExporter(unittest.TestCase):
+
+    def test_cortex_response(self):
+        from cortexCaller import CortexCaller
+
+        mock_df = pd.DataFrame([{"TestColumn": "TestValue"}])
+        mock_input = 'Test input from user'
+        cortexCaller = CortexCaller()
+        cortexCaller.call_complete = MagicMock(name='call_complete')
+        cortexCaller.call_complete.return_value = 'test answer from cortex'
+        response = cortexCaller.call_cortex(mock_df, mock_input)
+        cortexCaller.call_complete.assert_called_once()
+        assert response == cortexCaller.call_complete.return_value
diff --git a/snowflake-cortex/test/test_ui.py b/snowflake-cortex/test/test_ui.py
@@ -0,0 +1,102 @@
+from cortexCaller import CortexCaller
+from streamlit.testing.v1 import AppTest
+from snowflake.snowpark import Session
+from unittest.mock import patch, MagicMock
+import pytest as pytest
+
+@pytest.fixture(autouse=True)
+def session():
+    session = Session.builder.config('local_testing', True).create()
+    yield session
+    session.close()
+
+@pytest.fixture(autouse=True)
+def cortex_caller():
+    cortex_caller = CortexCaller()
+    yield cortex_caller
+
+
+@patch('snowflake.permissions.request_account_privileges')
+def test_setup(request_account_privileges: MagicMock, session, cortex_caller):
+
+    def script(session, cortex_caller):
+        from dashboard import Dashboard
+        dashboard = Dashboard(session, cortex_caller)
+        return dashboard.setup()
+        
+    with patch('snowflake.permissions.request_account_privileges', return_value=True) as request_account_privileges:
+        at = AppTest.from_function(script, kwargs={ "session": session, "cortex_caller": cortex_caller }).run()
+        assert at.header[0].value == 'Privileges setup'
+        assert at.caption[0].value == 'Follow the instructions below to set up your application.\nOnce you have completed the steps, you will be able to continue to the main example.'
+        assert at.button[0].label == 'Grant import privileges on snowflake DB ↗'  
+        assert 'privileges_granted' not in at.session_state
+        at.button(key='IMPORTED PRIVILEGES ON SNOWFLAKE DB').click().run()
+        request_account_privileges.assert_called_once_with(['IMPORTED PRIVILEGES ON SNOWFLAKE DB'])
+
+
+@patch('snowflake.permissions.get_held_account_privileges')
+def test_not_granted_privileges(get_held_account_privileges: MagicMock, session, cortex_caller):
+
+    get_held_account_privileges.return_value = False
+
+    at = AppTest.from_file("../src/ui/dashboard.py").run()
+    assert at.header[0].value == 'Privileges setup'
+    assert at.caption[0].value == 'Follow the instructions below to set up your application.\nOnce you have completed the steps, you will be able to continue to the main example.'
+    assert at.button[0].label == 'Grant import privileges on snowflake DB ↗'
+
+@patch('cortexCaller.CortexCaller.call_cortex')
+@patch('snowflake.snowpark.session.Session.table')
+@patch('snowflake.permissions.get_held_account_privileges')
+def test_cortex_call_mock(get_held_account_privileges: MagicMock, table: MagicMock, call_cortex: MagicMock, session, cortex_caller):
+    
+    get_held_account_privileges.return_value = True
+
+    def script(session, cortex_caller):
+        from dashboard import Dashboard
+        dashboard = Dashboard(session, cortex_caller)
+        return dashboard.run_streamlit()
+    
+    table.return_value=session.create_dataframe([{"TestColumn": "TestValue"}])
+    call_cortex.return_value='Mocked response from cortex'
+    at = AppTest.from_function(script, kwargs={ "session": session, "cortex_caller": cortex_caller }).run()
+    
+    at.chat_input(key='cortex_input').set_value('text for testing purposes').run()
+    call_cortex.assert_called_once()
+    assert at.chat_message[1].name == 'assistant'
+
+@patch('cortexCaller.CortexCaller.call_cortex')
+@patch('snowflake.snowpark.session.Session.table')
+@patch('snowflake.permissions.get_held_account_privileges')
+def test_input_showed_in_chat(get_held_account_privileges: MagicMock, table: MagicMock, call_cortex: MagicMock, session, cortex_caller):
+
+    get_held_account_privileges.return_value = True
+
+    def script(session, cortex_caller):
+        from dashboard import Dashboard
+        dashboard = Dashboard(session, cortex_caller)
+        return dashboard.run_streamlit()
+    
+    table.return_value=session.create_dataframe([{"TestColumn": "TestValue"}])
+    at = AppTest.from_function(script, kwargs={ "session": session, "cortex_caller": cortex_caller }).run()
+    at.chat_input(key='cortex_input').set_value('text for testing purposes').run()
+    assert at.chat_message[0].name == 'user'
+
+@patch('cortexCaller.CortexCaller.call_cortex')
+@patch('snowflake.snowpark.session.Session.table')
+@patch('snowflake.permissions.get_held_account_privileges')
+def test_input_empty(get_held_account_privileges: MagicMock, table: MagicMock, call_cortex: MagicMock, session, cortex_caller):
+    
+    get_held_account_privileges.return_value = True
+
+    def script(session, cortex_caller):
+        from dashboard import Dashboard
+        dashboard = Dashboard(session, cortex_caller)
+        return dashboard.run_streamlit()
+
+    table.return_value=session.create_dataframe([{"TestColumn": "TestValue"}])
+    at = AppTest.from_function(script, kwargs={ "session": session, "cortex_caller": cortex_caller }).run()
+    at.chat_input(key='cortex_input').set_value(' ').run()
+    assert not at.chat_message
+
+    at.chat_input(key='cortex_input').set_value('').run()
+    assert not at.chat_message

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# Snowflake Native App - Snowflake Cortex`
	`2`	`+`
	`3`	`+This is a sample Snowflake Native App that shows the use of Snowflake Cortex inside Native Apps.`