8000 GitHub - vineetp6/pgai: Helper functions for AI workflows
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content
/ pgai Public
forked from timescale/pgai

Helper functions for AI workflows

License

PostgreSQL license
1 star 251 forks Branches Tags Activity
Star
Notifications You must be signed in to change notification settings

vineetp6/pgai

BranchesTags
 
 

Repository files navigation

pgai

pgai brings AI workflows to your PostgreSQL database

Discord Try Timescale for free

pgai simplifies the process of building search, and Retrieval Augmented Generation(RAG) AI applications with PostgreSQL.

pgai brings embedding and generation AI models closer to the database. With pgai, you can now do the following directly from within PostgreSQL in a SQL query:

Here's how to get started with pgai:

Learn more about pgai: To learn more about the pgai extension and why we built it, read this blog post pgai: Giving PostgreSQL Developers AI Engineering Superpowers.

pgai Prerequisites

Before you start working with pgai, you need:

Installation

The fastest ways to run PostgreSQL with TimescaleDB and pgai are:

Enable pgai in a pre-built Docker container

  1. Run the TimescaleDB Docker image.

  2. Connect to your database:

    psql -d "postgres://<username>:<password>@<host>:<port>/<database-name>"

  3. Create the pgai extension:

    CREATE EXTENSION IF NOT EXISTS ai CASCADE;

    The CASCADE automatically installs the plpython3u and pgvector dependencies.

You now Provide your API key to pgai and Try out the AI models.

Installing from source

You can install pgai from source and install it in an existing PostgreSQL server

  1. Install the extension and dependencies

    make install

  2. Connect to your database:

    psql -d "postgres://<username>:<password>@<host>:<port>/<database-name>"

  3. Create the pgai extension:

    CREATE EXTENSION IF NOT EXISTS ai CASCADE;

    The CASCADE automatically installs pgvector and plpython3u.

You now Provide your API key to pgai and Try out the AI models.

Enable pgai in a Timescale Cloud service

To enable pgai:

  1. Create a new Timescale Service.

    If you want to use an existing service, pgai is added as an available extension on the first maintenance window after the pgai release date.

  2. Connect to your Timescale service:

    psql -d "postgres://<username>:<password>@<host>:<port>/<database-name>"

  3. Create the pgai extension:

    CREATE EXTENSION IF NOT EXISTS ai CASCADE;

    CASCADE automatically installs the plpython3u and pgvector dependencies.

You now Provide your API key to pgai and Try out the AI models.

Provide your API key to pgai

Most pgai functions require an OpenAI API key.

Handle API keys using pgai from psql

The api key is an optional parameter to pgai functions. You either:

Run AI queries by passing your API key implicitly as a session parameter

To use a session level parameter when connecting to your database with psql to run your AI queries:

  1. Set your OpenAI key as an environment variable in your shell:

    export OPENAI_API_KEY="this-is-my-super-secret-api-key-dont-tell"

  2. Use the session level parameter when you connect to your database:

    PGOPTIONS="-c ai.openai_api_key=$OPENAI_API_KEY" psql -d "postgres://<username>:<password>@<host>:<port>/<database-name>"

  3. Run your AI query:

    ai.openai_api_key is set for the duration of your psql session, you do not need to specify it for pgai functions.

    SELECT * 
FROM openai_list_models()
ORDER BY created DESC
;

Run AI queries by passing your API key explicitly as a function argument

  1. Set your OpenAI key as an environment variable in your shell:

    export OPENAI_API_KEY="this-is-my-super-secret-api-key-dont-tell"

  2. Connect to your database and set your api key as a psql variable.

    psql -d "postgres://<username>:<password>@<host>:<port>/<database-name>" -v openai_api_key=$OPENAI_API_KEY

    Your API key is now available as a psql variable named openai_api_key in your psql session.

    You can also log into the database, then set openai_api_key using the \getenv metacommand:

     \getenv openai_api_key OPENAI_API_KEY

  3. Pass your API key to your parameterized query:

    SELECT * 
FROM openai_list_models(_api_key=>$1)
ORDER BY created DESC
\bind :openai_api_key
\g

    Use \bind to pass the value of openai_api_key to the parameterized query.

    The \bind metacommand is available in psql version 16+.

Handle API keys using pgai from python

  1. In your Python environment, include the dotenv and postgres driver packages:

    pip install python-dotenv
pip install psycopg2-binary

  2. Set your OpenAI key in a .env file or as an environment variable:

    OPENAI_API_KEY="this-is-my-super-secret-api-key-dont-tell"
DB_URL="your connection string"

  3. Pass your API key as a parameter to your queries:

    import os
from dotenv import load_dotenv

load_d
8000
otenv()

OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
DB_URL = os.environ["DB_URL"]

import psycopg2

with psycopg2.connect(DB_URL) as conn:
    with conn.cursor() as cur:
        # pass the API key as a parameter to the query. don't use string manipulations
        cur.execute("SELECT * FROM openai_list_models(_api_key=>%s) ORDER BY created DESC", (OPENAI_API_KEY,))
        records = cur.fetchall()

    Do not use string manipulation to embed the key as a literal in the SQL query.

