Compare commits

...
Author SHA1 Message Date
Nicolò Boschi d913f0c1dc refactor: deduplicate text search query in retrieve_semantic_bm25_combined
Instead of maintaining 3 complete query copies (native, vchord, pg_textsearch),
now we:
- Build backend-specific parts (score_expr, order_by, where_filter)
- Use a single query template with injected backend-specific parts

This makes maintenance easier - changes to the semantic CTE or overall structure
only need to be made once.
2026-02-12 17:21:56 +01:00
Nicolò Boschi 0d4aca1803 feat: support timescale pg_textsearch as text search extension 2026-02-12 17:19:33 +01:00
10 changed files with 301 additions and 104 deletions
@@ -0,0 +1,32 @@
# PostgreSQL with pgvector and pg_textsearch extensions
# Note: pg_textsearch requires PostgreSQL 17+
FROM postgres:17
# Install build dependencies
RUN apt-get update && apt-get install -y \
build-essential \
git \
postgresql-server-dev-17 \
libpq-dev \
&& rm -rf /var/lib/apt/lists/*
# Install pgvector
RUN cd /tmp && \
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
cd pgvector && \
make && \
make install
# Install pg_textsearch
RUN cd /tmp && \
git clone https://github.com/timescale/pg_textsearch.git && \
cd pg_textsearch && \
make && \
make install
# Clean up source files and build dependencies
RUN rm -rf /tmp/pgvector /tmp/pg_textsearch && \
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
# Ensure extensions are preloaded
RUN echo "shared_preload_libraries = 'pg_textsearch'" >> /usr/share/postgresql/postgresql.conf.sample
@@ -0,0 +1,91 @@
name: hindsight
# Docker Compose file for Hindsight with PostgreSQL and Timescale pg_textsearch
# docker compose -f docker/docker-compose/pg_textsearch/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/pg_textsearch/docker-compose.yaml up -d
# Make sure to set the required environment variables before running:
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
# - Configure LLM provider variables as needed (see below in the hindsight service)
#
# Usage:
# docker compose up -d
#
# Optional environment variables with defaults:
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
services:
db:
# Use custom PostgreSQL image with pgvector and pg_textsearch extensions
build:
context: .
dockerfile: Dockerfile
container_name: hindsight-db
restart: always
# Expose PostgreSQL port
ports:
- "5437:5432"
environment:
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
volumes:
- pg_data:/var/lib/postgresql/data
networks:
- hindsight-net
pg-textsearch-init:
build:
context: .
dockerfile: Dockerfile
depends_on:
- db
environment:
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
command: >
bash -c "
echo 'Waiting for PostgreSQL to be ready...';
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
echo 'PostgreSQL is unavailable - sleeping';
sleep 2;
done;
echo 'PostgreSQL is ready - creating hindsight_db database';
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
echo 'Creating extensions in hindsight_db database';
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_textsearch CASCADE;';
echo 'Database and extensions created successfully';
"
restart: "no"
networks:
- hindsight-net
hindsight:
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
container_name: hindsight-app
ports:
- "8888:8888"
- "9999:9999"
environment:
# LLM Configuration
HINDSIGHT_API_LLM_PROVIDER: ${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}
# Vector and Text Search Extensions
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_textsearch
depends_on:
- db
networks:
- hindsight-net
networks:
hindsight-net:
driver: bridge
volumes:
pg_data:
@@ -51,7 +51,7 @@ def _detect_vector_extension() -> str:
def _detect_text_search_extension() -> str:
"""
Detect or validate text search extension: 'native' or 'vchord'.
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
Creates the extension if needed.
"""
@@ -69,11 +69,23 @@ def _detect_text_search_extension() -> str:
# Extension truly doesn't exist - re-raise the error
raise
return "vchord"
elif text_search_extension == "pg_textsearch":
# Create pg_textsearch extension if not exists
try:
op.execute("CREATE EXTENSION IF NOT EXISTS pg_textsearch 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 = 'pg_textsearch'")).fetchone()
if not result:
# Extension truly doesn't exist - re-raise the error
raise
return "pg_textsearch"
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'"
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
)
@@ -232,6 +244,12 @@ def upgrade() -> None:
ALTER TABLE memory_units
ADD COLUMN search_vector bm25_catalog.bm25vector
""")
elif text_search_ext == "pg_textsearch":
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
op.execute("""
ALTER TABLE memory_units
ADD COLUMN search_vector TEXT
""")
else: # native
# Native PostgreSQL: tsvector with automatic generation
op.execute("""
@@ -295,6 +313,14 @@ def upgrade() -> None:
CREATE INDEX idx_memory_units_text_search ON memory_units
USING bm25 (search_vector bm25_catalog.bm25_ops)
""")
elif text_search_ext == "pg_textsearch":
# Timescale pg_textsearch BM25 index on text column
# Note: pg_textsearch doesn't support expressions, so we index the main text column
op.execute("""
CREATE INDEX idx_memory_units_text_search ON memory_units
USING bm25(text)
WITH (text_config='english')
""")
else: # native
# Native PostgreSQL GIN index
op.execute("""
@@ -58,7 +58,7 @@ def _detect_vector_extension() -> str:
def _detect_text_search_extension() -> str:
"""
Detect or validate text search extension: 'native' or 'vchord'.
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
Creates the extension if needed.
"""
@@ -76,11 +76,23 @@ def _detect_text_search_extension() -> str:
# Extension truly doesn't exist - re-raise the error
raise
return "vchord"
elif text_search_extension == "pg_textsearch":
# Create pg_textsearch extension if not exists
try:
op.execute("CREATE EXTENSION IF NOT EXISTS pg_textsearch 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 = 'pg_textsearch'")).fetchone()
if not result:
# Extension truly doesn't exist - re-raise the error
raise
return "pg_textsearch"
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'"
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
)
@@ -146,6 +158,15 @@ def upgrade() -> None:
CREATE INDEX idx_learnings_text_search ON {schema}learnings
USING bm25 (search_vector bm25_catalog.bm25_ops)
""")
elif text_search_ext == "pg_textsearch":
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
op.execute(f"""
ALTER TABLE {schema}learnings ADD COLUMN search_vector TEXT
""")
op.execute(f"""
CREATE INDEX idx_learnings_text_search ON {schema}learnings
USING bm25(text) WITH (text_config='english')
""")
else: # native
# Native PostgreSQL: tsvector with automatic generation
op.execute(f"""
@@ -204,6 +225,16 @@ def upgrade() -> None:
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
USING bm25 (search_vector bm25_catalog.bm25_ops)
""")
elif text_search_ext == "pg_textsearch":
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
op.execute(f"""
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector TEXT
""")
op.execute(f"""
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
USING bm25(content)
WITH (text_config='english')
""")
else: # native
# Native PostgreSQL: tsvector with automatic generation
op.execute(f"""
+3 -3
View File
@@ -329,8 +329,8 @@ 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"
# Text search extension (native PostgreSQL, vchord BM25, or Timescale pg_textsearch)
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch"
# LiteLLM defaults
DEFAULT_LITELLM_API_BASE = "http://localhost:4000"
@@ -706,7 +706,7 @@ class HindsightConfig:
)
# Validate text_search_extension
valid_text_search = ("native", "vchord")
valid_text_search = ("native", "vchord", "pg_textsearch")
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)}"
@@ -1030,8 +1030,9 @@ async def _create_observation_directly(
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
RETURNING id
"""
else: # native
else: # native or pg_textsearch
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
query = f"""
INSERT INTO {fq_table("memory_units")} (
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
@@ -97,8 +97,9 @@ async def insert_facts_batch(
FROM input_data
RETURNING id
"""
else: # native
else: # native or pg_textsearch
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
query = f"""
WITH input_data AS (
SELECT * FROM unnest(
@@ -291,99 +291,74 @@ async def retrieve_semantic_bm25_combined(
# Build tags clause - param 6 if tags provided
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
# Build backend-specific BM25 parts
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)
bm25_score_expr = "search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($5, 'llmlingua2'))"
bm25_order_by = f"{bm25_score_expr} DESC"
bm25_where_filter = "" # No additional WHERE filter for vchord
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
"""
elif config.text_search_extension == "pg_textsearch":
# Timescale pg_textsearch: use <@> operator with to_bm25query
# Note: pg_textsearch scores are negative (lower/more negative = better, so -10 > -1)
# We negate the score to maintain API consistency (higher = better)
bm25_score_expr = "-(text <@> to_bm25query($5, 'idx_memory_units_text_search'))"
bm25_order_by = "text <@> to_bm25query($5, 'idx_memory_units_text_search') ASC"
bm25_where_filter = "" # No additional WHERE filter for pg_textsearch
params = [query_emb_str, bank_id, fact_types, limit, query_text]
else: # native
# Native PostgreSQL: use ts_rank_cd with to_tsquery
query_tsquery = " | ".join(tokens)
bm25_score_expr = "ts_rank_cd(search_vector, to_tsquery('english', $5))"
bm25_order_by = f"{bm25_score_expr} DESC"
bm25_where_filter = "AND search_vector @@ to_tsquery('english', $5)"
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
"""
if tags:
params.append(tags)
# Single query template with backend-specific parts injected
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,
{bm25_score_expr} AS bm25_score,
'bm25' AS source,
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY {bm25_order_by}) AS rn
FROM {fq_table("memory_units")}
WHERE bank_id = $2
AND fact_type = ANY($3)
{bm25_where_filter}
{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
+34 -1
View File
@@ -688,6 +688,9 @@ def ensure_text_search_extension(
if text_search_extension == "vchord":
target_column_type = "bm25vector"
target_index_type = "bm25"
elif text_search_extension == "pg_textsearch":
target_column_type = "text"
target_index_type = "bm25"
else: # native
target_column_type = "tsvector"
target_index_type = "gin"
@@ -775,7 +778,16 @@ def ensure_text_search_extension(
# 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"
# Detect current extension from column type
current_col_type = mismatched_tables[0][1]
if current_col_type == "tsvector":
current_ext = "native"
elif current_col_type == "bm25vector":
current_ext = "vchord"
elif current_col_type == "text":
current_ext = "pg_textsearch"
else:
current_ext = "unknown"
raise RuntimeError(
f"Cannot change text search extension from {current_ext} to {text_search_extension}: "
f"the following tables contain data: {table_list}. "
@@ -820,6 +832,27 @@ def ensure_text_search_extension(
USING bm25 (search_vector bm25_catalog.bm25_ops)
""")
)
elif text_search_extension == "pg_textsearch":
logger.info(f"Creating TEXT column on {table_name}")
# Dummy TEXT column for consistency (indexes operate on base columns)
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
# Create BM25 index on expression
logger.info(f"Creating BM25 index on {table_name}")
# Different expression for each table
if table_name == "memory_units":
index_expr = "(COALESCE(text, '') || ' ' || COALESCE(context, ''))"
else: # reflections
index_expr = "(COALESCE(name, '') || ' ' || content)"
conn.execute(
text(f"""
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
ON {schema_name}.{table_name}
USING bm25({index_expr})
WITH (text_config='english')
""")
)
else: # native
logger.info(f"Creating tsvector column on {table_name}")
# Different GENERATED expression for each table
+17 -10
View File
@@ -90,30 +90,37 @@ If you need to switch from one extension to another:
| Variable | Description | Default |
|----------|-------------|---------|
| `HINDSIGHT_API_TEXT_SEARCH_EXTENSION` | Text search backend: `native` or `vchord` | `native` |
| `HINDSIGHT_API_TEXT_SEARCH_EXTENSION` | Text search backend: `native`, `vchord`, or `pg_textsearch` | `native` |
Hindsight supports two text search backends for BM25 keyword retrieval:
Hindsight supports three 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)
- **pg_textsearch**: Timescale BM25 (text columns + BM25 indexes) - requires `pg_textsearch` extension
**When to use native:**
- Standard PostgreSQL deployment (no extra extensions)
- Simpler setup and wider compatibility
- Works well for most use cases
**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 pg_textsearch:**
- Want industry-standard BM25 ranking with better relevance than native PostgreSQL
- Need efficient top-K queries with Block-Max WAND optimization
- Prefer lower memory footprint compared to vchord
- Already using Timescale or have `pg_textsearch` available
**Switching backends:**
To switch from native to vchord (or vice versa):
1. Set `HINDSIGHT_API_TEXT_SEARCH_EXTENSION=vchord` (or `native`)
To switch between backends:
1. Set `HINDSIGHT_API_TEXT_SEARCH_EXTENSION` to your desired backend (`native`, `vchord`, or `pg_textsearch`)
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.
**Note:** VectorChord uses the `llmlingua2` tokenizer for multilingual support, while native and pg_textsearch use PostgreSQL's English tokenizer.
### LLM Provider