pgvector examples

Author: dariusz-przada

postgres/Makefile

@@ -12,7 +12,7 @@ ifndef UV_VERSION
 	UV_VERSION := 0.7.21
 endif
 
-.PHONY: uv_check venv sync lock update format lint migrations example
+.PHONY: uv_check venv sync lock update format lint migrations example example-index example-generate example-delete example-search database database-stop
 
 
 # Check installed UV version and install if needed
@@ -81,3 +81,32 @@ migrations:
 
 example:
 	@python -B -m example
+
+example-index:
+	@python -B -m example.vector
+
+example-generate:
+	@python -B -m example.vector.generate $(SIZE)
+
+example-search:
+	@python -B -m example.vector.search $(QUERY)
+
+example-delete:
+	@python -B -m example.vector.delete
+
+database:
+	@echo '# Starting PostgreSQL database with pgvector...'
+	@docker run -d \
+		--name postgres-draive-example \
+		-e POSTGRES_USER=$(POSTGRES_USER) \
+		-e POSTGRES_PASSWORD=$(POSTGRES_PASSWORD) \
+		-e POSTGRES_DB=$(POSTGRES_DATABASE) \
+		-p $(POSTGRES_PORT):5432 \
+		pgvector/pgvector:pg16 \
+		|| docker start postgres-draive-example
+	@echo '...database started!'
+
+database-stop:
+	@echo '# Stopping PostgreSQL database...'
+	@docker stop postgres-draive-example || true
+	@echo '...database stopped!'

postgres/pyproject.toml

@@ -5,8 +5,8 @@ build-backend = "uv_build"
 [project]
 name = "postgres-example"
 version = "0.1.0"
-requires-python = ">=3.12"
-dependencies = ["draive[openai, postgres]~=0.87.3"]
+requires-python = ">=3.13"
+dependencies = ["draive[openai, postgres]>=0.91.1", "pgvector"]
 
 [project.urls]
 Homepage = "https://miquido.com"
@@ -18,7 +18,7 @@ dev = ["bandit~=1.7", "pyright~=1.1", "ruff~=0.12"]
 namespace = true
 
 [tool.ruff]
-target-version = "py312"
+target-version = "py313"
 line-length = 100
 extend-exclude = [".venv", ".git", ".cache"]
 lint.select = ["E", "F", "A", "I", "B", "PL", "W", "C", "RUF", "UP"]
@@ -28,7 +28,7 @@ lint.ignore = ["A005"]
 "__init__.py" = ["F401", "E402"]
 
 [tool.pyright]
-pythonVersion = "3.12"
+pythonVersion = "3.13"
 venvPath = "./.venv"
 include = ["./src"]
 exclude = ["**/node_modules", "**/__pycache__"]

postgres/src/example/__main__.py

@@ -3,14 +3,14 @@
 from draive import (
     Conversation,
     ConversationMessage,
-    Instructions,
+    Template,
     ctx,
 )
 from draive.openai import OpenAI, OpenAIResponsesConfig
 from draive.postgres import (
     PostgresConfigurationRepository,
     PostgresConnectionPool,
-    PostgresInstructionsRepository,
+    PostgresTemplatesRepository,
     PostgresModelMemory,
 )
 
@@ -20,8 +20,8 @@ async def main() -> None:
         "example",
         # declare postgres as configuration provider
         PostgresConfigurationRepository(),
