Kumo

https://kumo.ai/docs
Admin
Kumo is an AI platform that enables data scientists and engineers to build predictive models and...
Tokens:179,826
Snippets:1,687
Trust Score:9.6
Update:10 months ago
Show doc for...
Tokens:
Raw
### Install Kumo AI Python SDK

Source: https://kumo.ai/docs/rfm/RFM-quickstart

Installs the Kumo AI Python SDK using pip. Ensure you have Python and pip installed on your system.

```python
pip install kumoai
```

--------------------------------

### Build and Link Tables into a LocalGraph in Python

Source: https://kumo.ai/docs/rfm/RFM-graph-creation

This example shows how to initialize a LocalGraph with specified tables and then establish a link between two tables using a foreign key. This is essential for creating the relational structure needed for graph-based predictions. It uses the 'rfm' library.

```python
# step 1: select tables
graph = rfm.LocalGraph(tables=[users,orders])

# step 2: link the tables 
# In the orders table (src_table), there exists a column named user_id (fkey), which we can use as a foreign key to link to the primary key in the users table (dst_table).
graph.link(src_table="orders", fkey="user_id", dst_table="users");
```

--------------------------------

### Start Kumo AI Application using SQL

Source: https://kumo.ai/docs/installing-kumo-on-spcs

This SQL script initiates the Kumo AI application, requiring user-defined warehouse and role configurations. It returns a URL for accessing the Kumo UI and can take 20-30 minutes to complete. Ensure the specified KUMO_USER_ROLE has necessary warehouse privileges.

```sql
-----------------------------
---  USER INPUT REQUIRED  ---
-----------------------------
-- Note that KUMO_USER_ROLE should be granted usage on this warehouse.
SET USER_WAREHOUSE = '<ADD WAREHOUSE TO USE>';
-- Modify the following if a different name was used to install the app.
SET KUMOAPP_NAME = 'KUMO';

USE ROLE KUMO_USER_ROLE;

USE WAREHOUSE IDENTIFIER($USER_WAREHOUSE);
SET KUMOAPP_COMPUTE = CONCAT(($KUMOAPP_NAME),'.PROCEDURES.START_APP');
CALL IDENTIFIER($KUMOAPP_COMPUTE)('USER_SCHEMA');

```

--------------------------------

### Initialize Kumo AI Client

Source: https://kumo.ai/docs/rfm/RFM-quickstart

Initializes the Kumo AI client by setting the KUMO_API_KEY environment variable with your generated API key. This client is used to interact with Kumo AI services for making predictions.

```python
import kumoai.experimental.rfm as rfm, os

os.environ["KUMO_API_KEY"] = "ENTER_YOUR_API_KEY_HERE"

rfm.init()
```

--------------------------------

### Predict Transactions within a Specific Window (PQL)

Source: https://kumo.ai/docs/pquery-structure

This example shows how to predict transaction values within a specific time range, from 10 days in the future up to 30 days in the future. The aggregation window is defined by the start and end parameters.

```PQL
PREDICT SUM(TRANSACTIONS.PRICE, 10, 30)
FOR EACH CUSTOMERS.CUSTOMER_ID
```

--------------------------------

### Example Data Splits Configuration

Source: https://kumo.ai/docs/reference/split

An example configuration showcasing various data splitting strategies available in Kumo AI, including ratio-based, date offset, random, temporal, and time range splits.

```text
split: ApproxDateOffsetSplit([0.8, 0.1, 0.1])
split: DateOffsetSplit([-360, -180], unit='days')
split: RandomSplit([0.7, 0.2, 0.1])
split: TemporalSplit([0.7, 0.2, 0.1])
split: TimeRangeSplit([("1994-01-01", "1997-01-01"), ("1997-01-01", "1998-01-01"), ("1998-01-01", "1999-01-01")])
```

--------------------------------

### PQL Entity Specification Example

Source: https://kumo.ai/docs/rfm/RFM-make-predictions

Shows how to specify the entity for which predictions will be made using the FOR clause in PQL. This example illustrates targeting predictions for each customer_id.

```python
FOR customer_id
```

--------------------------------

### PQL Entity Pinning Examples

Source: https://kumo.ai/docs/rfm/RFM-make-predictions

Provides examples of how to explicitly specify a single entity or a list of entities for which predictions should be generated using the FOR clause in PQL.

```python
FOR customer_id = 101
```

