Use AI_CreateModel to Register AI Models for In-Database Inference - PolarDB

Registers the metadata of an external AI model in PolarDB for PostgreSQL, making the model available for use with the polar_ai extension.

Before you begin

Enable the Polar_AI extension in your PolarDB for PostgreSQL cluster.
Prepare the endpoint and API key for the AI model you want to register.
Define the input and output transform functions for the model. See Input and output functions.

Syntax

SELECT polar_ai.AI_CreateModel(
    model_id,
    model_url,
    model_provider,
    model_type,
    model_name,
    model_config,
    model_headers_fn,
    model_in_transform_fn,
    model_out_transform_fn
);

Parameters

Parameter	Required	Description
`model_id`	Yes	A unique identifier for the AI model. Must differ from `model_name`. Cannot start with an underscore (`_`), which is reserved for built-in models created by the `Polar_AI` extension.
`model_url`	Yes	The URL used to access the AI model. Supports HTTP, HTTPS, and FILE protocols. For models deployed on Platform for AI (PAI), use the VPC endpoint from the model's call information, with the invocation method appended (for example, `/v1/completions`).
`model_provider`	No	The provider of the AI model. Examples: `AWS`, `Alibaba`, `Baidu`, `Tencent`.
`model_type`	No	The type of the AI model. Examples: `LSTM`, `GRU`.
`model_name`	Yes	The name of the AI model. Example: `text-embedding-v2`.
`model_config`	Yes	The authentication configuration for the AI model, in JSON format. Only token-based authentication is supported. Format: `{"author_type": "token", "token": "<YOUR_API_KEY>"}`. The API key is encrypted when stored.
`model_headers_fn`	No	A function that constructs custom request headers, returning a value of type JSONB. Omit this parameter if the model has no special header requirements. Default: `NULL`.
`model_in_transform_fn`	Yes	A function that converts user input into the JSONB-formatted HTTP request body required by the AI model.
`model_out_transform_fn`	Yes	A function that parses the model's JSON response and returns the required output in JSONB format.

AI_CreateModel only registers model metadata. It does not call or execute the model. To call a registered model, use AI_CallModel.

Return values

Returns the registration result as a table row with the following fields:

Field	Description
`model_seq`	The sequence number assigned to the AI model.
`model_schema`	The workspace (schema) to which the AI model belongs.
`model_id`	The unique identifier of the registered AI model.
`model_qname`	The qualified name of the AI model.
`created`	Whether registration succeeded. `t`: succeeded. `f`: failed.

Input and output functions

AI_CreateModel requires two transform functions to handle the model's request and response format.

Input function

The input function converts user-defined input into the JSONB-formatted HTTP request body that the AI model expects. For a completions API, the request body must contain model and prompt fields:

{
  "model": "DeepSeek-R1-Distill-Qwen-7B",
  "prompt": "Hello!"
}

Define the input function:

CREATE OR REPLACE FUNCTION ai_text_in_fn(model_name text, content text)
    RETURNS jsonb
    LANGUAGE plpgsql
    AS $function$
    BEGIN
        RETURN ('{"model": "' || model_name || '","prompt":"' || content || '"}')::jsonb;
    END;
    $function$;

Output function

The output function extracts the relevant field from the model's JSON response. For a completions API, the full response looks like this:

{
    "id": "8e44xxxx",
    "object": "text_completion",
    "created": 1744355891,
    "model": "DeepSeek-R1-Distill-Qwen-7B",
    "choices": [
        {
            "index": 0,
            "text": " I have a\n\n\n</think>\n\nHello! How can I assist you today? ",
            "logprobs": null,
            "finish_reason": "stop",
            "matched_stop": 151643
        }
    ],
    "usage": {
        "prompt_tokens": 3,
        "total_tokens": 21,
        "completion_tokens": 18,
        "prompt_tokens_details": null
    }
}

Define the output function to extract choices[0].text, returning a JSONB value:

CREATE OR REPLACE FUNCTION ai_text_out_fn(model_id text, response_json jsonb)
    RETURNS jsonb
    AS $$
    SELECT (response_json->'choices'->0->'text')::jsonb
    $$ LANGUAGE sql IMMUTABLE;

Examples

The following example registers the DeepSeek-R1-Distill-Qwen-7B model deployed on PAI.

Step 1: Create the input function.

CREATE OR REPLACE FUNCTION ai_text_in_fn(model_name text, content text)
    RETURNS jsonb
    LANGUAGE plpgsql
    AS $function$
    BEGIN
        RETURN ('{"model": "' || model_name || '","prompt":"' || content || '"}')::jsonb;
    END;
    $function$;

Step 2: Create the output function.

CREATE OR REPLACE FUNCTION ai_text_out_fn(model_id text, response_json jsonb)
    RETURNS jsonb
    AS $$
    SELECT (response_json->'choices'->0->'text')::jsonb
    $$ LANGUAGE sql IMMUTABLE;

Step 3: Register the model.

Replace <EAS Endpoint> with the VPC endpoint from the PAI model's call information, and <EAS Token> with the corresponding token.

SELECT polar_ai.AI_CreateModel(
    'my_text_pai_model',
    '<EAS Endpoint>/v1/completions',
    'Alibaba',
    'EAS LLM',
    'DeepSeek-R1-Distill-Qwen-32B',
    '{"author_type": "token", "token": "<EAS Token>"}',
    NULL,
    'ai_text_in_fn'::regproc,
    'ai_text_out_fn'::regproc
);

Expected output:

     ai_createmodel
-------------------------------------------------------------
 (1,my_text_pai_model,polar_ai,DeepSeek-R1-Distill-Qwen-7B,t)

Step 4 (optional): Verify the registration.

SELECT * FROM polar_ai._ai_models;

Sample output:

 model_seq |       model_id        | model_url | model_provider |          model_type          |      model_name      |                          model_config                           | model_headers_fn |         model_in_transform_fn         |          model_out_transform_fn
-----------+-----------------------+-----------+----------------+------------------------------+----------------------+-----------------------------------------------------------------+------------------+---------------------------------------+------------------------------------------
         1 | my_text_embedding_model | https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding | Alibaba | General-purpose text embedding model | text-embedding-v2 | {"token": "20E0BE9E5438xxxxxxxxxxxxxxxxxxxx", "author_type": "token"} | - | ai_text_embedding_in_fn(text,text) | ai_text_embedding_out_fn(text,jsonb)

Related operations

AI_CallModel: Calls a registered AI model.
View all registered models: SELECT * FROM polar_ai._ai_models;
View built-in models created by the Polar_AI extension: SELECT * FROM polar_ai._ai_models WHERE model_id LIKE '\_%';