Updated Streamlit UI and added functions for Cortex Agents

Author: sgsshankar

README.md

@@ -18,6 +18,9 @@ Get a hybrid (vector and keyword) search engine on your text data in minutes,
 - Use Cortex Search for your RAG Application
 - Use Cortex Search Powered Chat
 
+### Agent
+
+
 ## Prerequisites
 
 Before you begin, ensure you have met the following requirements:
@@ -60,6 +63,41 @@ To configure the application mode, update the `mode` parameter in `src/settings_
 - Set to `"debug"` for local development and editing.
 - Set to `"native"` for running natively in Snowflake.
 
+### Authentication for Cortex Agent – JWT Token Support Only
+
+This app currently supports **only JWT-based authentication** using RSA key pairs for Cortex Agents. OAuth is **not supported** at this time.
+
+#### Step 1: Generate RSA Key Pair
+
+```bash
+openssl genrsa 2048 | openssl pkcs8 -topk8 -inform PEM -out rsa_private_key.p8 -nocrypt
+openssl rsa -in rsa_private_key.p8 -pubout -out rsa_public_key.pub
+```
+
+#### Step 2: Upload Public Key to Snowflake
+
+1. Copy the content from `rsa_public_key.pub` (exclude `-----BEGIN PUBLIC KEY-----` headers).
+2. In Snowflake, run:
+
+```sql
+alter user <your_username> set rsa_public_key='<paste_public_key_here>';
+```
+
+3. Retrieve the public key fingerprint:
+
+```sql
+desc user <your_username>;
+```
+
+Use the fingerprint in `rsa_public_key_fp`.
+
+#### Step 3: Add to `settings_config.json`
+
+- Paste full contents of `rsa_private_key.p8` into `rsa_private_key` (escape newlines if needed).
+- Paste fingerprint into `rsa_public_key_fp`.
+
+For more details, refer to the [Snowflake documentation](https://docs.snowflake.com/en/user-guide/key-pair-auth.html#using-key-pair-authentication).
+
 ## Running the App
 
 ### Locally

requirements.txt

@@ -4,4 +4,5 @@ PyPDF2
 langchain
 pandas
 pyarrow
-snowflake
+snowflake
+PyJWT

snowflake.yml

@@ -1,4 +1,4 @@
-definition_version: "1"
+definition_version: '1'
 streamlit:
   name: snowflake_ai_toolkit
   stage: sf_ai_stage
@@ -8,6 +8,7 @@ streamlit:
   #pages_dir: pages/
   additional_source_files:
     - src/cortex_functions.py
+    - src/cortex_agent.py
     - src/utils.py
     - src/html_snippets.py
     - src/styles.css

src/search.py

@@ -7,7 +7,7 @@
 import threading
 import json
 from snowflake.core import Root
-
+import os
 
 config_path = Path("src/settings_config.json")
 with open(config_path, "r") as f:
@@ -29,7 +29,7 @@ def display_search(session):
     #     st.write(res)
 
 
-    warehouse = config["snowflake"]["warehouse"]
+    warehouse = config["warehouse"]
 
     if create_or_use == "Create":
 

src/utils.py

@@ -284,10 +284,6 @@ def create_prompt_for_rag(session, question: str, rag: bool, column: str, databa
         """
     return prompt
 
-def create_prompt_for_search_service(session, question: str, search_service: str, chat_history: list):
-
-    return
-
 def get_cortex_complete_result(session, query: str):
     """
     Executes a Cortex complete query and returns the result.
@@ -394,8 +390,8 @@ def create_database_and_stage_if_not_exists(session: Session):
         session (Session): Snowflake session object
     """
     # Fetch database and stage details from the config file
-    database_name = config["snowflake"]["database"]
-    stage_name = config["snowflake"]["stage"]
+    database_name = config["database"]
+    stage_name = config["stage"]
 
     # Check if the database exists, and create if it doesn't
     database_query = f"SHOW DATABASES LIKE '{database_name}'"

streamlit_app.py

@@ -7,6 +7,10 @@
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from src.notification import *
 from src.search import *
+from src.cortex_agent import *
+# from trulens.connectors.snowflake import SnowflakeConnector
+# from trulens.core.session import TruSession
+# from trulens.dashboard import run_dashboard
 
 # Load the config file
 config_path = Path("src/settings_config.json")
@@ -33,10 +37,15 @@
     elif config["mode"] == "debug":
         try:
             connection_parameters = {
-                "account": config["snowflake"]["account"],"user": config["snowflake"]["user"],"password": config["snowflake"]["password"],
-                "role": config["snowflake"]["role"],"warehouse": config["snowflake"]["warehouse"],"database": config["snowflake"]["database"],
-                "schema": config["snowflake"]["schema"]
+                "account": config["account"],"user": config["user"],"password": config["password"],
+                "role": config["role"],"warehouse": config["warehouse"],"database": config["database"],
+                "schema": config["schema"],
+                # "init_server_side": True
             }
+            # connector = SnowflakeConnector(**connection_parameters)
+            # truelens_session = TruSession(connector)
+            # st.session_state.truelens_session = truelens_session
+            # run_dashboard(truelens_session)
             st.session_state.snowflake_session = Session.builder.configs(connection_parameters).create()
         except Exception as e:
             st.error(f"Failed to create session in debug mode: {e}")
@@ -114,6 +123,8 @@
             st.session_state.page = "Build"
         if st.button("🔍 Cortex Search"):
             st.session_state.page = "Cortex Search"
+        if st.button("🤖 Cortex Agent"):
+            st.session_state.page = "Cortex Agent"
         if st.button("🔔 Notification"):
             st.session_state.page = "Notification"
 
@@ -124,6 +135,7 @@
     "Build": build.display_build,
     "Notification": notification.display_notification,
     "Cortex Search": display_search,
+    "Cortex Agent": display_cortex_agent,
     "Setup": setup.display_setup
 }