```python
FOR customer_id IN (101, 102,103,104)
```

--------------------------------

### Configure Kumo AI Compute Pool for Legacy AWS Users

Source: https://kumo.ai/docs/installing-kumo-on-spcs

This script is for legacy AWS users who utilize a single compute pool. It creates a GPU compute pool if one does not exist, grants usage permissions to the Kumo application, and then starts the application using the specified compute pool.

```sql
-- Modify the following if a different name was used to install the app.
SET KUMOAPP_NAME = 'KUMO';

SET KUMOAPP_COMPUTE = CONCAT(($KUMOAPP_NAME),'.PROCEDURES.START_APP');

-- Create a compute pool if none exist. Otherwise, use the name of an existing compute pool
-- in subsequent lines.
CREATE COMPUTE POOL IF NOT EXISTS KUMO_COMPUTE_PLANE_POOL_GPU_NV_M
    FOR APPLICATION IDENTIFIER($KUMOAPP_NAME)
    min_nodes = 1
    max_nodes = 1
    instance_family = 'GPU_NV_M'
    INITIALLY_SUSPENDED = TRUE
    AUTO_SUSPEND_SECS = 600;


GRANT USAGE ON COMPUTE POOL KUMO_COMPUTE_PLANE_POOL_GPU_NV_M to APPLICATION IDENTIFIER($KUMOAPP_NAME);

CALL IDENTIFIER($KUMOAPP_COMPUTE)('KUMO_COMPUTE_PLANE_POOL_GPU_NV_M', 'USER_SCHEMA', 'GPU_NV_M');

```

--------------------------------

### Initialize Kumo SDK

Source: https://kumo.ai/docs/examples/pzn-cold-start

Initializes the Kumo SDK client with the provided API endpoint and key. This is the first step before interacting with Kumo services.

```python
import kumoai as kumo

API_KEY = "YOUR_API_KEY"
kumo.init(url="https://<customer_id>.kumoai.cloud/api", api_key=API_KEY)
```

--------------------------------

### Create and Visualize a Local Graph

Source: https://kumo.ai/docs/rfm/RFM-quickstart

Creates a LocalGraph object from pandas DataFrames, automatically inferring relationships between tables. The graph can then be visualized using graphviz, which must be installed separately.

```python
import kumoai.experimental.rfm as rfm

graph = rfm.LocalGraph.from_data({
    "users": users_df,
    "items": items_df,
    "orders": orders_df,
})

# Inspect the graph - requires graphviz to be installed
graph.visualize()
```

--------------------------------

### Install KumoAI and Relbench Packages

Source: https://kumo.ai/docs/rfm/RFM-improve-model-performance

Installs the necessary KumoAI and Relbench Python packages using pip. This is a prerequisite for running the subsequent code examples.

```bash
!pip install kumoai --pre
!pip install relbench
```

--------------------------------

### Import Data and Setup Graph with KumoRFM (Python)

Source: https://kumo.ai/docs/rfm/RFM-batch-prediction

Imports necessary libraries (pandas, kumoai.experimental.rfm), loads user, item, and order data from S3 into pandas DataFrames, and initializes a KumoRFM graph using the loaded data. This is the initial setup for performing predictions.

```python
import kumoai.experimental.rfm as rfm
import pandas as pd

root = "s3://kumo-sdk-public/rfm-datasets/online-shopping"

df_users = pd.read_parquet(f"{root}/users.parquet")
df_items = pd.read_parquet(f"{root}/items.parquet")
df_orders = pd.read_parquet(f"{root}/orders.parquet")

graph = rfm.LocalGraph.from_data({
    'users': df_users,
    'items': df_items,
    'orders': df_orders,
})

model = rfm.KumoRFM(graph)
```

--------------------------------

### Retrieve Kumo Access URL (SQL)

Source: https://kumo.ai/docs/spcs/installing-kumo-on-spcs

This SQL command retrieves the access URL for the Kumo native application once it has been started. It calls a stored procedure to get the endpoint information.

```sql
SET KUMOAPP_NAME = 'KUMO';
SET KUMOAPP_STOP = CONCAT(($KUMOAPP_NAME),'.PROCEDURES.GET_END_POINT');
CALL IDENTIFIER($KUMOAPP_STOP)('USER_SCHEMA');
```

--------------------------------

### Temporal Recommendation Model Training (PQL)

