From 1c9622b7002b0a8eb85d055cefe41abfb84fc346 Mon Sep 17 00:00:00 2001 From: Alexandru Date: Thu, 12 Mar 2026 13:59:31 +0100 Subject: [PATCH 1/2] add configurable full-text search language --- py/all_possible_config.toml | 1 + py/core/base/providers/database.py | 1 + .../configs/r2r_azure_with_test_limits.toml | 1 + py/core/providers/database/chunks.py | 27 ++++--- py/core/providers/database/documents.py | 18 +++-- py/core/providers/database/postgres.py | 3 + py/core/providers/database/utils.py | 24 +++++- py/r2r/r2r.toml | 1 + py/tests/unit/conftest.py | 2 + .../test_chunks_handler_semantic_search.py | 2 +- .../test_full_text_search_language.py | 79 +++++++++++++++++++ 11 files changed, 140 insertions(+), 19 deletions(-) create mode 100644 py/tests/unit/database/test_full_text_search_language.py diff --git a/py/all_possible_config.toml b/py/all_possible_config.toml index 6929872cd0..3486c68f1c 100644 --- a/py/all_possible_config.toml +++ b/py/all_possible_config.toml @@ -115,6 +115,7 @@ default_collection_description = "Your default collection." collection_summary_system_prompt = "system" collection_summary_prompt = "collection_summary" disable_create_extension = false +full_text_search_language = "english" # PostgreSQL tuning settings [database.postgres_configuration_settings] diff --git a/py/core/base/providers/database.py b/py/core/base/providers/database.py index af754bba4f..1cd9d642df 100644 --- a/py/core/base/providers/database.py +++ b/py/core/base/providers/database.py @@ -140,6 +140,7 @@ class DatabaseConfig(ProviderConfig): collection_summary_system_prompt: str = "system" collection_summary_prompt: str = "collection_summary" disable_create_extension: bool = False + full_text_search_language: str = "english" # Graph settings batch_size: Optional[int] = 1 diff --git a/py/core/configs/r2r_azure_with_test_limits.toml b/py/core/configs/r2r_azure_with_test_limits.toml index 120ab94841..6edf2daf60 100644 --- a/py/core/configs/r2r_azure_with_test_limits.toml +++ b/py/core/configs/r2r_azure_with_test_limits.toml @@ -25,6 +25,7 @@ base_dimension = 512 base_model = "openai/text-embedding-3-small" [database] +full_text_search_language = "english" [database.limits] global_per_min = 10 # Small enough to test quickly monthly_limit = 20 # Small enough to test in one run diff --git a/py/core/providers/database/chunks.py b/py/core/providers/database/chunks.py index 016ccb8e64..ab276d0ae6 100644 --- a/py/core/providers/database/chunks.py +++ b/py/core/providers/database/chunks.py @@ -26,7 +26,7 @@ from .base import PostgresConnectionManager from .filters import apply_filters -from .utils import psql_quote_literal +from .utils import psql_quote_literal, psql_regconfig_literal logger = logging.getLogger() @@ -81,10 +81,15 @@ def __init__( connection_manager: PostgresConnectionManager, dimension: int | float, quantization_type: VectorQuantizationType, + full_text_search_language: str = "english", ): super().__init__(project_name, connection_manager) self.dimension = dimension self.quantization_type = quantization_type + self.full_text_search_language = full_text_search_language + self.full_text_search_regconfig = psql_regconfig_literal( + full_text_search_language + ) async def create_tables(self): # First check if table already exists and validate dimensions @@ -172,12 +177,12 @@ async def create_tables(self): {binary_col} text TEXT, metadata JSONB, - fts tsvector GENERATED ALWAYS AS (to_tsvector('english', text)) STORED + fts tsvector GENERATED ALWAYS AS (to_tsvector({self.full_text_search_regconfig}, text)) STORED ); CREATE INDEX IF NOT EXISTS idx_vectors_document_id ON {self._get_table_name(PostgresChunksHandler.TABLE_NAME)} (document_id); CREATE INDEX IF NOT EXISTS idx_vectors_owner_id ON {self._get_table_name(PostgresChunksHandler.TABLE_NAME)} (owner_id); CREATE INDEX IF NOT EXISTS idx_vectors_collection_ids ON {self._get_table_name(PostgresChunksHandler.TABLE_NAME)} USING GIN (collection_ids); - CREATE INDEX IF NOT EXISTS idx_vectors_text ON {self._get_table_name(PostgresChunksHandler.TABLE_NAME)} USING GIN (to_tsvector('english', text)); + CREATE INDEX IF NOT EXISTS idx_vectors_text ON {self._get_table_name(PostgresChunksHandler.TABLE_NAME)} USING GIN (to_tsvector({self.full_text_search_regconfig}, text)); """ await self.connection_manager.execute_query(query) @@ -554,7 +559,9 @@ async def full_text_search( conditions = [] params: list[str | int | bytes] = [query_text] - conditions.append("fts @@ websearch_to_tsquery('english', $1)") + conditions.append( + f"fts @@ websearch_to_tsquery({self.full_text_search_regconfig}, $1)" + ) if search_settings.filters: filter_condition, params = apply_filters( @@ -573,7 +580,7 @@ async def full_text_search( collection_ids, text, metadata, - ts_rank(fts, websearch_to_tsquery('english', $1), 32) as rank + ts_rank(fts, websearch_to_tsquery({self.full_text_search_regconfig}, $1), 32) as rank FROM {self._get_table_name(PostgresChunksHandler.TABLE_NAME)} {where_clause} ORDER BY rank DESC @@ -1257,8 +1264,8 @@ async def search_documents( CASE WHEN $1 = '' THEN 0.0 ELSE ts_rank_cd( - setweight(to_tsvector('english', {metadata_fields_expr}), 'A'), - websearch_to_tsquery('english', $1), + setweight(to_tsvector({self.full_text_search_regconfig}, {metadata_fields_expr}), 'A'), + websearch_to_tsquery({self.full_text_search_regconfig}, $1), 32 ) END as metadata_rank @@ -1272,14 +1279,14 @@ async def search_documents( document_id, AVG( ts_rank_cd( - setweight(to_tsvector('english', COALESCE(text, '')), 'B'), - websearch_to_tsquery('english', $1), + setweight(to_tsvector({self.full_text_search_regconfig}, COALESCE(text, '')), 'B'), + websearch_to_tsquery({self.full_text_search_regconfig}, $1), 32 ) ) as body_rank FROM {self._get_table_name(PostgresChunksHandler.TABLE_NAME)} WHERE $1 != '' - {"AND to_tsvector('english', text) @@ websearch_to_tsquery('english', $1)" if search_over_body else ""} + {f"AND to_tsvector({self.full_text_search_regconfig}, text) @@ websearch_to_tsquery({self.full_text_search_regconfig}, $1)" if search_over_body else ""} GROUP BY document_id ), -- Combined scores with document metadata diff --git a/py/core/providers/database/documents.py b/py/core/providers/database/documents.py index 6a8ff003bb..85b18e9983 100644 --- a/py/core/providers/database/documents.py +++ b/py/core/providers/database/documents.py @@ -24,6 +24,7 @@ from .base import PostgresConnectionManager from .filters import apply_filters +from .utils import psql_regconfig_literal logger = logging.getLogger() @@ -75,9 +76,14 @@ def __init__( project_name: str, connection_manager: PostgresConnectionManager, dimension: int | float, + full_text_search_language: str = "english", ): self.dimension = dimension super().__init__(project_name, connection_manager) + self.full_text_search_language = full_text_search_language + self.full_text_search_regconfig = psql_regconfig_literal( + full_text_search_language + ) async def create_tables(self): logger.info( @@ -108,9 +114,9 @@ async def create_tables(self): updated_at TIMESTAMPTZ DEFAULT NOW(), ingestion_attempt_number INT DEFAULT 0, raw_tsvector tsvector GENERATED ALWAYS AS ( - setweight(to_tsvector('english', COALESCE(title, '')), 'A') || - setweight(to_tsvector('english', COALESCE(summary, '')), 'B') || - setweight(to_tsvector('english', COALESCE((metadata->>'description')::text, '')), 'C') + setweight(to_tsvector({self.full_text_search_regconfig}, COALESCE(title, '')), 'A') || + setweight(to_tsvector({self.full_text_search_regconfig}, COALESCE(summary, '')), 'B') || + setweight(to_tsvector({self.full_text_search_regconfig}, COALESCE((metadata->>'description')::text, '')), 'C') ) STORED, total_tokens INT DEFAULT 0 ); @@ -846,7 +852,9 @@ async def full_text_document_search( ) -> list[DocumentResponse]: """Enhanced full-text search using generated tsvector.""" - where_clauses = ["raw_tsvector @@ websearch_to_tsquery('english', $1)"] + where_clauses = [ + f"raw_tsvector @@ websearch_to_tsquery({self.full_text_search_regconfig}, $1)" + ] params: list[str | int | bytes] = [query_text] filters = copy.deepcopy(search_settings.filters) @@ -877,7 +885,7 @@ async def full_text_document_search( summary, summary_embedding, total_tokens, - ts_rank_cd(raw_tsvector, websearch_to_tsquery('english', $1), 32) as text_score + ts_rank_cd(raw_tsvector, websearch_to_tsquery({self.full_text_search_regconfig}, $1), 32) as text_score FROM {self._get_table_name(PostgresDocumentsHandler.TABLE_NAME)} WHERE {where_clause} ORDER BY text_score DESC diff --git a/py/core/providers/database/postgres.py b/py/core/providers/database/postgres.py index b921316df4..0858c63f6a 100644 --- a/py/core/providers/database/postgres.py +++ b/py/core/providers/database/postgres.py @@ -130,6 +130,7 @@ def __init__( self.default_collection_description = ( config.default_collection_description ) + self.full_text_search_language = config.full_text_search_language self.connection_manager: PostgresConnectionManager = ( PostgresConnectionManager() @@ -138,6 +139,7 @@ def __init__( project_name=self.project_name, connection_manager=self.connection_manager, dimension=self.dimension, + full_text_search_language=self.full_text_search_language, ) self.token_handler = PostgresTokensHandler( self.project_name, self.connection_manager @@ -153,6 +155,7 @@ def __init__( connection_manager=self.connection_manager, dimension=self.dimension, quantization_type=(self.quantization_type), + full_text_search_language=self.full_text_search_language, ) self.conversations_handler = PostgresConversationsHandler( self.project_name, self.connection_manager diff --git a/py/core/providers/database/utils.py b/py/core/providers/database/utils.py index fce7662d97..eb293bc342 100644 --- a/py/core/providers/database/utils.py +++ b/py/core/providers/database/utils.py @@ -1,6 +1,11 @@ -""" -Database utility functions for PostgreSQL operations. -""" +"""Database utility functions for PostgreSQL operations.""" + +import re + + +_REGCONFIG_NAME_PATTERN = re.compile( + r"^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)?$" +) def psql_quote_literal(value: str) -> str: @@ -10,3 +15,16 @@ def psql_quote_literal(value: str) -> str: or your database driver's quoting functions. """ return "'" + value.replace("'", "''") + "'" + + +def psql_regconfig_literal(value: str) -> str: + """Return a validated PostgreSQL regconfig literal. + + Accepts standard text search configuration names like ``english`` or + schema-qualified names like ``pg_catalog.english``. + """ + if not _REGCONFIG_NAME_PATTERN.fullmatch(value): + raise ValueError( + f"Invalid PostgreSQL text search configuration: {value}" + ) + return f"{psql_quote_literal(value)}::regconfig" diff --git a/py/r2r/r2r.toml b/py/r2r/r2r.toml index 7e55b292d3..acdedd7522 100644 --- a/py/r2r/r2r.toml +++ b/py/r2r/r2r.toml @@ -65,6 +65,7 @@ provider = "postgres" default_collection_name = "Default" default_collection_description = "Your default collection." collection_summary_prompt = "collection_summary" +full_text_search_language = "english" [database.graph_creation_settings] graph_entity_description_prompt = "graph_entity_description" diff --git a/py/tests/unit/conftest.py b/py/tests/unit/conftest.py index ac74cd7719..a98b26ac9f 100644 --- a/py/tests/unit/conftest.py +++ b/py/tests/unit/conftest.py @@ -68,6 +68,7 @@ async def chunks_handler(db_provider): connection_manager=connection_manager, dimension=dimension, quantization_type=quantization_type, + full_text_search_language=db_provider.config.full_text_search_language, ) await handler.create_tables() return handler @@ -108,6 +109,7 @@ async def documents_handler(db_provider): project_name=project_name, connection_manager=connection_manager, dimension=dimension, + full_text_search_language=db_provider.config.full_text_search_language, ) await handler.create_tables() return handler diff --git a/py/tests/unit/database/test_chunks_handler_semantic_search.py b/py/tests/unit/database/test_chunks_handler_semantic_search.py index 9d8dcb0641..4c163d70c3 100644 --- a/py/tests/unit/database/test_chunks_handler_semantic_search.py +++ b/py/tests/unit/database/test_chunks_handler_semantic_search.py @@ -59,7 +59,7 @@ async def test_semantic_search_keeps_direct_knn_path_without_filters(): query = fetch_query.await_args.args[0] assert "WITH filtered AS MATERIALIZED" not in query assert "FROM filtered" not in query - assert "FROM test_project.chunks" in query + assert 'FROM "test_project"."chunks"' in query @pytest.mark.asyncio diff --git a/py/tests/unit/database/test_full_text_search_language.py b/py/tests/unit/database/test_full_text_search_language.py new file mode 100644 index 0000000000..2392680670 --- /dev/null +++ b/py/tests/unit/database/test_full_text_search_language.py @@ -0,0 +1,79 @@ +from types import SimpleNamespace +from typing import cast +from unittest.mock import AsyncMock + +import pytest + +from core.base import DatabaseConfig, SearchSettings, VectorQuantizationType +from core.providers.database.base import PostgresConnectionManager +from core.providers.database.chunks import PostgresChunksHandler +from core.providers.database.documents import PostgresDocumentsHandler +from core.providers.database.utils import psql_regconfig_literal + + +def test_database_config_defaults_full_text_language_to_english(): + assert DatabaseConfig().full_text_search_language == "english" + + +def test_psql_regconfig_literal_rejects_invalid_names(): + assert psql_regconfig_literal("simple") == "'simple'::regconfig" + + with pytest.raises(ValueError): + psql_regconfig_literal("english'; DROP TABLE chunks; --") + + +@pytest.mark.asyncio +async def test_chunks_handler_uses_configured_full_text_language(): + fetch_query = AsyncMock(side_effect=[[], [], []]) + execute_query = AsyncMock() + connection_manager = cast( + PostgresConnectionManager, + SimpleNamespace(fetch_query=fetch_query, execute_query=execute_query), + ) + handler = PostgresChunksHandler( + project_name="test_project", + connection_manager=connection_manager, + dimension=4, + quantization_type=VectorQuantizationType.FP32, + full_text_search_language="simple", + ) + + await handler.create_tables() + await handler.full_text_search("foo", SearchSettings()) + + assert execute_query.await_args is not None + create_query = execute_query.await_args.args[0] + assert "to_tsvector('simple'::regconfig, text)" in create_query + + assert fetch_query.await_args is not None + search_query = fetch_query.await_args.args[0] + assert "websearch_to_tsquery('simple'::regconfig, $1)" in search_query + + +@pytest.mark.asyncio +async def test_documents_handler_uses_configured_full_text_language(): + fetch_query = AsyncMock( + side_effect=[[{"column_name": "total_tokens"}], []] + ) + execute_query = AsyncMock() + connection_manager = cast( + PostgresConnectionManager, + SimpleNamespace(fetch_query=fetch_query, execute_query=execute_query), + ) + handler = PostgresDocumentsHandler( + project_name="test_project", + connection_manager=connection_manager, + dimension=4, + full_text_search_language="simple", + ) + + await handler.create_tables() + await handler.full_text_document_search("foo", SearchSettings()) + + assert execute_query.await_args is not None + create_query = execute_query.await_args.args[0] + assert "to_tsvector('simple'::regconfig, COALESCE(title, ''))" in create_query + + assert fetch_query.await_args is not None + search_query = fetch_query.await_args.args[0] + assert "websearch_to_tsquery('simple'::regconfig, $1)" in search_query From 3e59d0cd01f7ba385cdda535e0beecece4541631 Mon Sep 17 00:00:00 2001 From: Alexandru Date: Thu, 12 Mar 2026 14:19:24 +0100 Subject: [PATCH 2/2] apply ruff formatting fixes --- py/core/providers/database/utils.py | 1 - .../unit/database/test_chunks_handler_semantic_search.py | 8 ++++++-- py/tests/unit/database/test_full_text_search_language.py | 4 +++- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/py/core/providers/database/utils.py b/py/core/providers/database/utils.py index eb293bc342..dc25e2f7a4 100644 --- a/py/core/providers/database/utils.py +++ b/py/core/providers/database/utils.py @@ -2,7 +2,6 @@ import re - _REGCONFIG_NAME_PATTERN = re.compile( r"^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)?$" ) diff --git a/py/tests/unit/database/test_chunks_handler_semantic_search.py b/py/tests/unit/database/test_chunks_handler_semantic_search.py index 4c163d70c3..d18440d0bd 100644 --- a/py/tests/unit/database/test_chunks_handler_semantic_search.py +++ b/py/tests/unit/database/test_chunks_handler_semantic_search.py @@ -23,7 +23,9 @@ async def test_semantic_search_materializes_filtered_fp32_results(): quantization_type=VectorQuantizationType.FP32, ) settings = SearchSettings( - filters={"document_id": {"$in": ["11111111-1111-4111-8111-111111111111"]}}, + filters={ + "document_id": {"$in": ["11111111-1111-4111-8111-111111111111"]} + }, limit=10, offset=0, ) @@ -76,7 +78,9 @@ async def test_semantic_search_materializes_filtered_int1_candidates(): quantization_type=VectorQuantizationType.INT1, ) settings = SearchSettings( - filters={"document_id": {"$in": ["11111111-1111-4111-8111-111111111111"]}}, + filters={ + "document_id": {"$in": ["11111111-1111-4111-8111-111111111111"]} + }, limit=10, offset=0, ) diff --git a/py/tests/unit/database/test_full_text_search_language.py b/py/tests/unit/database/test_full_text_search_language.py index 2392680670..f78e07d61c 100644 --- a/py/tests/unit/database/test_full_text_search_language.py +++ b/py/tests/unit/database/test_full_text_search_language.py @@ -72,7 +72,9 @@ async def test_documents_handler_uses_configured_full_text_language(): assert execute_query.await_args is not None create_query = execute_query.await_args.args[0] - assert "to_tsvector('simple'::regconfig, COALESCE(title, ''))" in create_query + assert ( + "to_tsvector('simple'::regconfig, COALESCE(title, ''))" in create_query + ) assert fetch_query.await_args is not None search_query = fetch_query.await_args.args[0]