Snowflake-Labs · sfc-gh-dtaylor · May 22, 2024
diff --git a/01_querying_rest_api/README.md b/01_querying_rest_api/README.md
@@ -0,0 +1,61 @@
+# Querying the Cortex Search REST API
+
+## Prerequisites
+
+Before you can run the script, ensure you have the following prerequisites installed:
+
+- Python 3.x
+- pip (Python package installer)
+  Additionally, you must have access to a Snowflake account and the required permissions to query the Cortex Search Service at the specified database and schema.
+
+Installation
+First, clone this repository to your local machine using git:
+
+```
+git clone https://github.com/snowflake-labs/cortex-search.git
+cd cortex-search/01_querying_rest_api
+```
+
+Install the necessary Python dependencies by running:
+
+```
+pip install -r requirements.txt
+```
+
+## Key-pair auth configuration
+
+Additionally, you must generate a private key for JWT auth with Snowflake as described in [this document](https://docs.snowflake.com/user-guide/key-pair-auth#configuring-key-pair-authentication).
+
+**Note**: take note of the path to your generated RSA private key, e.g., `/path/to/my/rsa_key.p8` -- you will need to supply this as the `--private-key-path` parameter to query the service later from the command line, or list the path to the file from within a notebook.
+
+## Command line usage
+
+The `simple_query.py` example script can be executed from the command line. For instance:
+
+```
+python3 examples/simple_query.py -u https://my_org-my_account.us-west-2.aws.snowflakecomputing.com -s DB.SCHEMA.SERVICE_NAME -q "the sky is blue" -c "description,text" -l 10 -a my_account -k /path/to/my/rsa_key.p8 -n my_name
+```
+
+**Arguments:**
+
+- `-u`, `--url`: URL of the Snowflake instance. See [this guide](https://docs.snowflake.com/en/user-guide/admin-account-identifier#finding-the-organization-and-account-name-for-an-account) for finding your Account URL
+- `-s`, `--qualified-service-name`: The fully-qualified Cortex Search Service name, in the format DATABASE.SCHEMA.SERVICE
+- `-q`, `--query`: The search query string
+- `-c`, `--columns`: Comma-separated list of columns to return in the results
+- `-l`, `--limit`: The max number of results to return
+- `-a`, `--account`: Snowflake account name. See [this guide](https://docs.snowflake.com/en/user-guide/admin-account-identifier#finding-the-organization-and-account-name-for-an-account) for finding your Account name
+- `-k`, `--private-key-path`: Path to the RSA private key file for authentication.
+- `-n`, `--user-name`: Username for the Snowflake account
+- `-r`, `--role`: Role to use for the query. If provided, a session token scoped to this role will be created and used for authentication to the API.
+
+The `interactive_query.py` example provides an interactive CLI that demonstrates caching the JWT used for authentication between requests for better performance and implements retries when the JWT has expired. You can run it like the following:
+
+```
+python3 examples/interactive_query.py -u https://my_org-my_account.us-west-2.aws.snowflakecomputing.com -s DB.SCHEMA.SERVICE_NAME -c "description,text" -a my_account -k /path/to/my/rsa_key.p8 -n my_name
+```
+
+This will launch an interactive session, where you will be prompted repeatedly for search queries to your Cortex Search Service.
+
+## License
+
+Apache Version 2.0
diff --git a/examples/interactive_query.py → ...ng_rest_api/examples/interactive_query.py b/examples/interactive_query.py → ...ng_rest_api/examples/interactive_query.py
diff --git a/examples/simple_query.py → ...uerying_rest_api/examples/simple_query.py b/examples/simple_query.py → ...uerying_rest_api/examples/simple_query.py
diff --git a/requirements.txt → 01_querying_rest_api/requirements.txt b/requirements.txt → 01_querying_rest_api/requirements.txt
diff --git a/src/__init__.py → 01_querying_rest_api/src/__init__.py b/src/__init__.py → 01_querying_rest_api/src/__init__.py
diff --git a/src/utils/jwtutils.py → 01_querying_rest_api/src/utils/jwtutils.py b/src/utils/jwtutils.py → 01_querying_rest_api/src/utils/jwtutils.py
diff --git a/src/utils/queryutils.py → 01_querying_rest_api/src/utils/queryutils.py b/src/utils/queryutils.py → 01_querying_rest_api/src/utils/queryutils.py
diff --git a/02_querying_python_api/notebook_query.ipynb b/02_querying_python_api/notebook_query.ipynb
@@ -0,0 +1,141 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "15a04599",
+   "metadata": {},
+   "source": [
+    "# Querying a Cortex Search Service with the Python API\n",
+    "This notebook shows a simple demo of querying a Cortex Search Service with the Python API.  \n",
+    "The documentation for this query pattern can be [accessed here](https://docs.snowflake.com/LIMITEDACCESS/cortex-search/query-cortex-search-service).\n",
+    "\n",
+    "\n",
+    "## Prerequisites\n",
+    "To install the required packages in your python environment, run:  \n",
+    "    ``pip install snowflake snowflake-snowpark-python``\n",
+    "\n",
+    "\n",
+    "Note: Querying Cortex Search requires version 0.8.0 or later of the `snowflake` package."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8549cdb7-3b82-416b-b8d0-900732227212",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from snowflake.core import Root # snowflake >= 0.8.0\n",
+    "from snowflake.snowpark.session import Session\n",
+    "\n",
+    "# Set connection parameters\n",
+    "SNOWFLAKE_ACCOUNT = \"<ACCOUNT>\"\n",
+    "SNOWFLAKE_USER = \"<USER>\"\n",
+    "SNOWFLAKE_PASSWORD = \"<PASS>\"\n",
+    "SNOWFLAKE_WAREHOUSE = \"<WH>\"\n",
+    "\n",
+    "CORTEX_SEARCH_DATABASE = \"<DB>\"\n",
+    "CORTEX_SEARCH_SCHEMA = \"<SCHEMA>\"\n",
+    "CORTEX_SEARCH_SERVICE = \"<SERVICE>\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1ae0e607-a9aa-400e-8bb7-a9b6945cccaf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_session():\n",
+    "    \"\"\"\n",
+    "    Create Snowpark Session from connection parameters\n",
+    "    \"\"\"\n",
+    "    connection_parameters = {\n",
+    "        \"user\": SNOWFLAKE_USER,\n",
+    "        \"password\": SNOWFLAKE_PASSWORD,\n",
+    "        # \"authenticator\": \"externalbrowser\",\n",
+    "        \"account\": SNOWFLAKE_ACCOUNT,\n",
+    "        \"warehouse\": SNOWFLAKE_WAREHOUSE,\n",
+    "        \"database\": CORTEX_SEARCH_DATABASE,\n",
+    "        \"schema\": CORTEX_SEARCH_SCHEMA\n",
+    "    }\n",
+    "\n",
+    "    return Session.builder.configs(connection_parameters).create()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "00202c33-97e0-4860-b34a-fe599d3948b2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# make session and root objects\n",
+    "session = make_session()\n",
+    "root = Root(session)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "65286440",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def query_cortex_search_service(svc, query, columns=[], _filter = {}, limit=5):\n",
+    "    \"\"\"\n",
+    "    Query the specified Cortex Search service with the specified query string and search parameters.\n",
+    "    \"\"\"\n",
+    "    db, schema = session.get_current_database(), session.get_current_schema()\n",
+    "    \n",
+    "    cortex_search_service = (\n",
+    "        root\n",
+    "        .databases[db]\n",
+    "        .schemas[schema]\n",
+    "        .cortex_search_services[svc]\n",
+    "    )\n",
+    "    \n",
+    "    response = cortex_search_service.search(\n",
+    "        query,\n",
+    "        columns=columns,\n",
+    "        filter=_filter,\n",
+    "        limit=limit\n",
+    "    )\n",
+    "\n",
+    "    return response.results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c2bb1d8d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# query the service\n",
+    "print(query_cortex_search_service(svc=CORTEX_SEARCH_SERVICE, query=\"foo bar\"))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "cs-quickstart",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/README.md b/README.md
@@ -1,6 +1,12 @@
 # Cortex Search
 
-This repository contains example usage, including authentication, for the Cortex Search REST API, currently in Private Preview. The official preview documentation can be [found here](https://docs.snowflake.com/LIMITEDACCESS/cortex-search/cortex-search-overview).
+This repository contains example usage, including authentication, for Cortex Search, currently in Private Preview. The official preview documentation can be [found here](https://docs.snowflake.com/LIMITEDACCESS/cortex-search/cortex-search-overview).
+
+This repository contains the following examples:
+- [01. Querying the REST API](/01_querying_rest_api/README.md)
+- [02. Querying the Python API](/02_querying_python_api/README.md)
+- [03. Building an AI search app in Streamlit](/03_streamlit_ai_search/README.md)
+
 
 ## Prerequisites