Usage

This section shows you how to use AI directly from your database using SQL.

List models

List the models supported by OpenAI functions in pgai.

SELECT * 
FROM openai_list_models()
ORDER BY created DESC
;

The data returned looks like:

             id              |        created         |    owned_by     
-----------------------------+------------------------+-----------------
 gpt-4o-test-shared          | 2024-05-20 13:06:56-05 | system
 gpt-4o-2024-05-13           | 2024-05-10 14:08:52-05 | system
 gpt-4o                      | 2024-05-10 13:50:49-05 | system
 gpt-4-turbo-2024-04-09      | 2024-04-08 13:41:17-05 | system
 gpt-4-turbo                 | 2024-04-05 18:57:21-05 | system
 ...
(N rows)

Tokenize

To encode content and count the number of tokens returned:

  • Encode content into an array of tokens.

    SELECT openai_tokenize
( 'text-embedding-ada-002'
, 'Timescale is Postgres made Powerful'
);

    The data returned looks like:

                openai_tokenize             
----------------------------------------
 {19422,2296,374,3962,18297,1903,75458}
(1 row)

  • Count the number of tokens generated:

    SELECT array_length
( openai_tokenize
  ( 'text-embedding-ada-002'
  , 'Timescale is Postgres made Powerful'
  )
, 1
);

    The data returned looks like:

     array_length 
--------------
            7
(1 row)

Detokenize

Turn tokenized content into natural language:

SELECT openai_detokenize('text-embedding-ada-002', array[1820,25977,46840,23874,389,264,2579,58466]);

The data returned looks like:

             openai_detokenize              
--------------------------------------------
 the purple elephant sits on a red mushroom
(1 row)

Embed

Generate embeddings using a specified model.

  • Request an embedding using a specific model.

    SELECT openai_embed
( 'text-embedding-ada-002'
, 'the purple elephant sits on a red mushroom'
);

    The data returned looks like:

                          openai_embed                      
--------------------------------------------------------
 [0.005978798,-0.020522336,...-0.0022857306,-0.023699166]
(1 row)

  • Specify the number of dimensions you want in the returned embedding:

    SELECT openai_embed
( 'text-embedding-ada-002'
, 'the purple elephant sits on a red mushroom'
, _dimensions=>768
);

    This only works for certain models.

  • Pass a user identifier.

    SELECT openai_embed
( 'text-embedding-ada-002'
, 'the purple elephant sits on a red mushroom'
, _user=>'bac1aaf7-4460-42d3-bba5-2957b057f4a5'
);

  • Pass an array of text inputs.

    SELECT openai_embed
( 'text-embedding-ada-002'
, array['Timescale is Postgres made Powerful', 'the purple elephant sits on a red mushroom']
);

  • Provide tokenized input.

    select openai_embed
( 'text-embedding-ada-002'
, array[1820,25977,46840,23874,389,264,2579,58466]
);

Chat complete

Generate text or complete a chat:

  • Have an LLM generate text from a prompt:

    -- the following two metacommands cause the raw query results to be printed
-- without any decoration
\pset tuples_only on
\pset format unaligned

SELECT jsonb_pretty
(
  openai_chat_complete
  ( 'gpt-4o'
  , jsonb_build_array
    ( jsonb_build_object('role', 'system', 'content', 'you are a helpful assistant')
    , jsonb_build_object('role', 'user', 'content', 'what is the typical weather like in Alabama in June')
    )
  )
);

    The data returned looks like:

    {
    "id": "chatcmpl-9RgehyQ0aydAkQajrN6Oe0lepERKC",
    "model": "gpt-4o-2024-05-13",
    "usage": {
        "total_tokens": 332,
        "prompt_tokens": 26,
        "completion_tokens": 306
    },
    "object": "chat.completion",
    "choices": [
        {
            "index": 0,
            "message": {
                "role": "assistant",
                "content": "In Alabama, June typically ushers in the summer season with warm to hot temperatures and relatively high humidity. Here’s a general overview of what you can expect:\n\n1. **Temperature**: \n   - Average daytime highs usually range from the mid-80s to low 90s Fahrenheit (around 29-35°C).\n   - Nighttime temperatures often fall to the mid-60s to mid-70s Fahrenheit (18-24°C).\n\n2. **Humidity**:\n   - Humidity levels can be quite high, making the temperatures feel even warmer. The mix of heat and humidity can lead to a muggy atmosphere.\n\n3. **Rainfall**:\n   - June is part of the wet season for Alabama, so you can expect a fair amount of rainfall. Thunderstorms are relatively common, often in the afternoons and evenings.\n   - The precipitation can be sporadic, with sudden downpours that can clear up quickly.\n\n4. **Sunshine**:\n   - There are plenty of sunny days, though the sunshine can be intense. Ultraviolet (UV) levels are high, so sun protection is important.\n\n5. **Overall Climate**:\n   - Generally, the climate in Alabama in June is characterized by a typical Southeastern U.S. summer: hot, humid, and occasionally stormy. \n\nIf you’re planning a visit or activities in Alabama during June, it’s a good idea to stay hydrated, wear light clothing, and keep an eye on the weather forecast for any potential thunderstorms."
            },
            "logprobs": null,
            "finish_reason": "stop"
        }
    ],
    "created": 1716385851,
    "system_fingerprint": "fp_729ea513f7"
}

  • Return the content as text from a specific message in the choices array.

    openai_chat_complete returns a jsonb object containing the response from the API. You can use jsonb operators and functions to manipulate the object returned. For example, the following query returns the content as text from the first message in the choices array.

    -- the following two metacommands cause the raw query results to be printed