Source: https://kumo.ai/docs/examples/pzn-cold-start

Trains a temporal recommendation model to provide personalized suggestions, dynamically handling cold start items using item features. This is suitable when there's a mix of new and existing items. It supports new entities and can use feature or fusion embedding modes.

```pql
PREDICT LIST_DISTINCT(orders.item_id, 0, 7) RANK TOP K
FOR EACH customers.customer_id

// If most items are new:
// module: link_prediction_embedding
// handle_new_entities: true
// target_embedding_mode: feature

// If a mix of old and new items:
// module: link_prediction_ranking
// handle_new_entities: false
// target_embedding_mode: fusion
```

--------------------------------

### Enable Cold-Start Handling in Kumo AI

Source: https://kumo.ai/docs/examples/pzn-related-items

This configuration option for Kumo AI models enables the system to handle new items (cold-start problem) by learning from item attributes when no prior purchase history is available. This is crucial for dynamic product catalogs.

```pql
handle_new_target_entities: true
```

--------------------------------

### Visualize a LocalGraph in Python

Source: https://kumo.ai/docs/rfm/RFM-graph-creation

This code snippet demonstrates how to visualize the structure of a LocalGraph. Visualization requires the 'graphviz' library to be installed. It provides a graphical representation of the tables and their connections.

```python
# Requires graphviz to be installed
graph.visualize();
```

--------------------------------

### Launch Kumo AI App using SQL

Source: https://kumo.ai/docs/spcs/installing-kumo-on-spcs

This SQL script launches the Kumo AI application. It requires user input for the warehouse to be used and optionally the app name if it differs from 'KUMO'. The role executing the script must have usage privileges on the specified warehouse. The command returns a URL to access the Kumo UI and may take 20-30 minutes to complete.

```sql
-----------------------------
---  USER INPUT REQUIRED  ---
-----------------------------
-- Note that KUMO_USER_ROLE should be granted usage on this warehouse.
SET USER_WAREHOUSE = '<ADD WAREHOUSE TO USE>';
-- Modify the following if a different name was used to install the app.
SET KUMOAPP_NAME = 'KUMO';

USE ROLE KUMO_USER_ROLE;

USE WAREHOUSE IDENTIFIER($USER_WAREHOUSE);
SET KUMOAPP_COMPUTE = CONCAT(($KUMOAPP_NAME),'.PROCEDURES.START_APP');
CALL IDENTIFIER($KUMOAPP_COMPUTE)('USER_SCHEMA');
```

--------------------------------

### Initialize Kumo SDK (SaaS/Databricks)

Source: https://kumo.ai/docs/reference/sdk

Initializes the Kumo SDK for SaaS and Databricks editions using a URL and API key. Ensure you have your API key and the correct URL for your Kumo platform.

```python
import kumoai as kumo
kumo.init(url="<url>/api", api_key="<api_key>")
```

--------------------------------

### Define and Train Predictive Model with Kumo SDK

Source: https://kumo.ai/docs/examples/pzn-cold-start

Defines a predictive query for item recommendations, validates it, suggests a model plan, and initiates the training process using the Kumo SDK. It prints the training metrics upon completion.

```python
pquery = kumo.PredictiveQuery(
    graph=graph,
    query="PREDICT LIST_DISTINCT(orders.item_id, 0, X, days) RANK TOP 50 FOR EACH customers.customer_id"
)
pquery.validate(verbose=True)

model_plan = pquery.suggest_model_plan()
trainer = kumo.Trainer(model_plan)
training_job = trainer.fit(
    graph=graph,
    train_table=pquery.generate_training_table(non_blocking=True),
    non_blocking=False,
)
print(f"Training metrics: {training_job.metrics()}")
```

--------------------------------

### Retrieve Kumo AI Access URL

Source: https://kumo.ai/docs/installing-kumo-on-spcs

This SQL command retrieves the current access URL for the Kumo native application after it has been started. This procedure is useful for quickly obtaining the URL without needing to restart or update the application.

```sql
SET KUMOAPP_NAME = 'KUMO';
SET KUMOAPP_STOP = CONCAT(($KUMOAPP_NAME),'.PROCEDURES.GET_END_POINT');
CALL IDENTIFIER($KUMOAPP_STOP)('USER_SCHEMA');

```

--------------------------------

### Remove and Re-add Links in a LocalGraph in Python

