Implement the vchord / pgvector support (#350)

* Implement the vchord / pgvector support

* feat(alembic): detect vector extension and create appropriate index
This commit is contained in:
Damien 2026-02-12 03:47:07 -06:00 committed by GitHub
parent 93ddd41621
commit a713b68b1f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 241 additions and 27 deletions

View file

@ -0,0 +1,98 @@
name: hindsight
# Docker Compose file for Hindsight with PostgreSQL and vectorchord
# docker compose -f docker/docker-compose/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/docker-compose.yaml up -d
# Make sure to set the required environment variables before running:
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
# - Configure LLM provider variables as needed (see below in the hindsight service)
#
# Usage:
# docker compose up -d
#
# Optional environment variables with defaults:
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
# - HINDSIGHT_DB_VERSION: PostgreSQL version (default: 18)
services:
db:
# Use a PostgreSQL-Image with vectorchord extension pre-installed
image: tensorchord/vchord-suite:pg${HINDSIGHT_DB_VERSION:-18-latest}
container_name: hindsight-db
restart: always
# Expose PostgreSQL port
ports:
- "5436:5432"
environment:
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
volumes:
- pg_data:/var/lib/postgresql/${HINDSIGHT_DB_VERSION:-18}/docker
networks:
- hindsight-net
vectorchord-init:
image: tensorchord/vchord-suite:pg18-latest
#container_name: vectorchord-init
depends_on:
- db
environment:
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
command: >
bash -c "
echo 'Waiting for PostgreSQL to be ready...';
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
echo 'PostgreSQL is unavailable - sleeping';
sleep 2;
done;
echo 'PostgreSQL is ready - creating hindsight_db database';
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
echo 'Creating extensions in hindsight_db database';
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;';
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_tokenizer CASCADE;';
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE;';
echo 'Database and extensions created successfully';
"
restart: "no"
networks:
- hindsight-net
hindsight:
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
container_name: hindsight-app
ports:
- "8888:8888"
- "9999:9999"
environment:
# LLM Configuration
- HINDSIGHT_API_LLM_PROVIDER=openai
- HINDSIGHT_API_LLM_MODEL=gpt-5-mini
# LiteLLM Configuration (shared by embeddings and reranker)
# Embeddings Configuration
# NOTE: OpenRouter does support embeddings endpoints
- HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
- HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-large
- DEFAULT_EMBEDDING_DIMENSION=3072
# Reranker Configuration
- HINDSIGHT_API_RERANKER_PROVIDER=litellm
- HINDSIGHT_API_RERANKER_LITELLM_MODEL=deepinfra/Qwen3-Reranker-8B
# Database Configuration
- HINDSIGHT_API_DATABASE_URL=postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
- HINDSIGHT_API_OTEL_TRACES_ENABLED=false
depends_on:
- db
networks:
- hindsight-net
networks:
hindsight-net:
driver: bridge
volumes:
pg_data:

View file

@ -21,6 +21,25 @@ branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _detect_vector_extension() -> str:
"""
Detect available vector extension: 'vchord' or 'pgvector'.
Prefers vchord if both available. Raises error if neither found.
"""
conn = op.get_bind()
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
if vchord_check:
return "vchord"
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
if pgvector_check:
return "pgvector"
raise RuntimeError(
"Neither vchord nor pgvector extension found. Install one: CREATE EXTENSION vchord; or CREATE EXTENSION vector;"
)
def upgrade() -> None:
"""Upgrade schema - create all tables from scratch."""
@ -200,6 +219,17 @@ def upgrade() -> None:
["bank_id", sa.text("event_date DESC")],
postgresql_where=sa.text("fact_type = 'observation'"),
)
# Create vector index - conditional based on available extension
vector_ext = _detect_vector_extension()
if vector_ext == "vchord":
# Use vchordrq index for vchord (supports high-dimensional embeddings)
op.execute("""
CREATE INDEX idx_memory_units_embedding ON memory_units
USING vchordrq (embedding vector_l2_ops)
""")
else: # pgvector
# Use HNSW index for pgvector
op.create_index(
"idx_memory_units_embedding",
"memory_units",

View file

@ -13,6 +13,7 @@ This migration:
from collections.abc import Sequence
from alembic import context, op
from sqlalchemy import text
# revision identifiers, used by Alembic.
revision: str = "n9i0j1k2l3m4"
@ -27,10 +28,32 @@ def _get_schema_prefix() -> str:
return f'"{schema}".' if schema else ""
def _detect_vector_extension() -> str:
"""
Detect available vector extension: 'vchord' or 'pgvector'.
Prefers vchord if both available. Raises error if neither found.
"""
conn = op.get_bind()
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
if vchord_check:
return "vchord"
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
if pgvector_check:
return "pgvector"
raise RuntimeError(
"Neither vchord nor pgvector extension found. Install one: CREATE EXTENSION vchord; or CREATE EXTENSION vector;"
)
def upgrade() -> None:
"""Create learnings and pinned_reflections tables."""
schema = _get_schema_prefix()
# Detect which vector extension is available
vector_ext = _detect_vector_extension()
# 1. Create learnings table
op.execute(f"""
CREATE TABLE {schema}learnings (
@ -57,10 +80,19 @@ def upgrade() -> None:
# Indexes for learnings
op.execute(f"CREATE INDEX idx_learnings_bank_id ON {schema}learnings(bank_id)")
# Create vector index based on detected extension
if vector_ext == "vchord":
op.execute(f"""
CREATE INDEX idx_learnings_embedding ON {schema}learnings
USING vchordrq (embedding vector_l2_ops)
""")
else: # pgvector
op.execute(f"""
CREATE INDEX idx_learnings_embedding ON {schema}learnings
USING hnsw (embedding vector_cosine_ops)
""")
op.execute(f"CREATE INDEX idx_learnings_tags ON {schema}learnings USING GIN(tags)")
# Full-text search for learnings
@ -94,10 +126,19 @@ def upgrade() -> None:
# Indexes for pinned_reflections
op.execute(f"CREATE INDEX idx_pinned_reflections_bank_id ON {schema}pinned_reflections(bank_id)")
# Create vector index based on detected extension
if vector_ext == "vchord":
op.execute(f"""
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
USING vchordrq (embedding vector_l2_ops)
""")
else: # pgvector
op.execute(f"""
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
USING hnsw (embedding vector_cosine_ops)
""")
op.execute(f"CREATE INDEX idx_pinned_reflections_tags ON {schema}pinned_reflections USING GIN(tags)")
# Full-text search for pinned_reflections

View file

@ -33,6 +33,37 @@ logger = logging.getLogger(__name__)
MIGRATION_LOCK_ID = 123456789
def _detect_vector_extension(conn) -> str:
"""
Detect available vector extension: 'vchord' or 'pgvector'.
Prefers vchord if both available. Raises error if neither found.
Args:
conn: SQLAlchemy connection object
Returns:
"vchord" or "pgvector"
Raises:
RuntimeError: If neither extension is installed
"""
# Check vchord first (preferred for high-dimensional embeddings)
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
if vchord_check:
logger.debug("Detected vector extension: vchord")
return "vchord"
# Fall back to pgvector
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
if pgvector_check:
logger.debug("Detected vector extension: pgvector")
return "pgvector"
raise RuntimeError(
"Neither vchord nor pgvector extension found. Install one: CREATE EXTENSION vchord; or CREATE EXTENSION vector;"
)
def _get_schema_lock_id(schema: str) -> int:
"""
Generate a unique advisory lock ID for a schema.
@ -361,6 +392,10 @@ def ensure_embedding_dimension(
logger.debug(f"memory_units table does not exist in schema '{schema_name}', skipping dimension check")
return
# Detect which vector extension is available
vector_ext = _detect_vector_extension(conn)
logger.info(f"Detected vector extension: {vector_ext}")
# Get current column dimension from pg_attribute
# pgvector stores dimension in atttypmod
current_dim = conn.execute(
@ -408,8 +443,7 @@ def ensure_embedding_dimension(
# Table is empty, safe to alter column
logger.info(f"Altering embedding column dimension from {current_dimension} to {required_dimension}")
# Drop the HNSW index on embedding column if it exists
# Only drop indexes that use 'hnsw' and reference the 'embedding' column
# Drop existing vector index (works for both HNSW and vchordrq)
conn.execute(
text(f"""
DO $$
@ -419,7 +453,7 @@ def ensure_embedding_dimension(
SELECT indexname FROM pg_indexes
WHERE schemaname = '{schema_name}'
AND tablename = 'memory_units'
AND indexdef LIKE '%hnsw%'
AND (indexdef LIKE '%hnsw%' OR indexdef LIKE '%vchordrq%')
AND indexdef LIKE '%embedding%'
LOOP
EXECUTE 'DROP INDEX IF EXISTS {schema_name}.' || idx_name;
@ -434,7 +468,17 @@ def ensure_embedding_dimension(
)
conn.commit()
# Recreate the HNSW index
# Recreate index with appropriate type based on detected extension
if vector_ext == "vchord":
conn.execute(
text(f"""
CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_vchordrq
ON {schema_name}.memory_units
USING vchordrq (embedding vector_l2_ops)
""")
)
logger.info(f"Created vchordrq index for {required_dimension}-dimensional embeddings")
else: # pgvector
conn.execute(
text(f"""
CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_hnsw
@ -443,6 +487,7 @@ def ensure_embedding_dimension(
WITH (m = 16, ef_construction = 64)
""")
)
logger.info(f"Created HNSW index for {required_dimension}-dimensional embeddings")
conn.commit()
logger.info(f"Successfully changed embedding dimension to {required_dimension}")