-        # declare postgres as instructons repository
-        PostgresInstructionsRepository(),
+        # declare postgres as templates repository
+        PostgresTemplatesRepository(),
         disposables=(
             OpenAI(),  # use OpenAI for LLM
             PostgresConnectionPool(),  # use postgres connection pool
@@ -31,7 +31,7 @@ async def main() -> None:
             memory = PostgresModelMemory("example_session")
             await memory.maintenance()  # initialize session if needed
             result: ConversationMessage = await Conversation.completion(
-                instructions=Instructions.of("example"),
+                instructions=Template.of("example"),
                 input="Hello!",
                 # declare postgres as conversation memory
                 # use actual session ID to distinct multiple sessions

postgres/src/example/vector/__init__.py

@@ -0,0 +1,4 @@
+from draive import load_env, setup_logging
+
+load_env()
+setup_logging("vector_index")

postgres/src/example/vector/__main__.py

@@ -0,0 +1,103 @@
+"""Vector Indexing and Search Example
+
+This script demonstrates a simple end-to-end example of vector indexing and semantic search
+using PostgreSQL with pgvector extension. It's a quick demonstration that shows both indexing
+and searching in a single script.
+
+What it does:
+    1. Creates 5 sample text chunks about various technologies
+    2. Converts them into vector embeddings using OpenAI
+    3. Indexes the chunks in PostgreSQL with pgvector
+    4. Performs a semantic search query: "What is Python?"
+    5. Displays the top 3 most similar results
+
+Use cases:
+    - Quick demonstration of vector search capabilities
+    - Testing vector indexing and search functionality
+    - Learning how to use the VectorIndex API
+    - Smoke test for database and OpenAI connectivity
+
+Usage:
+    # Run using make
+    make example-index
+
+    # Run directly with Python
+    python -m example.vector
+
+Prerequisites:
+    - PostgreSQL database with pgvector extension running (use 'make database' to start)
+    - OpenAI API key configured in .env file:
+        # OpenAI: Key for draive-examples
+        OPENAI_API_KEY=sk-proj-xxxxxxxx
+        OPENAI_MODEL=text-embedding-3-small
+    - Dependencies installed (use 'make venv' or 'make sync')
+
+Example output:
+    The script will:
+    1. Index 5 sample chunks with progress message
+    2. Perform a search for "What is Python?"
+    3. Display the 3 most relevant results
+
+For more advanced examples:
+    - Use 'make example-generate' to index larger datasets (10 to 100,000 documents)
+    - Use 'make example-search' to perform custom searches
+    - Use 'make example-delete' to clean up the vector index
+"""
+
+from asyncio import run
+
+from draive import ctx
+from draive.openai import OpenAI
+from draive.postgres import PostgresConnectionPool, PostgresVectorIndex
+from draive.utils import VectorIndex
+
+from .common.connection import build_postgres_dsn, initialize_pgvector
+from .common.model import Chunk
+
+
+async def main() -> None:
+    async with ctx.scope(
+        "index",
+        PostgresVectorIndex(),
+        disposables=(
+            OpenAI(),  # use OpenAI for embeddings
+            PostgresConnectionPool.of(
+                dsn=build_postgres_dsn(),  # you can pass custom dsn string ie. "postgresql://user:password@host:port/database"
+                initialize=initialize_pgvector,
+            ),  # use postgres connection pool with vector support
+        ),
+    ):
+        # Sample chunks to index
+        chunks = [
+            Chunk(text="Python is a high-level programming language"),
+            Chunk(text="PostgreSQL is a powerful relational database"),
+            Chunk(text="Machine learning is a subset of artificial intelligence"),
+            Chunk(text="Docker is a containerization platform"),
+            Chunk(text="FastAPI is a modern Python web framework"),
+        ]
+
+        print("Indexing chunks...")
+        await VectorIndex.index(
+            Chunk,
+            values=chunks,
+            attribute=Chunk._.text,
+        )
+        print(f"...indexed {len(chunks)} chunks!")
+
+        # Search for similar chunks
+        query = "What is Python?"
+        print(f"\nSearching for: '{query}'")
+        results = await VectorIndex.search(
+            Chunk,
+            query=query,
+            limit=3,
+            score_threshold=0.0,
+            rerank=False,
+        )
+
+        print(f"Found {len(results)} results:")
+        for idx, result in enumerate(results, 1):
+            print(f"{idx}. {result.text}")
+
+
+run(main())

postgres/src/example/vector/common/__init__.py

@@ -0,0 +1,4 @@
+from .dsn import build_postgres_dsn
+from .initialize import initialize_pgvector
+
+__all__ = ["build_postgres_dsn", "initialize_pgvector"]

postgres/src/example/vector/common/connection/__init__.py

@@ -0,0 +1,24 @@
+from os import getenv
+
+__all__ = ["build_postgres_dsn"]
+
+
+def build_postgres_dsn() -> str:
+    """Build PostgreSQL DSN from environment variables.
+
+    Reads connection parameters from environment variables:
+    - POSTGRES_USER (default: postgres)
+    - POSTGRES_PASSWORD (default: postgres)
+    - POSTGRES_HOST (default: localhost)
+    - POSTGRES_PORT (default: 5432)
+    - POSTGRES_DATABASE (default: postgres)
+
+    Returns:
+        PostgreSQL DSN in format: postgresql://user:password@host:port/database
+    """
+    pg_user = getenv("POSTGRES_USER", "postgres")
+    pg_password = getenv("POSTGRES_PASSWORD", "postgres")
+    pg_host = getenv("POSTGRES_HOST", "localhost")
+    pg_port = getenv("POSTGRES_PORT", "5432")
+    pg_database = getenv("POSTGRES_DATABASE", "postgres")
+    return f"postgresql://{pg_user}:{pg_password}@{pg_host}:{pg_port}/{pg_database}"

postgres/src/example/vector/common/connection/dsn.py

@@ -0,0 +1,13 @@
+from asyncpg.connection import Connection
+from pgvector.asyncpg import register_vector
+
+__all__ = ["initialize_pgvector"]
+
+
+async def initialize_pgvector(connection: Connection) -> None:
+    """Initialize pgvector extension for the connection.
+
+    This function registers the vector type with asyncpg to enable
+    vector operations on PostgreSQL connections.
+    """
+    await register_vector(connection)

postgres/src/example/vector/common/connection/initialize.py

@@ -0,0 +1,6 @@
+"""Generator package for creating test documents and queries."""
+
+from .document_generator import generate_documents
+from .query_generator import generate_query_variations
+
+__all__ = ["generate_documents", "generate_query_variations"]