Source: https://kumo.ai/docs/rfm/RFM-graph-creation

This example illustrates how to remove an existing link between tables and then re-add it. This is useful for modifying the graph's structure dynamically. It assumes a 'graph' object and specifies source table, foreign key, and destination table.

```python
# Remove link:
graph.unlink(src_table="orders", fkey="user_id", dst_table="users")

# Re-add link:
graph.link(src_table="orders", fkey="user_id", dst_table="users");
```

--------------------------------

### Create Kumo Objects and Roles (SQL)

Source: https://kumo.ai/docs/installing-kumo-on-spcs

This SQL script sets up the foundational objects for Kumo AI, including databases, schemas, and roles. It requires user input for the warehouse name and optionally the Kumo application name. It grants necessary privileges like BIND SERVICE ENDPOINT and CREATE COMPUTE POOL to the Kumo application. It also defines network rules for egress traffic and creates an external access integration.

```sql
-----------------------------
---  USER INPUT REQUIRED  ---
-----------------------------
SET USER_WAREHOUSE = '<ADD WAREHOUSE TO USE>';

-- (Optional) Modify the following if a different name was used to install the app.
SET KUMOAPP_NAME = 'KUMO';
SET KUMOAPP_DB = 'KUMO_APP_DB';

SET KUMOAPP_USR = CONCAT(($KUMOAPP_NAME),'.APP_USER');

USE ROLE ACCOUNTADMIN;
USE WAREHOUSE IDENTIFIER($USER_WAREHOUSE);

-- Grant BIND SERVICE ENDPOINT privilege to the Kumo application to enable
-- network ingress and access to the Kumo UI for users of Kumo in your account.
-- Details of the privilige can be found here
-- https://other-docs.snowflake.com/LIMITEDACCESS/native-apps/na-spcs-consumer#set-up-access-to-network-objects
GRANT BIND SERVICE ENDPOINT ON ACCOUNT TO APPLICATION IDENTIFIER($KUMOAPP_NAME);

-- Grant CREATE COMPUTE POOL privilege to allow the Kumo application to create compute pools
-- This allows for efficient management of compute pools by Kumo.
GRANT CREATE COMPUTE POOL ON ACCOUNT TO APPLICATION IDENTIFIER($KUMOAPP_NAME);

-- Share events from the application with the provider, KUMO.AI
ALTER APPLICATION IDENTIFIER($KUMOAPP_NAME) SET AUTHORIZE_TELEMETRY_EVENT_SHARING=true;

-- Database and schema to hold network rules for the Kumo application along
-- with usage for the Kumo application.
CREATE DATABASE IF NOT EXISTS IDENTIFIER($KUMOAPP_DB);
USE DATABASE IDENTIFIER($KUMOAPP_DB);
CREATE SCHEMA IF NOT EXISTS KUMO_SCHEMA;
USE SCHEMA KUMO_SCHEMA;

---------------------
-- EGRESS RULES FOR KUMO
---------------------

CREATE OR REPLACE PROCEDURE CREATE_EGRESS_RULES(kumo_app_name varchar)
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION = '3.10'
PACKAGES = ('snowflake-snowpark-python')
HANDLER = 'create_egress_rule'
EXECUTE AS CALLER
AS
$$
import json


def create_egress_rule(session, kumo_app_name):
    df = session.sql(f"CALL {kumo_app_name}.PROCEDURES.GET_CONFIGURATION('KUMO_EXTERNAL_ACCESS_INTEGRATION')").collect()
    result = json.loads(df[0][0])
    if "payload" not in result or "host_ports" not in result["payload"] or len(result["payload"]["host_ports"]) == 0:
        raise RuntimeError(f'Error creating egress rules: {result}')

    hosts = result["payload"]["host_ports"]
    session.sql(f"""CREATE OR REPLACE NETWORK RULE kumo_network_rule
    MODE = EGRESS
    TYPE = HOST_PORT
    VALUE_LIST = ({', '.join(f"'{item}'" for item in hosts)})
  """).collect()

    session.sql("""CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION KUMO_EXTERNAL_ACCESS_INTEGRATION_APP
  ALLOWED_NETWORK_RULES = (kumo_network_rule)
  ENABLED = true""").collect()
    session.sql(f"GRANT USAGE ON INTEGRATION KUMO_EXTERNAL_ACCESS_INTEGRATION_APP TO APPLICATION {kumo_app_name}").collect()
    session.sql(f"CALL {kumo_app_name}.PROCEDURES.register_single_callback('KUMO_EXTERNAL_ACCESS_INTEGRATION', 'ADD', SYSTEM$REFERENCE('external_access_integration', 'KUMO_EXTERNAL_ACCESS_INTEGRATION_APP', 'PERSISTENT'))").collect()
    return f"Created egress rules and external access integration for {hosts}"
$$;

CALL CREATE_EGRESS_RULES($KUMOAPP_NAME);

------------------------

-- Create KUMO_USER_ROLE and grant necessary privileges. This role can be assigned to
-- any Snowflake user who is expected to use the Kumo application.
CREATE ROLE IF NOT EXISTS KUMO_USER_ROLE;

-- Grant the Application role to KUMO_USER_ROLE.
GRANT APPLICATION ROLE IDENTIFIER($KUMOAPP_USR) TO ROLE KUMO_USER_ROLE;

-- Grant USAGE on the database, schema and warehouse to KUMO_USER_ROLE.
GRANT USAGE ON DATABASE IDENTIFIER($KUMOAPP_DB) TO ROLE KUMO_USER_ROLE;
GRANT USAGE ON ALL SCHEMAS IN DATABASE IDENTIFIER($KUMOAPP_DB) TO ROLE KUMO_USER_ROLE;
GRANT USAGE ON WAREHOUSE IDENTIFIER($USER_WAREHOUSE) TO ROLE KUMO_USER_ROLE;

```

