From c029807addcf2b84d38d04b52e11414593150dee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Thu, 12 Feb 2026 14:13:04 +0100 Subject: [PATCH] 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 --- .../docker-compose.yaml} | 29 +- .../versions/5a366d414dce_initial_schema.py | 98 ++++- ...k2l3m4_learnings_and_pinned_reflections.py | 110 ++++- hindsight-api/hindsight_api/config.py | 27 ++ .../engine/consolidation/consolidator.py | 35 +- .../hindsight_api/engine/memory_engine.py | 66 +-- .../engine/retain/fact_storage.py | 73 +++- .../hindsight_api/engine/search/retrieval.py | 205 ++++++--- .../hindsight_api/extensions/context.py | 21 +- hindsight-api/hindsight_api/main.py | 2 + hindsight-api/hindsight_api/migrations.py | 395 +++++++++++++++++- hindsight-api/tests/test_load_large_batch.py | 2 +- .../docs/developer/configuration.md | 58 +++ 13 files changed, 917 insertions(+), 204 deletions(-) rename docker/docker-compose/{docker-compose-vchord.yaml => vchord/docker-compose.yaml} (78%) diff --git a/docker/docker-compose/docker-compose-vchord.yaml b/docker/docker-compose/vchord/docker-compose.yaml similarity index 78% rename from docker/docker-compose/docker-compose-vchord.yaml rename to docker/docker-compose/vchord/docker-compose.yaml index 32d91649..6ca29ac6 100644 --- a/docker/docker-compose/docker-compose-vchord.yaml +++ b/docker/docker-compose/vchord/docker-compose.yaml @@ -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 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 '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'; " restart: "no" @@ -65,25 +67,18 @@ services: - "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 + # LLM Configuration (uses OpenAI for testing vchord) + # LLM configuration + HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai} + HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key} # 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 + HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db} + + # Vector and Text Search Extensions + HINDSIGHT_API_VECTOR_EXTENSION: vchord + HINDSIGHT_API_TEXT_SEARCH_EXTENSION: vchord + depends_on: - db networks: diff --git a/hindsight-api/hindsight_api/alembic/versions/5a366d414dce_initial_schema.py b/hindsight-api/hindsight_api/alembic/versions/5a366d414dce_initial_schema.py index 8589593b..f527d04f 100644 --- a/hindsight-api/hindsight_api/alembic/versions/5a366d414dce_initial_schema.py +++ b/hindsight-api/hindsight_api/alembic/versions/5a366d414dce_initial_schema.py @@ -6,6 +6,7 @@ Create Date: 2025-11-27 11:54:19.228030 """ +import os from collections.abc import Sequence import sqlalchemy as sa @@ -23,21 +24,57 @@ 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. + Detect or validate vector extension: 'vchord' or 'pgvector'. + Respects HINDSIGHT_API_VECTOR_EXTENSION env var if set. """ conn = op.get_bind() - vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar() - if vchord_check: + 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() + if not vchord_check: + raise RuntimeError( + "Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;" + ) return "vchord" - - pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar() - if pgvector_check: + 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'") - raise RuntimeError( - "Neither vchord nor pgvector extension found. Install one: CREATE EXTENSION vchord; or CREATE EXTENSION vector;" - ) + +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'" + ) def upgrade() -> None: @@ -185,11 +222,23 @@ def upgrade() -> None: ) # Add search_vector column for full-text search - op.execute(""" - ALTER TABLE memory_units - ADD COLUMN search_vector tsvector - GENERATED ALWAYS AS (to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))) STORED - """) + # 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(""" + ALTER TABLE memory_units + ADD COLUMN search_vector tsvector + GENERATED ALWAYS AS (to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))) STORED + """) op.create_index("idx_memory_units_bank_id", "memory_units", ["bank_id"]) op.create_index("idx_memory_units_document_id", "memory_units", ["document_id"]) @@ -238,11 +287,20 @@ def upgrade() -> None: postgresql_ops={"embedding": "vector_cosine_ops"}, ) - # Create BM25 full-text search index on search_vector - op.execute(""" - CREATE INDEX idx_memory_units_text_search ON memory_units - USING gin(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(""" + CREATE INDEX idx_memory_units_text_search ON memory_units + USING gin(search_vector) + """) op.execute(""" CREATE MATERIALIZED VIEW memory_units_bm25 AS diff --git a/hindsight-api/hindsight_api/alembic/versions/n9i0j1k2l3m4_learnings_and_pinned_reflections.py b/hindsight-api/hindsight_api/alembic/versions/n9i0j1k2l3m4_learnings_and_pinned_reflections.py index b3e2ff5c..69f51d45 100644 --- a/hindsight-api/hindsight_api/alembic/versions/n9i0j1k2l3m4_learnings_and_pinned_reflections.py +++ b/hindsight-api/hindsight_api/alembic/versions/n9i0j1k2l3m4_learnings_and_pinned_reflections.py @@ -10,6 +10,7 @@ This migration: 3. Adds consolidation tracking columns to the 'banks' table """ +import os from collections.abc import Sequence from alembic import context, op @@ -30,21 +31,57 @@ def _get_schema_prefix() -> str: def _detect_vector_extension() -> str: """ - Detect available vector extension: 'vchord' or 'pgvector'. - Prefers vchord if both available. Raises error if neither found. + Detect or validate vector extension: 'vchord' or 'pgvector'. + Respects HINDSIGHT_API_VECTOR_EXTENSION env var if set. """ conn = op.get_bind() - vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar() - if vchord_check: + 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() + if not vchord_check: + raise RuntimeError( + "Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;" + ) return "vchord" - - pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar() - if pgvector_check: + 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'") - raise RuntimeError( - "Neither vchord nor pgvector extension found. Install one: CREATE EXTENSION vchord; or CREATE EXTENSION vector;" - ) + +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'" + ) def upgrade() -> None: @@ -54,6 +91,9 @@ def upgrade() -> None: # Detect which vector extension is available vector_ext = _detect_vector_extension() + # Detect which text search extension to use + text_search_ext = _detect_text_search_extension() + # 1. Create learnings table op.execute(f""" CREATE TABLE {schema}learnings ( @@ -96,11 +136,23 @@ def upgrade() -> None: op.execute(f"CREATE INDEX idx_learnings_tags ON {schema}learnings USING GIN(tags)") # Full-text search for learnings - op.execute(f""" - ALTER TABLE {schema}learnings ADD COLUMN search_vector tsvector - GENERATED ALWAYS AS (to_tsvector('english', text)) STORED - """) - op.execute(f"CREATE INDEX idx_learnings_text_search ON {schema}learnings USING gin(search_vector)") + 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""" + ALTER TABLE {schema}learnings ADD COLUMN search_vector tsvector + GENERATED ALWAYS AS (to_tsvector('english', text)) STORED + """) + op.execute(f"CREATE INDEX idx_learnings_text_search ON {schema}learnings USING gin(search_vector)") # 2. Create pinned_reflections table op.execute(f""" @@ -142,14 +194,26 @@ def upgrade() -> None: op.execute(f"CREATE INDEX idx_pinned_reflections_tags ON {schema}pinned_reflections USING GIN(tags)") # Full-text search for pinned_reflections - op.execute(f""" - ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector tsvector - GENERATED ALWAYS AS (to_tsvector('english', COALESCE(name, '') || ' ' || content)) STORED - """) - op.execute(f""" - CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections - USING gin(search_vector) - """) + 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""" + ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector tsvector + GENERATED ALWAYS AS (to_tsvector('english', COALESCE(name, '') || ' ' || content)) STORED + """) + op.execute(f""" + CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections + USING gin(search_vector) + """) # 3. Add consolidation tracking columns to banks table op.execute(f""" diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 18e1bea5..96c7940c 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/hindsight-api/hindsight_api/config.py @@ -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_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_PORT = "HINDSIGHT_API_PORT" 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_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 DEFAULT_LITELLM_API_BASE = "http://localhost:4000" DEFAULT_EMBEDDINGS_LITELLM_MODEL = "text-embedding-3-small" @@ -460,6 +469,8 @@ class HindsightConfig: # Database database_url: 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_provider: str @@ -687,6 +698,20 @@ class HindsightConfig: def validate(self) -> None: """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 # to ensure the LLM has enough output capacity to extract facts from chunks if self.retain_max_completion_tokens <= self.retain_chunk_size: @@ -712,6 +737,8 @@ class HindsightConfig: # Database database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL), 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_provider=llm_provider, llm_api_key=os.getenv(ENV_LLM_API_KEY), diff --git a/hindsight-api/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api/hindsight_api/engine/consolidation/consolidator.py index f45271a1..2051750b 100644 --- a/hindsight-api/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api/hindsight_api/engine/consolidation/consolidator.py @@ -18,6 +18,7 @@ import uuid from datetime import datetime, timezone from typing import TYPE_CHECKING, Any +from ...config import get_config from ..memory_engine import fq_table from ..retain import embedding_utils from .prompts import ( @@ -1015,15 +1016,33 @@ async def _create_observation_directly( t0 = time.time() observation_id = uuid.uuid4() + + # 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")} ( + id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history, + tags, event_date, occurred_start, occurred_end, mentioned_at + ) + VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10) + RETURNING id + """ + row = await conn.fetchrow( - 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 - ) - VALUES ($1, $2, $3, 'observation', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9, $10) - RETURNING id - """, + query, observation_id, bank_id, observation_text, diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index fef75afe..da9f2a2c 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -968,7 +968,12 @@ class MemoryEngine(MemoryEngineInterface): # Run database migrations if enabled 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: raise ValueError("Database URL is required for migrations") @@ -976,30 +981,43 @@ class MemoryEngine(MemoryEngineInterface): # Migrate all schemas from the tenant extension # The tenant extension is the single source of truth for which schemas exist logger.info("Running database migrations...") - try: - tenants = await self._tenant_extension.list_tenants() - if tenants: - logger.info(f"Running migrations on {len(tenants)} schema(s)...") - for tenant in tenants: - schema = tenant.schema - if schema: - try: - 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") + tenants = await self._tenant_extension.list_tenants() + if tenants: + logger.info(f"Running migrations on {len(tenants)} schema(s)...") + for tenant in tenants: + schema = tenant.schema + if schema: + run_migrations(self.db_url, schema=schema) + logger.info("Schema migrations completed") - # Ensure embedding column dimension matches the model's dimension - # This is done after migrations and after embeddings.initialize() - for tenant in tenants: - schema = tenant.schema - if schema: - try: - ensure_embedding_dimension(self.db_url, self.embeddings.dimension, schema=schema) - except Exception as e: - logger.warning(f"Failed to ensure embedding dimension for schema {schema}: {e}") - except Exception as e: - logger.warning(f"Failed to run schema migrations: {e}") + # Get config for vector extension setting + config = get_config() + + # Ensure embedding column dimension matches the model's dimension + # This is done after migrations and after embeddings.initialize() + for tenant in tenants: + schema = tenant.schema + if schema: + ensure_embedding_dimension( + self.db_url, + self.embeddings.dimension, + schema=schema, + vector_extension=config.vector_extension, + ) + + # 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)}") diff --git a/hindsight-api/hindsight_api/engine/retain/fact_storage.py b/hindsight-api/hindsight_api/engine/retain/fact_storage.py index ef6880b5..17749569 100644 --- a/hindsight-api/hindsight_api/engine/retain/fact_storage.py +++ b/hindsight-api/hindsight_api/engine/retain/fact_storage.py @@ -7,6 +7,7 @@ Handles insertion of facts into the database. import json import logging +from ...config import get_config from ..memory_engine import fq_table from .fact_extraction import _sanitize_text from .types import ProcessedFact @@ -70,28 +71,58 @@ async def insert_facts_batch( # Batch insert all facts # Note: tags are passed as JSON strings and converted back to varchar[] via jsonb_array_elements_text + array_agg - results = await conn.fetch( - 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) - 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[] + # 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""" + 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) ) - FROM input_data - RETURNING id - """, + 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 ( + 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) + 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[] + ) + FROM input_data + RETURNING id + """ + + results = await conn.fetch( + query, bank_id, fact_texts, embeddings, diff --git a/hindsight-api/hindsight_api/engine/search/retrieval.py b/hindsight-api/hindsight_api/engine/search/retrieval.py index 95f207a3..44f8e61d 100644 --- a/hindsight-api/hindsight_api/engine/search/retrieval.py +++ b/hindsight-api/hindsight_api/engine/search/retrieval.py @@ -158,7 +158,7 @@ async def retrieve_bm25( 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 sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower()) @@ -169,29 +169,46 @@ async def retrieve_bm25( # If no valid tokens, return empty results return [] - # Convert query to tsquery using OR for more flexible matching - # This prevents empty results when some terms are missing - query_tsquery = " | ".join(tokens) - + # Build query based on text search backend + config = get_config() tags_clause = build_tags_where_clause_simple(tags, 5) - params = [query_tsquery, bank_id, fact_type, limit] - if tags: - params.append(tags) - results = await conn.fetch( - f""" - 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 - FROM {fq_table("memory_units")} - WHERE bank_id = $2 - AND fact_type = $3 - AND search_vector @@ to_tsquery('english', $1) - {tags_clause} - ORDER BY bm25_score DESC - LIMIT $4 - """, - *params, - ) + 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] + if tags: + params.append(tags) + + query = f""" + 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 + FROM {fq_table("memory_units")} + WHERE bank_id = $2 + AND fact_type = $3 + AND search_vector @@ to_tsquery('english', $1) + {tags_clause} + ORDER BY bm25_score DESC + LIMIT $4 + """ + + results = await conn.fetch(query, *params) return [RetrievalResult.from_db_row(dict(r)) for r in results] @@ -268,59 +285,109 @@ async def retrieve_semantic_bm25_combined( result_dict[ft][0].append(RetrievalResult.from_db_row(row)) 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 tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match) - params = [query_emb_str, bank_id, fact_types, limit, query_tsquery] - if tags: - params.append(tags) + + 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] + 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, + ts_rank_cd(search_vector, to_tsquery('english', $5)) AS bm25_score, + 'bm25' AS source, + ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY ts_rank_cd(search_vector, to_tsquery('english', $5)) DESC) AS rn + FROM {fq_table("memory_units")} + WHERE bank_id = $2 + AND fact_type = ANY($3) + AND search_vector @@ to_tsquery('english', $5) + {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 + """ # 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( - 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, - ts_rank_cd(search_vector, to_tsquery('english', $5)) AS bm25_score, - 'bm25' AS source, - ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY ts_rank_cd(search_vector, to_tsquery('english', $5)) DESC) AS rn - FROM {fq_table("memory_units")} - WHERE bank_id = $2 - AND fact_type = ANY($3) - AND search_vector @@ to_tsquery('english', $5) - {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 - """, - *params, - ) + results = await conn.fetch(query, *params) # Group results by fact_type and source result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types} diff --git a/hindsight-api/hindsight_api/extensions/context.py b/hindsight-api/hindsight_api/extensions/context.py index e470f60b..ad702dde 100644 --- a/hindsight-api/hindsight_api/extensions/context.py +++ b/hindsight-api/hindsight_api/extensions/context.py @@ -96,7 +96,13 @@ class DefaultExtensionContext(ExtensionContext): async def run_migration(self, schema: str) -> None: """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) db_url = self._database_url @@ -107,6 +113,9 @@ class DefaultExtensionContext(ExtensionContext): run_migrations(db_url, schema=schema) + # Get config for vector extension setting + config = get_config() + # Ensure embedding column dimension matches the model's dimension # This is needed because migrations create columns with default dimension if self._memory_engine is not None: @@ -114,7 +123,15 @@ class DefaultExtensionContext(ExtensionContext): if embeddings is not None: dimension = getattr(embeddings, "dimension", 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": """Get the memory engine interface.""" diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index b69d891b..b77ca7e0 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/hindsight-api/hindsight_api/main.py @@ -155,6 +155,8 @@ def main(): config = HindsightConfig( database_url=config.database_url, database_schema=config.database_schema, + vector_extension=config.vector_extension, + text_search_extension=config.text_search_extension, llm_provider=config.llm_provider, llm_api_key=config.llm_api_key, llm_model=config.llm_model, diff --git a/hindsight-api/hindsight_api/migrations.py b/hindsight-api/hindsight_api/migrations.py index 14a8c7ee..02642614 100644 --- a/hindsight-api/hindsight_api/migrations.py +++ b/hindsight-api/hindsight_api/migrations.py @@ -33,35 +33,39 @@ logger = logging.getLogger(__name__) 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'. - Prefers vchord if both available. Raises error if neither found. + Validate vector extension: 'vchord' or 'pgvector'. Args: conn: SQLAlchemy connection object + vector_extension: Configured extension ("pgvector" or "vchord") Returns: "vchord" or "pgvector" Raises: - RuntimeError: If neither extension is installed + RuntimeError: If configured extension is not 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") + # Verify the configured extension is installed + if vector_extension == "vchord": + vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar() + if not vchord_check: + raise RuntimeError( + "Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;" + ) + logger.debug("Using configured 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") + 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" - - raise RuntimeError( - "Neither vchord nor pgvector extension found. Install one: CREATE EXTENSION vchord; or CREATE EXTENSION vector;" - ) + else: + raise ValueError(f"Invalid vector_extension: {vector_extension}. Must be 'pgvector' or 'vchord'") def _get_schema_lock_id(schema: str) -> int: @@ -355,6 +359,7 @@ def ensure_embedding_dimension( database_url: str, required_dimension: int, schema: str | None = None, + vector_extension: str = "pgvector", ) -> None: """ Ensure the embedding column dimension matches the model's dimension. @@ -369,6 +374,7 @@ def ensure_embedding_dimension( database_url: SQLAlchemy database URL required_dimension: The embedding dimension required by the model schema: Target PostgreSQL schema name (None for public) + vector_extension: Configured vector extension ("pgvector" or "vchord") Raises: RuntimeError: If dimension mismatch with existing data @@ -393,8 +399,8 @@ def ensure_embedding_dimension( return # Detect which vector extension is available - vector_ext = _detect_vector_extension(conn) - logger.info(f"Detected vector extension: {vector_ext}") + vector_ext = _detect_vector_extension(conn, vector_extension) + logger.info(f"Using vector extension: {vector_ext}") # Get current column dimension from pg_attribute # pgvector stores dimension in atttypmod @@ -491,3 +497,354 @@ def ensure_embedding_dimension( conn.commit() 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}") diff --git a/hindsight-api/tests/test_load_large_batch.py b/hindsight-api/tests/test_load_large_batch.py index 9a413948..4f499aa1 100644 --- a/hindsight-api/tests/test_load_large_batch.py +++ b/hindsight-api/tests/test_load_large_batch.py @@ -209,7 +209,7 @@ class TestLargeBatchRetain: raise @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): """ Test that large batches are properly chunked into sub-batches. diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 34336051..ebd091e4 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -57,6 +57,64 @@ hindsight-admin run-db-migration 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 | Variable | Description | Default |