-- without any decoration
\pset tuples_only on
\pset format unaligned

select openai_chat_complete
( 'gpt-4o'
, jsonb_build_array
  ( jsonb_build_object('role', 'system', 'content', 'you are a helpful assistant')
  , jsonb_build_object('role', 'user', 'content', 'what is the typical weather like in Alabama in June')
  )
)->'choices'->0->'message'->>'content'
;

    The data returned looks like:

    In June, Alabama generally experiences warm to hot weather as it transitions into summer. Typical conditions include:

1. **Temperatures**: Daytime highs usually range from the mid-80s to low 90s Fahrenheit (around 29-34°C). Nighttime lows typically range from the mid-60s to low 70s Fahrenheit (around 18-23°C).

2. **Humidity**: June tends to be quite humid, which can make the temperatures feel even warmer. High humidity levels are characteristic of Alabama summers.

3. **Precipitation**: June is part of the wetter season in Alabama, with regular afternoon thunderstorms being common. Rainfall can vary, but you can expect an average of about 4 to 5 inches (around 100-125 mm) of rain for the month.

4. **Sunshine**: There are usually plenty of sunny days, although the frequent thunderstorms can lead to overcast skies at times.

Overall, if you're planning to visit Alabama in June, be prepared for hot and humid conditions, and keep an umbrella or rain jacket handy for those afternoon storms.

Moderate

Check if content is classified as potentially harmful:

-- the following two metacommands cause the raw query results to be printed
-- without any decoration
\pset tuples_only on
\pset format unaligned

select jsonb_pretty
(
  openai_moderate
  ( 'text-moderation-stable'
  , 'I want to kill them.'
  )
);

The data returned looks like:

{
    "id": "modr-9RsN6qZWoZYm1AK4mtrKuEjfOcMWp",
    "model": "text-moderation-007",
    "results": [
        {
            "flagged": true,
            "categories": {
                "hate": false,
                "sexual": false,
                "violence": true,
                "self-harm": false,
                "self_harm": false,
                "harassment": true,
                "sexual/minors": false,
                "sexual_minors": false,
                "hate/threatening": false,
                "hate_threatening": false,
                "self-harm/intent": false,
                "self_harm_intent": false,
                "violence/graphic": false,
                "violence_graphic": false,
                "harassment/threatening": true,
                "harassment_threatening": true,
                "self-harm/instructions": false,
                "self_harm_instructions": false
            },
            "category_scores": {
                "hate": 0.2324090600013733,
                "sexual": 0.00001205232911161147,
                "violence": 0.997192919254303,
                "self-harm": 0.0000023696395601291442,
                "self_harm": 0.0000023696395601291442,
                "harassment": 0.5278584957122803,
                "sexual/minors": 0.00000007506431387582779,
                "sexual_minors": 0.00000007506431387582779,
                "hate/threatening": 0.024183575063943863,
                "hate_threatening": 0.024183575063943863,
                "self-harm/intent": 0.0000017161115692942985,
                "self_harm_intent": 0.0000017161115692942985,
                "violence/graphic": 0.00003399916022317484,
                "violence_graphic": 0.00003399916022317484,
                "harassment/threatening": 0.5712487697601318,
                "harassment_threatening": 0.5712487697601318,
                "self-harm/instructions": 0.000000001132860139030356,
                "self_harm_instructions": 0.000000001132860139030356
            }
        }
    ]
}

Advanced examples

For more advanced usage, the Advanced examples use pgai to embed, moderate, and summarize a git commit history.

Get involved

pgai is still at an early stage. Now is a great time to help shape the direction of this project; we are currently deciding priorities. Have a look at the list of features we're thinking of working on. Feel free to comment, expand the list, or hop on the Discussions forum.

To get started, take a look at how to contribute and how to set up a dev/test environment.

About Timescale

Timescale is a PostgreSQL database company. To learn more visit the timescale.com.

Timescale Cloud is a high-performance, developer focused, cloud platform that provides PostgreSQL services for the most demanding AI, time-series, analytics, and event workloads. Timescale Cloud is ideal for production applications and provides high availability, streaming backups, upgrades over time, roles and permissions, and great security.

About

Helper functions for AI workflows

Resources

Readme

License

PostgreSQL license
Activity

Stars

1 star

Watchers

0 watching

Forks

0 forks
Report repository

Releases

No releases published

Packages

No packages published

Languages

  • Makefile 49.6%
  • Shell 34.6%
  • Dockerfile 15.8%
0