--------------------------------

### Temporal Recommendation Predictive Query (PQL)

Source: https://kumo.ai/docs/cold-start-recommendations-cold-start-items

Trains a personalized recommendation model per-customer. It supports two main cases: when the majority of items are new (cold start) and when there's a balanced mixture of cold and non-cold items. For cold start items, `link_prediction_embedding` with specific settings is recommended. For balanced mixtures, `link_prediction_ranking` is suggested.

```PQL
PREDICT LIST_DISTINCT(orders.item_id, 0, 7) RANK TOP K
FOR EACH customers.customer_id
// Case 1. Majority cold start items
// module: link_prediction_embedding
// handle_new_entities: true
// target_embedding_mode: feature

// Case 2. Balanced mixture of cold and non-cold items
// module: link_prediction_ranking
// handle_new_entities: false
// target_embedding_mode: fusion
```

--------------------------------

### SUM Operator Example in PQL

Source: https://kumo.ai/docs/reference/sum

Calculates the sum of loan amounts from day 0 (excluded) to day 30 (included). This operator is applicable to numerical columns and can be used in target or temporal filter fields. The default unit for start and end times is 'days'.

```text
SUM(LOAN.AMOUNT, 0, 30)
```

--------------------------------

### Import E-Commerce Data using Pandas

Source: https://kumo.ai/docs/rfm/RFM-quickstart

Imports an E-Commerce dataset from S3 into pandas DataFrames. This dataset includes user, item, and order information, which will be used to build a graph for analysis. Requires pandas and access to the specified S3 path.

```python
import pandas as pd

dataset_url = "s3://kumo-sdk-public/rfm-datasets/online-shopping"

users_df = pd.read_parquet(f"{dataset_url}/users.parquet")
items_df = pd.read_parquet(f"{dataset_url}/items.parquet")
orders_df = pd.read_parquet(f"{dataset_url}/orders.parquet")
```

--------------------------------

### Predictive Query Language (PQL) Lead Ranking Example

Source: https://kumo.ai/docs/examples/growthmarketing

An example of a Predictive Query Language (PQL) query to predict lead conversion within a specified timeframe. PQL supports inline filters, boolean expressions, and aggregation functions for flexible machine learning problem definition.

```sql
PREDICT COUNT(events.* WHERE events.type = 'conversion', 0, N, days ) > 0
FOR EACH leads.lead_id
WHERE COUNT(triggers.*,-1, 0, days) > 0
ASSUMING COUNT(events.* WHERE events.source = 'sales', 0, 1, days ) > 0
```

--------------------------------

### Initialize KumoRFM MCP Server

Source: https://kumo.ai/docs/rfm/RFM-build-agent-from-kumorfm-mcp

Initializes the KumoRFM MCP server using MCPServerStdio, setting its name, arguments, and environment variables including the KUMO_API_KEY. This is a crucial step for establishing communication with the KumoRFM service.

