Skip to content

tensorchord/pgvecto.rs-py

Repository files navigation

PGVecto.rs support for Python

discord invitation link trackgit-views trackgit-views trackgit-views

Usage

Install from PyPI:

pip install pgvecto_rs

See the usage of SDK

Or use it as an extension of postgres clients:

Requirements

To initialize a pgvecto.rs instance, you can run our official image by Quick start:

You can get the latest tags from the Release page. For example, it might be:

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.2.1

SDK

Our SDK is designed to use the pgvecto.rs out-of-box. You can exploit the power of pgvecto.rs to do similarity search or retrieve with filters, without writing any SQL code.

Install dependencies:

pip install "pgvecto_rs[sdk]"

A minimal example:

from pgvecto_rs.sdk import PGVectoRs, Record

# Create a client
client = PGVectoRs(
    db_url="postgresql+psycopg://postgres:mysecretpassword@localhost:5432/postgres",
    table_name="example",
    dimension=3,
)

try:
    # Add some records
    client.add_records(
        [
            Record.from_text("hello 1", [1, 2, 3]),
            Record.from_text("hello 2", [1, 2, 4]),
        ]
    )

    # Search with default operator (sqrt_euclid).
    # The results is sorted by distance
    for rec, dis in client.search([1, 2, 5]):
        print(rec.text)
        print(dis)
finally:
    # Clean up (i.e. drop the table)
    client.drop()

Output:

hello 2
1.0
hello 1
4.0

See examples/sdk_example.py and tests/test_sdk.py for more examples.

SQLAlchemy

Install dependencies:

pip install "pgvecto_rs[sqlalchemy]"

Then write your code. See examples/sqlalchemy_example.py and tests/test_sqlalchemy.py for example.

All the operators include:

  • squared_euclidean_distance
  • negative_dot_product_distance
  • negative_cosine_distance

psycopg3

Install dependencies:

pip install "pgvecto_rs[psycopg3]"

Then write your code. See examples/psycopg_example.py and tests/test_psycopg.py for example.

Development

This package is managed by PDM.

Set up things:

pdm venv create
pdm use # select the venv inside the project path
pdm sync

Run lint:

pdm run format
pdm run fix
pdm run check

Run test in current environment:

pdm run test

Test

Tox is used to test the package locally.

Run test in all environment:

tox run