Working with Python UDFs

import typing
import numpy as np
import pandas as pd
from singlestoredb.functions import udf, Table
import numpy.typing as npt

# Define output schema using NamedTuple
class SalesOutput(typing.NamedTuple):
    total_sales: float
    average_sales: float
    category: str

@udf(returns=SalesOutput)
async def vector_sales(
    units_sold: npt.NDArray[np.float64],  # Vector of units sold
    unit_price: npt.NDArray[np.float64],  # Vector of unit prices
    category: npt.NDArray[np.str_]          # Vector of categories
) -> Table[pd.DataFrame]:
    """
    Calculate total and average sales per row with category.
    
    Parameters
    ----------
    units_sold : np.ndarray[np.float64]
        Number of units sold
    unit_price : np.ndarray[np.float64]
        Price per unit
    category : np.ndarray[str]
        Category name for each row

    Returns
    -------
    pd.DataFrame
        Table with columns:
        - total_sales
        - average_sales
        - category
    """
    # Compute vectorized total and average
    total_sales = units_sold * unit_price
    average_sales = total_sales / np.maximum(units_sold, 1)  # Avoid divide by zero
    
    # Return as a DataFrame with named columns
    df = pd.DataFrame({
        'total_sales': total_sales,
        'average_sales': average_sales,
        'category': category
    })
    
    return Table(df)

# Start Python UDF server

import singlestoredb.apps as apps
connection_info = await apps.run_udf_app()

CREATE OR REPLACE EXTERNAL FUNCTION `vector_sales`(
    `units_sold` DOUBLE NOT NULL,
    `unit_price` DOUBLE NOT NULL,
    `category` TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL
)
RETURNS TABLE(
    `total_sales` DOUBLE NOT NULL,
    `average_sales` DOUBLE NOT NULL,
    `category` TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL
)
DEFINER = '<function_owner>'@'%'
AS MANAGED SERVICE '<managed_service_url>'
FORMAT ROWDAT_1;

import base64
import sys

import httpx
import numpy as np
import numpy.typing as npt
from singlestoredb.functions import udf

MODEL_NAME = 'name-of-embedding-service'
MODEL_URL = 'url-of-embedding-server'
TOKEN = 'api-token-of-embedding-service'

HEADERS = {
    'accept': 'application/json',
    'Content-Type': 'application/json',
    'Authorization': f'Bearer {TOKEN}',
}


@udf
async def mixedbread_embeddings(
    text: npt.NDArray[np.str_],
) -> npt.NDArray[np.bytes_]:
    async with httpx.AsyncClient() as client:
        res = await client.post(
            MODEL_URL,
            headers=HEADERS,
            json=dict(
                model=MODEL_NAME,
                input=text.tolist(),
                encoding_format='base64',
            ),
        )

        if res.status_code >= 400:
            print(res.content.decode('utf8'), file=sys.stderr)
            raise RuntimeError(res.content.decode('utf8'))

        return np.array(
            [base64.b64decode(x['embedding']) for x in res.json()['data']],
            dtype=object,
        )


# Start Python UDF server.
import singlestoredb.apps as apps

connection_info = await apps.run_udf_app()

CREATE OR REPLACE EXTERNAL FUNCTION `mixedbread_embeddings`(
    `text` TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL
)
RETURNS BLOB NOT NULL
DEFINER = '<function_owner>'@'%'
AS MANAGED SERVICE '<managed_service_url>'
FORMAT ROWDAT_1;

import os
from typing import List
from pydantic import BaseModel, Field
from singlestoredb.functions import udf, Table
import instructor
from openai import OpenAI

# Create OpenAI client via instructor wrapper
client = instructor.from_openai(OpenAI(api_key=os.getenv("OPENAI_API_KEY")))

# Define output schema
class Synonym(BaseModel):
    word: str = Field(description='Synonym of the given word')
    score: float = Field(
        description='Closeness score from 0.0 to 1.0'
    )

# Define the TVF
@udf
async def get_synonyms(word: str) -> Table[List[Synonym]]:
    """Return a list of synonyms of the given word and a score."""
    return Table(client.create(
        model='gpt-4o-mini',
        messages=[
            dict(role='system', content='You are a helpful assistant'),
            dict(
                role='user',
                content=f'''
                    * Get a list of synonyms of the word "{word}"
                    * Limit the number of results to 10
                '''
            )
        ],
        response_model=List[Synonym],
        max_retries=0,
    ))

# Start Python UDF server

import singlestoredb.apps as apps
connection_info = await apps.run_udf_app()

CREATE OR REPLACE EXTERNAL FUNCTION `get_synonyms`(
    `word` TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL
)
RETURNS TABLE(
    `word` TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL,
    `score` DOUBLE NOT NULL
)
DEFINER = '<function_owner>'@'%'
AS MANAGED SERVICE '<managed_service_url>'
FORMAT ROWDAT_1;

SELECT * FROM get_synonyms('danger');

                    word   score
0                 hazard   0.90
1                   risk   0.85
2                 threat   0.88
3                  peril   0.90
4               jeopardy   0.87
5            threatening   0.80
6  threatening situation   0.70
7             threatened   0.75
8           insecurities   0.60
9                 unsafe   0.65

import numpy as np
import pandas as pd
import numpy.typing as npt
import singlestoredb.notebook as nb
from tensorflow.keras.models import load_model
from singlestoredb.functions import udf

# Download model from Stage
nb.stage.download_file('my_model.keras', local_path='my_model.keras')

# Load the model
model = load_model('my_model.keras')

@udf
async def keras_score(
    param_1: npt.NDArray[np.float32],
    param_2: npt.NDArray[np.float32],
    param_3: npt.NDArray[np.float32],
) -> npt.NDArray[np.float32]:
    """
    Score rows using a Keras model.

    Parameters
    ----------
    param_1 : np.ndarray[np.float32]
        First input parameter.
    param_2 : np.ndarray[np.float32]
        Second input parameter.
    param_3 : np.ndarray[np.float32]
        Third input parameter.

    Returns
    -------
    np.ndarray[np.float32]
        Predicted scores from the model.
    """
    X_test = pd.DataFrame({
        'param_1': param_1,
        'param_2': param_2,
        'param_3': param_3
    })
    return model.predict(X_test).reshape((-1,))

# Start Python UDF server

import singlestoredb.apps as apps
connection_info = await apps.run_udf_app()