```python
server = MCPServerStdio(
    name="kumo_rfm",
    args=["-m", "kumo_rfm_mcp.server"],
    env={"KUMO_API_KEY": os.environ["KUMO_API_KEY"]},
)
```

--------------------------------

### Batch Prediction with KumoRFM (Python)

Source: https://kumo.ai/docs/rfm/RFM-batch-prediction

Enables scoring for the entire dataset or large batches by automatically handling batching and retries. This example collects all user IDs and uses `model.batch_mode` to perform predictions for each user, with options for batch size and retries.

```python
# collect all user IDs
indices = df_users["user_id"].tolist()

with model.batch_mode(batch_size="max", num_retries=1):
    result = model.predict(
        "PREDICT COUNT(orders.*, 0, 30, days) FOR EACH users.user_id",
        indices=indices
    )
```

--------------------------------

### Subgraph Node Example

Source: https://kumo.ai/docs/rfm/RFM-understand-explanations

Represents a single row within a table in a Kumo AI subgraph. It includes cell data with values and scores, and links to other related nodes. Node IDs are internal indices, with ID 0 always representing the seed entity.

```python
# Node 0 in 'orders' table (the seed entity) 
{
	'cells': {
		'date': {'value': '2018-09-20', 'score': 1.0},
	    'price': {'value': 8.98, 'score': 0.291},
		'sales_channel_id': {'value': 2, 'score': 0.222}
		},
	'links': {
		'item_id->items': {5},
		'user_id->users': {17}
	}
}
```

--------------------------------

### Create Kumo Objects and Roles (SQL)

Source: https://kumo.ai/docs/spcs/installing-kumo-on-spcs

This script sets up the foundational Kumo environment within your Snowflake account. It requires user input for the warehouse name and optionally the Kumo application name. The script creates a database, schema, network rule, external access integration, and grants necessary privileges to the Kumo application and a default user role (KUMO_USER_ROLE). It also includes a Python stored procedure to dynamically create egress rules based on Kumo's configuration.

