feat: support for other text and vector search pg extensions (#355)

* feat: support for other text and vector search pg extensions

* test: increase timeout for test_batch_chunking_behavior to account for VectorChord BM25 tokenization overhead

* feat: support for other text and vector search pg extensions
This commit is contained in:
Nicolò Boschi 2026-02-12 14:13:04 +01:00 committed by GitHub
parent 8d731f2e5f
commit c029807add
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 917 additions and 204 deletions

View file

@ -52,6 +52,8 @@ services:
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 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 pg_tokenizer CASCADE;';
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE;'; psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE;';
echo 'Creating llmlingua2 tokenizer';
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c \"SELECT create_tokenizer('llmlingua2', \\$\\$ model = \\\"llmlingua2\\\" \\$\\$);\" 2>/dev/null || echo 'Tokenizer already exists or creation skipped';
echo 'Database and extensions created successfully'; echo 'Database and extensions created successfully';
" "
restart: "no" restart: "no"
@ -65,25 +67,18 @@ services:
- "8888:8888" - "8888:8888"
- "9999:9999" - "9999:9999"
environment: environment:
# LLM Configuration # LLM Configuration (uses OpenAI for testing vchord)
- HINDSIGHT_API_LLM_PROVIDER=openai # LLM configuration
- HINDSIGHT_API_LLM_MODEL=gpt-5-mini HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
# 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 # 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_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
# Vector and Text Search Extensions
HINDSIGHT_API_VECTOR_EXTENSION: vchord
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: vchord
depends_on: depends_on:
- db - db
networks: networks:

View file

@ -6,6 +6,7 @@ Create Date: 2025-11-27 11:54:19.228030
""" """
import os
from collections.abc import Sequence from collections.abc import Sequence
import sqlalchemy as sa import sqlalchemy as sa
@ -23,20 +24,56 @@ depends_on: str | Sequence[str] | None = None
def _detect_vector_extension() -> str: def _detect_vector_extension() -> str:
""" """
Detect available vector extension: 'vchord' or 'pgvector'. Detect or validate vector extension: 'vchord' or 'pgvector'.
Prefers vchord if both available. Raises error if neither found. Respects HINDSIGHT_API_VECTOR_EXTENSION env var if set.
""" """
conn = op.get_bind() conn = op.get_bind()
vector_extension = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
# Validate configured extension is installed
if vector_extension == "vchord":
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar() vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
if vchord_check: if not 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( raise RuntimeError(
"Neither vchord nor pgvector extension found. Install one: CREATE EXTENSION vchord; or CREATE EXTENSION vector;" "Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;"
)
return "vchord"
elif vector_extension == "pgvector":
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
if not pgvector_check:
raise RuntimeError(
"Configured vector extension 'pgvector' not found. Install it with: CREATE EXTENSION vector;"
)
return "pgvector"
else:
raise ValueError(f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {vector_extension}. Must be 'pgvector' or 'vchord'")
def _detect_text_search_extension() -> str:
"""
Detect or validate text search extension: 'native' or 'vchord'.
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
Creates the extension if needed.
"""
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
if text_search_extension == "vchord":
# Create vchord_bm25 extension if not exists
try:
op.execute("CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE")
except Exception:
# Extension might already exist or user lacks permissions - verify it exists
conn = op.get_bind()
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord_bm25'")).fetchone()
if not result:
# Extension truly doesn't exist - re-raise the error
raise
return "vchord"
elif text_search_extension == "native":
return "native"
else:
raise ValueError(
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native' or 'vchord'"
) )
@ -185,6 +222,18 @@ def upgrade() -> None:
) )
# Add search_vector column for full-text search # Add search_vector column for full-text search
# Type depends on configured text search backend
text_search_ext = _detect_text_search_extension()
if text_search_ext == "vchord":
# VectorChord BM25: bm25vector type (no GENERATED - tokenization happens on INSERT)
# Note: vchord_bm25 extension creates types in bm25_catalog schema
op.execute("""
ALTER TABLE memory_units
ADD COLUMN search_vector bm25_catalog.bm25vector
""")
else: # native
# Native PostgreSQL: tsvector with automatic generation
op.execute(""" op.execute("""
ALTER TABLE memory_units ALTER TABLE memory_units
ADD COLUMN search_vector tsvector ADD COLUMN search_vector tsvector
@ -238,7 +287,16 @@ def upgrade() -> None:
postgresql_ops={"embedding": "vector_cosine_ops"}, postgresql_ops={"embedding": "vector_cosine_ops"},
) )
# Create BM25 full-text search index on search_vector # Create full-text search index on search_vector
# Index type depends on text search backend
if text_search_ext == "vchord":
# VectorChord BM25 index
op.execute("""
CREATE INDEX idx_memory_units_text_search ON memory_units
USING bm25 (search_vector bm25_catalog.bm25_ops)
""")
else: # native
# Native PostgreSQL GIN index
op.execute(""" op.execute("""
CREATE INDEX idx_memory_units_text_search ON memory_units CREATE INDEX idx_memory_units_text_search ON memory_units
USING gin(search_vector) USING gin(search_vector)

View file

@ -10,6 +10,7 @@ This migration:
3. Adds consolidation tracking columns to the 'banks' table 3. Adds consolidation tracking columns to the 'banks' table
""" """
import os
from collections.abc import Sequence from collections.abc import Sequence
from alembic import context, op from alembic import context, op
@ -30,20 +31,56 @@ def _get_schema_prefix() -> str:
def _detect_vector_extension() -> str: def _detect_vector_extension() -> str:
""" """
Detect available vector extension: 'vchord' or 'pgvector'. Detect or validate vector extension: 'vchord' or 'pgvector'.
Prefers vchord if both available. Raises error if neither found. Respects HINDSIGHT_API_VECTOR_EXTENSION env var if set.
""" """
conn = op.get_bind() conn = op.get_bind()
vector_extension = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
# Validate configured extension is installed
if vector_extension == "vchord":
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar() vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
if vchord_check: if not 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( raise RuntimeError(
"Neither vchord nor pgvector extension found. Install one: CREATE EXTENSION vchord; or CREATE EXTENSION vector;" "Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;"
)
return "vchord"
elif vector_extension == "pgvector":
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
if not pgvector_check:
raise RuntimeError(
"Configured vector extension 'pgvector' not found. Install it with: CREATE EXTENSION vector;"
)
return "pgvector"
else:
raise ValueError(f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {vector_extension}. Must be 'pgvector' or 'vchord'")
def _detect_text_search_extension() -> str:
"""
Detect or validate text search extension: 'native' or 'vchord'.
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
Creates the extension if needed.
"""
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
if text_search_extension == "vchord":
# Create vchord_bm25 extension if not exists
try:
op.execute("CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE")
except Exception:
# Extension might already exist or user lacks permissions - verify it exists
conn = op.get_bind()
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord_bm25'")).fetchone()
if not result:
# Extension truly doesn't exist - re-raise the error
raise
return "vchord"
elif text_search_extension == "native":
return "native"
else:
raise ValueError(
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native' or 'vchord'"
) )
@ -54,6 +91,9 @@ def upgrade() -> None:
# Detect which vector extension is available # Detect which vector extension is available
vector_ext = _detect_vector_extension() vector_ext = _detect_vector_extension()
# Detect which text search extension to use
text_search_ext = _detect_text_search_extension()
# 1. Create learnings table # 1. Create learnings table
op.execute(f""" op.execute(f"""
CREATE TABLE {schema}learnings ( CREATE TABLE {schema}learnings (
@ -96,6 +136,18 @@ def upgrade() -> None:
op.execute(f"CREATE INDEX idx_learnings_tags ON {schema}learnings USING GIN(tags)") op.execute(f"CREATE INDEX idx_learnings_tags ON {schema}learnings USING GIN(tags)")
# Full-text search for learnings # Full-text search for learnings
if text_search_ext == "vchord":
# VectorChord BM25: bm25vector type (no GENERATED - tokenization happens on INSERT)
# Note: vchord_bm25 extension creates types in bm25_catalog schema
op.execute(f"""
ALTER TABLE {schema}learnings ADD COLUMN search_vector bm25_catalog.bm25vector
""")
op.execute(f"""
CREATE INDEX idx_learnings_text_search ON {schema}learnings
USING bm25 (search_vector bm25_catalog.bm25_ops)
""")
else: # native
# Native PostgreSQL: tsvector with automatic generation
op.execute(f""" op.execute(f"""
ALTER TABLE {schema}learnings ADD COLUMN search_vector tsvector ALTER TABLE {schema}learnings ADD COLUMN search_vector tsvector
GENERATED ALWAYS AS (to_tsvector('english', text)) STORED GENERATED ALWAYS AS (to_tsvector('english', text)) STORED
@ -142,6 +194,18 @@ def upgrade() -> None:
op.execute(f"CREATE INDEX idx_pinned_reflections_tags ON {schema}pinned_reflections USING GIN(tags)") op.execute(f"CREATE INDEX idx_pinned_reflections_tags ON {schema}pinned_reflections USING GIN(tags)")
# Full-text search for pinned_reflections # Full-text search for pinned_reflections
if text_search_ext == "vchord":
# VectorChord BM25: bm25vector type (no GENERATED - tokenization happens on INSERT/UPDATE)
# Note: vchord_bm25 extension creates types in bm25_catalog schema
op.execute(f"""
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector bm25_catalog.bm25vector
""")
op.execute(f"""
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
USING bm25 (search_vector bm25_catalog.bm25_ops)
""")
else: # native
# Native PostgreSQL: tsvector with automatic generation
op.execute(f""" op.execute(f"""
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector tsvector ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector tsvector
GENERATED ALWAYS AS (to_tsvector('english', COALESCE(name, '') || ' ' || content)) STORED GENERATED ALWAYS AS (to_tsvector('english', COALESCE(name, '') || ' ' || content)) STORED

View file

@ -205,6 +205,9 @@ ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES"
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL" ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR" ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
ENV_HOST = "HINDSIGHT_API_HOST" ENV_HOST = "HINDSIGHT_API_HOST"
ENV_PORT = "HINDSIGHT_API_PORT" ENV_PORT = "HINDSIGHT_API_PORT"
ENV_BASE_PATH = "HINDSIGHT_API_BASE_PATH" ENV_BASE_PATH = "HINDSIGHT_API_BASE_PATH"
@ -323,6 +326,12 @@ DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0" DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0"
DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0" DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
# Vector extension (pgvector vs vchord)
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord"
# Text search extension (native PostgreSQL vs vchord BM25)
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord"
# LiteLLM defaults # LiteLLM defaults
DEFAULT_LITELLM_API_BASE = "http://localhost:4000" DEFAULT_LITELLM_API_BASE = "http://localhost:4000"
DEFAULT_EMBEDDINGS_LITELLM_MODEL = "text-embedding-3-small" DEFAULT_EMBEDDINGS_LITELLM_MODEL = "text-embedding-3-small"
@ -460,6 +469,8 @@ class HindsightConfig:
# Database # Database
database_url: str database_url: str
database_schema: str database_schema: str
vector_extension: str # "pgvector" or "vchord"
text_search_extension: str # "native" or "vchord"
# LLM (default, used as fallback for per-operation config) # LLM (default, used as fallback for per-operation config)
llm_provider: str llm_provider: str
@ -687,6 +698,20 @@ class HindsightConfig:
def validate(self) -> None: def validate(self) -> None:
"""Validate configuration values and raise errors for invalid combinations.""" """Validate configuration values and raise errors for invalid combinations."""
# Validate vector_extension
valid_extensions = ("pgvector", "vchord")
if self.vector_extension not in valid_extensions:
raise ValueError(
f"Invalid vector_extension: {self.vector_extension}. Must be one of: {', '.join(valid_extensions)}"
)
# Validate text_search_extension
valid_text_search = ("native", "vchord")
if self.text_search_extension not in valid_text_search:
raise ValueError(
f"Invalid text_search_extension: {self.text_search_extension}. Must be one of: {', '.join(valid_text_search)}"
)
# RETAIN_MAX_COMPLETION_TOKENS must be greater than RETAIN_CHUNK_SIZE # RETAIN_MAX_COMPLETION_TOKENS must be greater than RETAIN_CHUNK_SIZE
# to ensure the LLM has enough output capacity to extract facts from chunks # to ensure the LLM has enough output capacity to extract facts from chunks
if self.retain_max_completion_tokens <= self.retain_chunk_size: if self.retain_max_completion_tokens <= self.retain_chunk_size:
@ -712,6 +737,8 @@ class HindsightConfig:
# Database # Database
database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL), database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL),
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA), database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
# LLM # LLM
llm_provider=llm_provider, llm_provider=llm_provider,
llm_api_key=os.getenv(ENV_LLM_API_KEY), llm_api_key=os.getenv(ENV_LLM_API_KEY),

View file

@ -18,6 +18,7 @@ import uuid
from datetime import datetime, timezone from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
from ...config import get_config
from ..memory_engine import fq_table from ..memory_engine import fq_table
from ..retain import embedding_utils from ..retain import embedding_utils
from .prompts import ( from .prompts import (
@ -1015,15 +1016,33 @@ async def _create_observation_directly(
t0 = time.time() t0 = time.time()
observation_id = uuid.uuid4() observation_id = uuid.uuid4()
row = await conn.fetchrow(
f""" # Query varies based on text search backend
config = get_config()
if config.text_search_extension == "vchord":
# VectorChord: manually tokenize and insert search_vector
query = f"""
INSERT INTO {fq_table("memory_units")} (
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
tags, event_date, occurred_start, occurred_end, mentioned_at, search_vector
)
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10,
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
RETURNING id
"""
else: # native
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
query = f"""
INSERT INTO {fq_table("memory_units")} ( INSERT INTO {fq_table("memory_units")} (
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history, id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
tags, event_date, occurred_start, occurred_end, mentioned_at tags, event_date, occurred_start, occurred_end, mentioned_at
) )
VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10) VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10)
RETURNING id RETURNING id
""", """
row = await conn.fetchrow(
query,
observation_id, observation_id,
bank_id, bank_id,
observation_text, observation_text,

View file

@ -968,7 +968,12 @@ class MemoryEngine(MemoryEngineInterface):
# Run database migrations if enabled # Run database migrations if enabled
if self._run_migrations: if self._run_migrations:
from ..migrations import ensure_embedding_dimension, run_migrations from ..migrations import (
ensure_embedding_dimension,
ensure_text_search_extension,
ensure_vector_extension,
run_migrations,
)
if not self.db_url: if not self.db_url:
raise ValueError("Database URL is required for migrations") raise ValueError("Database URL is required for migrations")
@ -976,30 +981,43 @@ class MemoryEngine(MemoryEngineInterface):
# Migrate all schemas from the tenant extension # Migrate all schemas from the tenant extension
# The tenant extension is the single source of truth for which schemas exist # The tenant extension is the single source of truth for which schemas exist
logger.info("Running database migrations...") logger.info("Running database migrations...")
try:
tenants = await self._tenant_extension.list_tenants() tenants = await self._tenant_extension.list_tenants()
if tenants: if tenants:
logger.info(f"Running migrations on {len(tenants)} schema(s)...") logger.info(f"Running migrations on {len(tenants)} schema(s)...")
for tenant in tenants: for tenant in tenants:
schema = tenant.schema schema = tenant.schema
if schema: if schema:
try:
run_migrations(self.db_url, schema=schema) run_migrations(self.db_url, schema=schema)
except Exception as e:
logger.warning(f"Failed to migrate schema {schema}: {e}")
logger.info("Schema migrations completed") logger.info("Schema migrations completed")
# Get config for vector extension setting
config = get_config()
# Ensure embedding column dimension matches the model's dimension # Ensure embedding column dimension matches the model's dimension
# This is done after migrations and after embeddings.initialize() # This is done after migrations and after embeddings.initialize()
for tenant in tenants: for tenant in tenants:
schema = tenant.schema schema = tenant.schema
if schema: if schema:
try: ensure_embedding_dimension(
ensure_embedding_dimension(self.db_url, self.embeddings.dimension, schema=schema) self.db_url,
except Exception as e: self.embeddings.dimension,
logger.warning(f"Failed to ensure embedding dimension for schema {schema}: {e}") schema=schema,
except Exception as e: vector_extension=config.vector_extension,
logger.warning(f"Failed to run schema migrations: {e}") )
# Ensure vector indexes match the configured extension
for tenant in tenants:
schema = tenant.schema
if schema:
ensure_vector_extension(self.db_url, vector_extension=config.vector_extension, schema=schema)
# Ensure text search columns/indexes match the configured extension
for tenant in tenants:
schema = tenant.schema
if schema:
ensure_text_search_extension(
self.db_url, text_search_extension=config.text_search_extension, schema=schema
)
logger.info(f"Connecting to PostgreSQL at {mask_network_location(self.db_url)}") logger.info(f"Connecting to PostgreSQL at {mask_network_location(self.db_url)}")

View file

@ -7,6 +7,7 @@ Handles insertion of facts into the database.
import json import json
import logging import logging
from ...config import get_config
from ..memory_engine import fq_table from ..memory_engine import fq_table
from .fact_extraction import _sanitize_text from .fact_extraction import _sanitize_text
from .types import ProcessedFact from .types import ProcessedFact
@ -70,8 +71,35 @@ async def insert_facts_batch(
# Batch insert all facts # Batch insert all facts
# Note: tags are passed as JSON strings and converted back to varchar[] via jsonb_array_elements_text + array_agg # Note: tags are passed as JSON strings and converted back to varchar[] via jsonb_array_elements_text + array_agg
results = await conn.fetch( # Query varies based on text search backend
f""" config = get_config()
if config.text_search_extension == "vchord":
# VectorChord: manually tokenize and insert search_vector
query = f"""
WITH input_data AS (
SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[]
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json)
)
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags, search_vector)
SELECT
$1,
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, metadata, chunk_id, document_id,
COALESCE(
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
'{{}}'::varchar[]
),
tokenize(COALESCE(text, '') || ' ' || COALESCE(context, ''), 'llmlingua2')::bm25_catalog.bm25vector
FROM input_data
RETURNING id
"""
else: # native
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
query = f"""
WITH input_data AS ( WITH input_data AS (
SELECT * FROM unnest( SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[], $2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
@ -91,7 +119,10 @@ async def insert_facts_batch(
) )
FROM input_data FROM input_data
RETURNING id RETURNING id
""", """
results = await conn.fetch(
query,
bank_id, bank_id,
fact_texts, fact_texts,
embeddings, embeddings,

View file

@ -158,7 +158,7 @@ async def retrieve_bm25(
from .tags import TagsMatch, build_tags_where_clause_simple from .tags import TagsMatch, build_tags_where_clause_simple
# Sanitize query text: remove special characters that have meaning in tsquery # Sanitize query text for native backend: remove special characters that have meaning in tsquery
# Keep only alphanumeric characters and spaces # Keep only alphanumeric characters and spaces
sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower()) sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower())
@ -169,17 +169,34 @@ async def retrieve_bm25(
# If no valid tokens, return empty results # If no valid tokens, return empty results
return [] return []
# Convert query to tsquery using OR for more flexible matching # Build query based on text search backend
# This prevents empty results when some terms are missing config = get_config()
query_tsquery = " | ".join(tokens)
tags_clause = build_tags_where_clause_simple(tags, 5) tags_clause = build_tags_where_clause_simple(tags, 5)
if config.text_search_extension == "vchord":
# VectorChord BM25: use <&> operator with to_bm25query and tokenize
params = [bank_id, fact_type, limit, query_text] # Use raw query_text for tokenization
if tags:
params.append(query_text) # VectorChord doesn't need sanitization
query = f"""
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($4, 'llmlingua2')) AS bm25_score
FROM {fq_table("memory_units")}
WHERE bank_id = $1
AND fact_type = $2
{tags_clause}
ORDER BY bm25_score DESC
LIMIT $3
"""
else: # native
# Native PostgreSQL: use ts_rank_cd with to_tsquery
query_tsquery = " | ".join(tokens)
params = [query_tsquery, bank_id, fact_type, limit] params = [query_tsquery, bank_id, fact_type, limit]
if tags: if tags:
params.append(tags) params.append(tags)
results = await conn.fetch( query = f"""
f"""
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags, SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score
FROM {fq_table("memory_units")} FROM {fq_table("memory_units")}
@ -189,9 +206,9 @@ async def retrieve_bm25(
{tags_clause} {tags_clause}
ORDER BY bm25_score DESC ORDER BY bm25_score DESC
LIMIT $4 LIMIT $4
""", """
*params,
) results = await conn.fetch(query, *params)
return [RetrievalResult.from_db_row(dict(r)) for r in results] return [RetrievalResult.from_db_row(dict(r)) for r in results]
@ -268,18 +285,66 @@ async def retrieve_semantic_bm25_combined(
result_dict[ft][0].append(RetrievalResult.from_db_row(row)) result_dict[ft][0].append(RetrievalResult.from_db_row(row))
return result_dict return result_dict
query_tsquery = " | ".join(tokens) # Build BM25 query based on text search backend
config = get_config()
# Build tags clause - param 6 if tags provided # Build tags clause - param 6 if tags provided
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match) tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
if config.text_search_extension == "vchord":
# VectorChord BM25: use <&> operator with to_bm25query and tokenize
# Note: VectorChord scores are negative (higher = better, so -1 > -10)
params = [query_emb_str, bank_id, fact_types, limit, query_text] # Pass raw query_text for tokenization
if tags:
params.append(tags)
query = f"""
WITH semantic_ranked AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
1 - (embedding <=> $1::vector) AS similarity,
NULL::float AS bm25_score,
'semantic' AS source,
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY embedding <=> $1::vector) AS rn
FROM {fq_table("memory_units")}
WHERE bank_id = $2
AND embedding IS NOT NULL
AND fact_type = ANY($3)
AND (1 - (embedding <=> $1::vector)) >= 0.3
{tags_clause}
),
bm25_ranked AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
NULL::float AS similarity,
search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($5, 'llmlingua2')) AS bm25_score,
'bm25' AS source,
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($5, 'llmlingua2')) DESC) AS rn
FROM {fq_table("memory_units")}
WHERE bank_id = $2
AND fact_type = ANY($3)
{tags_clause}
),
semantic AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
similarity, bm25_score, source
FROM semantic_ranked WHERE rn <= $4
),
bm25 AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
similarity, bm25_score, source
FROM bm25_ranked WHERE rn <= $4
)
SELECT * FROM semantic
UNION ALL
SELECT * FROM bm25
"""
else: # native
# Native PostgreSQL: use ts_rank_cd with to_tsquery
query_tsquery = " | ".join(tokens)
params = [query_emb_str, bank_id, fact_types, limit, query_tsquery] params = [query_emb_str, bank_id, fact_types, limit, query_tsquery]
if tags: if tags:
params.append(tags) params.append(tags)
# Combined CTE query for both semantic and BM25 across all fact types query = f"""
# Uses window functions to limit per fact_type per method
results = await conn.fetch(
f"""
WITH semantic_ranked AS ( WITH semantic_ranked AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags, SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, embedding, fact_type, document_id, chunk_id, tags,
1 - (embedding <=> $1::vector) AS similarity, 1 - (embedding <=> $1::vector) AS similarity,
@ -318,9 +383,11 @@ async def retrieve_semantic_bm25_combined(
SELECT * FROM semantic SELECT * FROM semantic
UNION ALL UNION ALL
SELECT * FROM bm25 SELECT * FROM bm25
""", """
*params,
) # Combined CTE query for both semantic and BM25 across all fact types
# Uses window functions to limit per fact_type per method
results = await conn.fetch(query, *params)
# Group results by fact_type and source # Group results by fact_type and source
result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types} result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types}

View file

@ -96,7 +96,13 @@ class DefaultExtensionContext(ExtensionContext):
async def run_migration(self, schema: str) -> None: async def run_migration(self, schema: str) -> None:
"""Run migrations for a specific schema.""" """Run migrations for a specific schema."""
from hindsight_api.migrations import ensure_embedding_dimension, run_migrations from hindsight_api.config import get_config
from hindsight_api.migrations import (
ensure_embedding_dimension,
ensure_text_search_extension,
ensure_vector_extension,
run_migrations,
)
# Prefer getting URL from memory engine (handles pg0 case where URL is set after init) # Prefer getting URL from memory engine (handles pg0 case where URL is set after init)
db_url = self._database_url db_url = self._database_url
@ -107,6 +113,9 @@ class DefaultExtensionContext(ExtensionContext):
run_migrations(db_url, schema=schema) run_migrations(db_url, schema=schema)
# Get config for vector extension setting
config = get_config()
# Ensure embedding column dimension matches the model's dimension # Ensure embedding column dimension matches the model's dimension
# This is needed because migrations create columns with default dimension # This is needed because migrations create columns with default dimension
if self._memory_engine is not None: if self._memory_engine is not None:
@ -114,7 +123,15 @@ class DefaultExtensionContext(ExtensionContext):
if embeddings is not None: if embeddings is not None:
dimension = getattr(embeddings, "dimension", None) dimension = getattr(embeddings, "dimension", None)
if dimension is not None: if dimension is not None:
ensure_embedding_dimension(db_url, dimension, schema=schema) ensure_embedding_dimension(
db_url, dimension, schema=schema, vector_extension=config.vector_extension
)
# Ensure vector indexes match the configured extension
ensure_vector_extension(db_url, vector_extension=config.vector_extension, schema=schema)
# Ensure text search columns/indexes match the configured extension
ensure_text_search_extension(db_url, text_search_extension=config.text_search_extension, schema=schema)
def get_memory_engine(self) -> "MemoryEngineInterface": def get_memory_engine(self) -> "MemoryEngineInterface":
"""Get the memory engine interface.""" """Get the memory engine interface."""

View file

@ -155,6 +155,8 @@ def main():
config = HindsightConfig( config = HindsightConfig(
database_url=config.database_url, database_url=config.database_url,
database_schema=config.database_schema, database_schema=config.database_schema,
vector_extension=config.vector_extension,
text_search_extension=config.text_search_extension,
llm_provider=config.llm_provider, llm_provider=config.llm_provider,
llm_api_key=config.llm_api_key, llm_api_key=config.llm_api_key,
llm_model=config.llm_model, llm_model=config.llm_model,

View file

@ -33,35 +33,39 @@ logger = logging.getLogger(__name__)
MIGRATION_LOCK_ID = 123456789 MIGRATION_LOCK_ID = 123456789
def _detect_vector_extension(conn) -> str: def _detect_vector_extension(conn, vector_extension: str = "pgvector") -> str:
""" """
Detect available vector extension: 'vchord' or 'pgvector'. Validate vector extension: 'vchord' or 'pgvector'.
Prefers vchord if both available. Raises error if neither found.
Args: Args:
conn: SQLAlchemy connection object conn: SQLAlchemy connection object
vector_extension: Configured extension ("pgvector" or "vchord")
Returns: Returns:
"vchord" or "pgvector" "vchord" or "pgvector"
Raises: Raises:
RuntimeError: If neither extension is installed RuntimeError: If configured extension is not installed
""" """
# Check vchord first (preferred for high-dimensional embeddings) # Verify the configured extension is installed
if vector_extension == "vchord":
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar() vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
if vchord_check: if not 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( raise RuntimeError(
"Neither vchord nor pgvector extension found. Install one: CREATE EXTENSION vchord; or CREATE EXTENSION vector;" "Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;"
) )
logger.debug("Using configured vector extension: vchord")
return "vchord"
elif vector_extension == "pgvector":
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
if not pgvector_check:
raise RuntimeError(
"Configured vector extension 'pgvector' not found. Install it with: CREATE EXTENSION vector;"
)
logger.debug("Using configured vector extension: pgvector")
return "pgvector"
else:
raise ValueError(f"Invalid vector_extension: {vector_extension}. Must be 'pgvector' or 'vchord'")
def _get_schema_lock_id(schema: str) -> int: def _get_schema_lock_id(schema: str) -> int:
@ -355,6 +359,7 @@ def ensure_embedding_dimension(
database_url: str, database_url: str,
required_dimension: int, required_dimension: int,
schema: str | None = None, schema: str | None = None,
vector_extension: str = "pgvector",
) -> None: ) -> None:
""" """
Ensure the embedding column dimension matches the model's dimension. Ensure the embedding column dimension matches the model's dimension.
@ -369,6 +374,7 @@ def ensure_embedding_dimension(
database_url: SQLAlchemy database URL database_url: SQLAlchemy database URL
required_dimension: The embedding dimension required by the model required_dimension: The embedding dimension required by the model
schema: Target PostgreSQL schema name (None for public) schema: Target PostgreSQL schema name (None for public)
vector_extension: Configured vector extension ("pgvector" or "vchord")
Raises: Raises:
RuntimeError: If dimension mismatch with existing data RuntimeError: If dimension mismatch with existing data
@ -393,8 +399,8 @@ def ensure_embedding_dimension(
return return
# Detect which vector extension is available # Detect which vector extension is available
vector_ext = _detect_vector_extension(conn) vector_ext = _detect_vector_extension(conn, vector_extension)
logger.info(f"Detected vector extension: {vector_ext}") logger.info(f"Using vector extension: {vector_ext}")
# Get current column dimension from pg_attribute # Get current column dimension from pg_attribute
# pgvector stores dimension in atttypmod # pgvector stores dimension in atttypmod
@ -491,3 +497,354 @@ def ensure_embedding_dimension(
conn.commit() conn.commit()
logger.info(f"Successfully changed embedding dimension to {required_dimension}") logger.info(f"Successfully changed embedding dimension to {required_dimension}")
def ensure_vector_extension(
database_url: str,
vector_extension: str = "pgvector",
schema: str | None = None,
) -> None:
"""
Ensure the vector indexes match the configured vector extension.
This function checks the current vector index type in the database
and adjusts it if necessary:
- If index type matches configured extension: no action needed
- If they differ and tables are empty: drop old indexes, recreate with new type
- If they differ and tables have data: raise error with migration guidance
Args:
database_url: SQLAlchemy database URL
vector_extension: Configured vector extension ("pgvector" or "vchord")
schema: Target PostgreSQL schema name (None for public)
Raises:
RuntimeError: If extension mismatch with existing data
"""
schema_name = schema or "public"
engine = create_engine(database_url)
with engine.connect() as conn:
# Detect which vector extension should be used
target_ext = _detect_vector_extension(conn, vector_extension)
logger.info(f"Target vector extension: {target_ext}")
# Tables with vector indexes to check
tables_to_check = [
("memory_units", "idx_memory_units_embedding"),
("learnings", "idx_learnings_embedding"),
("pinned_reflections", "idx_pinned_reflections_embedding"),
]
# Determine target index type
target_index_type = "vchordrq" if target_ext == "vchord" else "hnsw"
mismatched_tables = []
tables_with_data = []
for table_name, index_name in tables_to_check:
# Check if table exists
table_exists = conn.execute(
text("""
SELECT EXISTS (
SELECT 1 FROM information_schema.tables
WHERE table_schema = :schema AND table_name = :table_name
)
"""),
{"schema": schema_name, "table_name": table_name},
).scalar()
if not table_exists:
logger.debug(f"Table {table_name} does not exist in schema '{schema_name}', skipping")
continue
# Check current index type by querying pg_indexes
current_index_info = conn.execute(
text("""
SELECT indexdef
FROM pg_indexes
WHERE schemaname = :schema
AND tablename = :table_name
AND indexname LIKE :index_pattern
"""),
{"schema": schema_name, "table_name": table_name, "index_pattern": "%embedding%"},
).fetchone()
if not current_index_info:
logger.warning(f"No embedding index found for {table_name}, will create it")
mismatched_tables.append((table_name, index_name, None))
continue
indexdef = current_index_info[0].lower()
if "vchordrq" in indexdef:
current_index_type = "vchordrq"
elif "hnsw" in indexdef:
current_index_type = "hnsw"
else:
logger.warning(f"Unknown index type for {table_name}: {indexdef}")
continue
# Check if index type matches target
if current_index_type != target_index_type:
logger.info(
f"Index type mismatch on {table_name}: current={current_index_type}, target={target_index_type}"
)
mismatched_tables.append((table_name, index_name, current_index_type))
# Check if table has data
row_count = conn.execute(
text(f"SELECT COUNT(*) FROM {schema_name}.{table_name} WHERE embedding IS NOT NULL")
).scalar()
if row_count > 0:
tables_with_data.append((table_name, row_count))
else:
logger.debug(f"Index type OK for {table_name}: {current_index_type}")
# If no mismatches, we're done
if not mismatched_tables:
logger.debug(f"All vector indexes match configured extension: {target_ext}")
return
# If there's data in any mismatched table, raise error
if tables_with_data:
table_list = ", ".join([f"{table}({count} rows)" for table, count in tables_with_data])
raise RuntimeError(
f"Cannot change vector extension from {current_index_type} to {target_index_type}: "
f"the following tables contain data: {table_list}. "
f"To change vector extension, you must either:\n"
f" 1. Re-embed all data: DELETE FROM {schema_name}.memory_units; "
f"DELETE FROM {schema_name}.learnings; DELETE FROM {schema_name}.pinned_reflections; then restart\n"
f" 2. Use the current vector extension (set HINDSIGHT_API_VECTOR_EXTENSION='{current_index_type.replace('vchordrq', 'vchord').replace('hnsw', 'pgvector')}')"
)
# Tables are empty, safe to recreate indexes
logger.info(f"Recreating vector indexes for {target_ext}")
for table_name, index_name, current_type in mismatched_tables:
# Drop existing index if it exists
if current_type:
logger.info(f"Dropping {current_type} index on {table_name}")
conn.execute(text(f"DROP INDEX IF EXISTS {schema_name}.{index_name}"))
# Create new index with appropriate type
if target_ext == "vchord":
logger.info(f"Creating vchordrq index on {table_name}")
conn.execute(
text(f"""
CREATE INDEX IF NOT EXISTS {index_name}
ON {schema_name}.{table_name}
USING vchordrq (embedding vector_l2_ops)
""")
)
else: # pgvector
logger.info(f"Creating HNSW index on {table_name}")
conn.execute(
text(f"""
CREATE INDEX IF NOT EXISTS {index_name}
ON {schema_name}.{table_name}
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64)
""")
)
conn.commit()
logger.info(f"Successfully migrated vector indexes to {target_ext}")
def ensure_text_search_extension(
database_url: str,
text_search_extension: str = "native",
schema: str | None = None,
) -> None:
"""
Ensure the text search columns and indexes match the configured extension.
This function checks the current search_vector column type and index type
in the database and adjusts them if necessary:
- If they match configured extension: no action needed
- If they differ and tables are empty: drop old column/index, recreate with new type
- If they differ and tables have data: raise error with migration guidance
Args:
database_url: SQLAlchemy database URL
text_search_extension: Configured text search extension ("native" or "vchord")
schema: Target PostgreSQL schema name (None for public)
Raises:
RuntimeError: If extension mismatch with existing data
"""
schema_name = schema or "public"
engine = create_engine(database_url)
with engine.connect() as conn:
# Tables with search_vector columns to check
tables_to_check = [
"memory_units",
"reflections", # Renamed from pinned_reflections in p1k2l3m4n5o6 migration
]
# Determine target column type and index type
if text_search_extension == "vchord":
target_column_type = "bm25vector"
target_index_type = "bm25"
else: # native
target_column_type = "tsvector"
target_index_type = "gin"
mismatched_tables = []
tables_with_data = []
for table_name in tables_to_check:
# Check if table exists
table_exists = conn.execute(
text("""
SELECT EXISTS (
SELECT 1 FROM information_schema.tables
WHERE table_schema = :schema AND table_name = :table_name
)
"""),
{"schema": schema_name, "table_name": table_name},
).scalar()
if not table_exists:
logger.debug(f"Table {table_name} does not exist in schema '{schema_name}', skipping")
continue
# Get current column type from information_schema
current_column_info = conn.execute(
text("""
SELECT data_type, udt_name
FROM information_schema.columns
WHERE table_schema = :schema
AND table_name = :table_name
AND column_name = 'search_vector'
"""),
{"schema": schema_name, "table_name": table_name},
).fetchone()
if not current_column_info:
logger.warning(f"No search_vector column found for {table_name}, will create it")
mismatched_tables.append((table_name, None, None))
continue
# Check column type (udt_name contains the actual type: tsvector, bm25vector, etc.)
current_column_type = current_column_info[1] # udt_name
# Get current index type
current_index_info = conn.execute(
text("""
SELECT am.amname
FROM pg_indexes pi
JOIN pg_class c ON c.relname = pi.indexname
JOIN pg_am am ON am.oid = c.relam
WHERE pi.schemaname = :schema
AND pi.tablename = :table_name
AND pi.indexname LIKE '%text_search%'
"""),
{"schema": schema_name, "table_name": table_name},
).fetchone()
current_index_type = current_index_info[0] if current_index_info else None
# Check if column and index types match target
column_matches = current_column_type == target_column_type
index_matches = current_index_type == target_index_type if current_index_type else False
if not (column_matches and index_matches):
logger.info(
f"Text search mismatch on {table_name}: "
f"column={current_column_type} (want {target_column_type}), "
f"index={current_index_type} (want {target_index_type})"
)
mismatched_tables.append((table_name, current_column_type, current_index_type))
# Check if table has data
row_count = conn.execute(text(f"SELECT COUNT(*) FROM {schema_name}.{table_name}")).scalar()
if row_count > 0:
tables_with_data.append((table_name, row_count))
else:
logger.debug(f"Text search OK for {table_name}: {current_column_type}/{current_index_type}")
# If no mismatches, we're done
if not mismatched_tables:
logger.debug(f"All text search columns/indexes match configured extension: {text_search_extension}")
return
# If there's data in any mismatched table, raise error
if tables_with_data:
table_list = ", ".join([f"{table}({count} rows)" for table, count in tables_with_data])
current_ext = "native" if mismatched_tables[0][1] == "tsvector" else "vchord"
raise RuntimeError(
f"Cannot change text search extension from {current_ext} to {text_search_extension}: "
f"the following tables contain data: {table_list}. "
f"To change text search extension, you must either:\n"
f" 1. Clear all data: DELETE FROM {schema_name}.memory_units; "
f"DELETE FROM {schema_name}.reflections; then restart\n"
f" 2. Use the current text search extension (set HINDSIGHT_API_TEXT_SEARCH_EXTENSION='{current_ext}')"
)
# Tables are empty, safe to recreate columns/indexes
logger.info(f"Recreating text search columns/indexes for {text_search_extension}")
for table_name, current_col_type, current_idx_type in mismatched_tables:
# Drop existing index if it exists
if current_idx_type:
logger.info(f"Dropping {current_idx_type} index on {table_name}")
conn.execute(
text(f"""
DROP INDEX IF EXISTS {schema_name}.idx_{table_name.replace(".", "_")}_text_search
""")
)
# Drop existing column if it exists
if current_col_type:
logger.info(f"Dropping {current_col_type} column on {table_name}")
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} DROP COLUMN IF EXISTS search_vector"))
# Create new column with appropriate type
if text_search_extension == "vchord":
logger.info(f"Creating bm25vector column on {table_name}")
# Note: vchord_bm25 extension creates types in bm25_catalog schema
conn.execute(
text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector bm25_catalog.bm25vector")
)
# Create BM25 index
logger.info(f"Creating BM25 index on {table_name}")
conn.execute(
text(f"""
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
ON {schema_name}.{table_name}
USING bm25 (search_vector bm25_catalog.bm25_ops)
""")
)
else: # native
logger.info(f"Creating tsvector column on {table_name}")
# Different GENERATED expression for each table
if table_name == "memory_units":
generated_expr = "to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))"
else: # reflections
generated_expr = "to_tsvector('english', COALESCE(name, '') || ' ' || content)"
conn.execute(
text(f"""
ALTER TABLE {schema_name}.{table_name}
ADD COLUMN search_vector tsvector
GENERATED ALWAYS AS ({generated_expr}) STORED
""")
)
# Create GIN index
logger.info(f"Creating GIN index on {table_name}")
conn.execute(
text(f"""
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
ON {schema_name}.{table_name}
USING gin(search_vector)
""")
)
conn.commit()
logger.info(f"Successfully migrated text search to {text_search_extension}")

View file

@ -209,7 +209,7 @@ class TestLargeBatchRetain:
raise raise
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.timeout(120) @pytest.mark.timeout(240) # Increased timeout for VectorChord BM25 tokenization
async def test_batch_chunking_behavior(self, memory_with_mock_llm, request_context): async def test_batch_chunking_behavior(self, memory_with_mock_llm, request_context):
""" """
Test that large batches are properly chunked into sub-batches. Test that large batches are properly chunked into sub-batches.

View file

@ -57,6 +57,64 @@ hindsight-admin run-db-migration
hindsight-admin run-db-migration --schema tenant_acme hindsight-admin run-db-migration --schema tenant_acme
``` ```
### Vector Extension
| Variable | Description | Default |
|----------|-------------|---------|
| `HINDSIGHT_API_VECTOR_EXTENSION` | Vector extension to use: `auto`, `pgvector`, or `vchord` | `auto` |
Hindsight supports two PostgreSQL vector extensions:
- **pgvector**: Standard extension, works well for most embeddings (up to ~2000 dimensions)
- **vchord**: Optimized for high-dimensional embeddings (3000+ dimensions), includes BM25 search
When set to `auto` (default), Hindsight automatically detects which extension is installed, preferring vchord if both are available.
**When to use vchord:**
- Using high-dimensional embeddings (e.g., `text-embedding-3-large` with 3072 dimensions)
- Need better performance with large embedding dimensions
- Want to use vchord's BM25 search capabilities
**When to use pgvector:**
- Using standard embedding dimensions (384-1536)
- Prefer the widely-adopted pgvector extension
- Simpler deployment (pgvector is more commonly available)
**Switching extensions:**
If you need to switch from one extension to another:
1. Set `HINDSIGHT_API_VECTOR_EXTENSION` to your desired extension (`pgvector` or `vchord`)
2. If your database has existing data, you'll get an error with migration instructions
3. For empty databases, indexes will be automatically recreated on startup
### Text Search Extension
| Variable | Description | Default |
|----------|-------------|---------|
| `HINDSIGHT_API_TEXT_SEARCH_EXTENSION` | Text search backend: `native` or `vchord` | `native` |
Hindsight supports two text search backends for BM25 keyword retrieval:
- **native**: PostgreSQL's built-in full-text search (`tsvector` + GIN indexes)
- **vchord**: VectorChord BM25 (`bm25vector` + BM25 indexes) - requires `vchord_bm25` extension
**When to use vchord:**
- Already using vchord for vector search (good integration)
- Want better BM25 ranking performance
- Need advanced tokenization (uses `llmlingua2` tokenizer)
**When to use native:**
- Standard PostgreSQL deployment (no extra extensions)
- Simpler setup and wider compatibility
- Works well for most use cases
**Switching backends:**
To switch from native to vchord (or vice versa):
1. Set `HINDSIGHT_API_TEXT_SEARCH_EXTENSION=vchord` (or `native`)
2. If your database has existing data, you'll get an error with migration instructions
3. For empty databases, the columns/indexes will be automatically recreated on startup
**Note:** VectorChord text search uses the `llmlingua2` tokenizer for multilingual support, while native uses PostgreSQL's English tokenizer.
### LLM Provider ### LLM Provider
| Variable | Description | Default | | Variable | Description | Default |