```sql
-----------------------------
---  USER INPUT REQUIRED  ---
-----------------------------
SET USER_WAREHOUSE = '<ADD WAREHOUSE TO USE>';

-- (Optional) Modify the following if a different name was used to install the app.
SET KUMOAPP_NAME = 'KUMO';
SET KUMOAPP_DB = 'KUMO_APP_DB';

SET KUMOAPP_USR = CONCAT(($KUMOAPP_NAME),'.APP_USER');

USE ROLE ACCOUNTADMIN;
USE WAREHOUSE IDENTIFIER($USER_WAREHOUSE);

-- Grant BIND SERVICE ENDPOINT privilege to the Kumo application to enable
-- network ingress and access to the Kumo UI for users of Kumo in your account.
-- Details of the privilige can be found here
-- https://other-docs.snowflake.com/LIMITEDACCESS/native-apps/na-spcs-consumer#set-up-access-to-network-objects
GRANT BIND SERVICE ENDPOINT ON ACCOUNT TO APPLICATION IDENTIFIER($KUMOAPP_NAME);

-- Grant CREATE COMPUTE POOL privilege to allow the Kumo application to create compute pools
-- This allows for efficient management of compute pools by Kumo.
GRANT CREATE COMPUTE POOL ON ACCOUNT TO APPLICATION IDENTIFIER($KUMOAPP_NAME);

-- Share events from the application with the provider, KUMO.AI
ALTER APPLICATION IDENTIFIER($KUMOAPP_NAME) SET AUTHORIZE_TELEMETRY_EVENT_SHARING=true;

-- Database and schema to hold network rules for the Kumo application along
-- with usage for the Kumo application.
CREATE DATABASE IF NOT EXISTS IDENTIFIER($KUMOAPP_DB);
USE DATABASE IDENTIFIER($KUMOAPP_DB);
CREATE SCHEMA IF NOT EXISTS KUMO_SCHEMA;
USE SCHEMA KUMO_SCHEMA;

---------------------
-- EGRESS RULES FOR KUMO
---------------------

CREATE OR REPLACE PROCEDURE CREATE_EGRESS_RULES(kumo_app_name varchar)
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION = '3.10'
PACKAGES = ('snowflake-snowpark-python')
HANDLER = 'create_egress_rule'
EXECUTE AS CALLER
AS
$$
import json


def create_egress_rule(session, kumo_app_name):
  df = session.sql(f"CALL {kumo_app_name}.PROCEDURES.GET_CONFIGURATION('KUMO_EXTERNAL_ACCESS_INTEGRATION')").collect()
  result = json.loads(df[0][0])
  if "payload" not in result or "host_ports" not in result["payload"] or len(result["payload"]["host_ports"]) == 0:
    raise RuntimeError(f'Error creating egress rules: {result}')

  hosts = result["payload"]["host_ports"]
  session.sql(f"""CREATE OR REPLACE NETWORK RULE kumo_network_rule
      MODE = EGRESS
      TYPE = HOST_PORT
      VALUE_LIST = ({', '.join(f"'{item}'" for item in hosts)})
    """).collect()

  session.sql("""CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION KUMO_EXTERNAL_ACCESS_INTEGRATION_APP
    ALLOWED_NETWORK_RULES = (kumo_network_rule)
    ENABLED = true""").collect()
  session.sql(f"GRANT USAGE ON INTEGRATION KUMO_EXTERNAL_ACCESS_INTEGRATION_APP TO APPLICATION {kumo_app_name}").collect()
  session.sql(f"CALL {kumo_app_name}.PROCEDURES.register_single_callback('KUMO_EXTERNAL_ACCESS_INTEGRATION', 'ADD', SYSTEM$REFERENCE('external_access_integration', 'KUMO_EXTERNAL_ACCESS_INTEGRATION_APP', 'PERSISTENT'))").collect()
  return f"Created egress rules and external access integration for {hosts}"
$$
;

CALL CREATE_EGRESS_RULES($KUMOAPP_NAME);

------------------------

-- Create KUMO_USER_ROLE and grant necessary privileges. This role can be assigned to
-- any Snowflake user who is expected to use the Kumo application.
CREATE ROLE IF NOT EXISTS KUMO_USER_ROLE;

-- Grant the Application role to KUMO_USER_ROLE.
GRANT APPLICATION ROLE IDENTIFIER($KUMOAPP_USR) TO ROLE KUMO_USER_ROLE;

-- Grant USAGE on the database, schema and warehouse to KUMO_USER_ROLE.
GRANT USAGE ON DATABASE IDENTIFIER($KUMOAPP_DB) TO ROLE KUMO_USER_ROLE;
GRANT USAGE ON ALL SCHEMAS IN DATABASE IDENTIFIER($KUMOAPP_DB) TO ROLE KUMO_USER_ROLE;
GRANT USAGE ON WAREHOUSE IDENTIFIER($USER_WAREHOUSE) TO ROLE KUMO_USER_ROLE;

```

--------------------------------

### Make Predictive Queries with KumoRFM

Source: https://kumo.ai/docs/rfm/RFM-quickstart

Utilizes the KumoRFM model to make various predictive queries on the graph. This includes forecasting product demand, predicting customer churn, recommending items, and imputing missing values. Each prediction requires a specific query string.

```python
import kumoai.experimental.rfm as rfm

# Assuming graph is already created and initialized
model = rfm.KumoRFM(graph)

# Forecast 30-day product demand
query1 = "PREDICT SUM(orders.price, 0, 30, days) FOR items.item_id=1"
result1 = model.predict(query1)
display(result1)

# Predict customer churn
query2 = "PREDICT COUNT(orders.*, 0, 90, days)=0 FOR users.user_id IN (42, 123)"
result2 = model.predict(query2)
display(result2)

# Item recommendation
query3 = "PREDICT LIST_DISTINCT(orders.item_id, 0, 30, days) RANK TOP 10 FOR users.user_id=123"
result3 = model.predict(query3)
display(result3)

# Missing value imputation
query4 = "PREDICT users.age FOR users.user_id=8"
result4 = model.predict(query4)
display(result4)
```

--------------------------------

### Batch Recommendations Storage Format

Source: https://kumo.ai/docs/examples/pzn-related-items

This SQL-like example illustrates how to store precomputed item-to-item recommendations in a key-value store. The format maps a left-hand side item ID to a list of its recommended right-hand side item IDs.

```sql
ITEM_ID_1 : {ITEM_ID_A, ITEM_ID_B, ..., ITEM_ID_N} 
# LHS_ITEM : {Recommended items for LHS_ITEM}
```
Tokens:
Raw