From 94b5cf4fa651742fc6b3ad4c2f0d03614b5798cb Mon Sep 17 00:00:00 2001 From: are-ces <195810094+are-ces@users.noreply.github.com> Date: Fri, 5 Jun 2026 14:12:16 +0200 Subject: [PATCH] LCORE-2437: support pgvector in BYOK RAG enrichment Add pgvector as a supported rag_type in the BYOK RAG enrichment pipeline. When rag_type is set to remote::pgvector, the enrichment script generates a pgvector provider config with PostgreSQL connection fields instead of the FAISS-style kv_sqlite storage backend. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/byok_guide.md | 91 +++++------ docs/openapi.json | 78 +++++++++- docs/rag_guide.md | 41 ++--- src/llama_stack_configuration.py | 81 ++++++++-- src/models/config.py | 62 +++++++- tests/unit/models/config/test_byok_rag.py | 69 +++++++++ .../models/config/test_dump_configuration.py | 5 + tests/unit/test_llama_stack_configuration.py | 144 ++++++++++++++++++ 8 files changed, 481 insertions(+), 90 deletions(-) diff --git a/docs/byok_guide.md b/docs/byok_guide.md index aae702b29..f63c4f01f 100644 --- a/docs/byok_guide.md +++ b/docs/byok_guide.md @@ -189,18 +189,33 @@ byok_rag: score_multiplier: 1.0 # Weight for Inline RAG result ranking (default: 1.0) ``` -**`byok_rag` field reference:** +**Common fields (all providers):** | Field | Required | Default | Description | |-----------------------|----------|-------------------------------------------|-------------------------------------------------------------------------------------------| | `rag_id` | Yes | — | Unique identifier for the knowledge source | -| `rag_type` | No | `inline::faiss` | Vector store provider type | +| `rag_type` | No | `inline::faiss` | Vector store provider type (`inline::faiss` or `remote::pgvector`) | | `embedding_model` | No | `sentence-transformers/all-mpnet-base-v2` | Embedding model identifier or path | | `embedding_dimension` | No | `768` | Embedding vector dimensionality | | `vector_db_id` | Yes | — | Vector store ID generated by rag-content (e.g. `vs_8c94967b-81cc-4028-a294-9cfac6fd9ae2`) | -| `db_path` | Yes | — | Path to the vector database file | | `score_multiplier` | No | `1.0` | Weight for Inline RAG ranking (values > 1.0 boost; < 1.0 reduce) | +**FAISS fields** (`rag_type: inline::faiss`): + +| Field | Required | Default | Description | +|-----------|----------|---------|----------------------------------| +| `db_path` | Yes | — | Path to the vector database file | + +**pgvector fields** (`rag_type: remote::pgvector`): + +| Field | Required | Default | Description | +|------------|----------|--------------------------|---------------------| +| `host` | No | `${env.POSTGRES_HOST}` | PostgreSQL host | +| `port` | No | `${env.POSTGRES_PORT}` | PostgreSQL port | +| `db` | No | `${env.POSTGRES_DATABASE}` | PostgreSQL database | +| `user` | No | `${env.POSTGRES_USER}` | PostgreSQL user | +| `password` | No | `${env.POSTGRES_PASSWORD}` | PostgreSQL password | + **Multiple knowledge sources:** You can configure multiple BYOK sources. When using Inline RAG, `score_multiplier` adjusts the relative importance of each store's results: @@ -283,28 +298,27 @@ byok_rag: ### 2. pgvector (PostgreSQL) - **Type**: PostgreSQL with pgvector extension - **Best for**: Large-scale deployments, shared knowledge bases -- **Configuration**: `remote::pgvector` +- **Configuration**: `rag_type: remote::pgvector` - **Requirements**: PostgreSQL with pgvector extension -> [!NOTE] -> pgvector is not yet supported via `byok_rag` in `lightspeed-stack.yaml` (see [LCORE-2437](https://redhat.atlassian.net/browse/LCORE-2437)). -> It must be configured directly in the Llama Stack configuration file. - ```yaml -vector_io: -- provider_id: pgvector-knowledge - provider_type: remote::pgvector - config: - host: localhost - port: 5432 - db: knowledge_db - user: lightspeed_user +byok_rag: + - rag_id: pgvector-knowledge + rag_type: remote::pgvector + embedding_model: sentence-transformers/all-mpnet-base-v2 + embedding_dimension: 768 + vector_db_id: rhdocs + host: ${env.POSTGRES_HOST} + port: ${env.POSTGRES_PORT} + db: ${env.POSTGRES_DATABASE} + user: ${env.POSTGRES_USER} password: ${env.POSTGRES_PASSWORD} - kvstore: - type: sqlite - db_path: .llama/distributions/pgvector/registry.db ``` +> [!NOTE] +> Connection fields (`host`, `port`, `db`, `user`, `password`) default to +> `${env.POSTGRES_*}` environment variable references when omitted. + **pgvector Table Schema:** - `id` (text): UUID identifier of the chunk - `document` (jsonb): JSON containing content and metadata @@ -342,13 +356,7 @@ rag: ### Example 2: Multiple Knowledge Sources with pgvector -A configuration combining a local FAISS store (via `byok_rag`) with a remote pgvector store (configured directly in the Llama Stack configuration file): - -> [!NOTE] -> pgvector is not yet supported via `byok_rag` in `lightspeed-stack.yaml` (see [LCORE-2437](https://redhat.atlassian.net/browse/LCORE-2437)). -> The pgvector provider must be configured directly in the Llama Stack configuration file. - -**`lightspeed-stack.yaml`** — FAISS store and RAG strategy: +A configuration combining a local FAISS store with a remote pgvector store: ```yaml name: Lightspeed Core Service (LCS) @@ -365,34 +373,29 @@ byok_rag: vector_db_id: vs_e9d8c7b6-43af-4b2d-8e1f-0a9b8c7d6e5f db_path: /data/vector_dbs/local/faiss_store.db score_multiplier: 1.0 + - rag_id: enterprise-kb + rag_type: remote::pgvector + embedding_model: sentence-transformers/all-mpnet-base-v2 + embedding_dimension: 768 + vector_db_id: enterprise_docs + host: ${env.POSTGRES_HOST} + port: ${env.POSTGRES_PORT} + db: ${env.POSTGRES_DATABASE} + user: ${env.POSTGRES_USER} + password: ${env.POSTGRES_PASSWORD} rag: inline: - local-docs + - enterprise-kb tool: - local-docs -``` - -**Llama Stack configuration file** — pgvector provider: - -```yaml -vector_io: -- provider_id: enterprise-kb - provider_type: remote::pgvector - config: - host: localhost - port: 5432 - db: knowledge_db - user: lightspeed_user - password: ${env.POSTGRES_PASSWORD} - kvstore: - type: sqlite - db_path: .llama/distributions/pgvector/registry.db + - enterprise-kb ``` > [!NOTE] > For pgvector, ensure your PostgreSQL credentials are available via environment variables -> (e.g., `POSTGRES_PASSWORD`). +> (e.g., `POSTGRES_HOST`, `POSTGRES_PASSWORD`). > [!TIP] > A complete working example combining BYOK and OKP is available at diff --git a/docs/openapi.json b/docs/openapi.json index 0150e2055..b0c88ad2f 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -11818,7 +11818,7 @@ "type": "string", "minLength": 1, "title": "RAG type", - "description": "Type of RAG database.", + "description": "Type of RAG database (e.g. 'inline::faiss', 'remote::pgvector').", "default": "inline::faiss" }, "embedding_model": { @@ -11842,9 +11842,16 @@ "description": "Vector database identification." }, "db_path": { - "type": "string", + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], "title": "DB path", - "description": "Path to RAG database." + "description": "Path to RAG database. Required for inline::faiss." }, "score_multiplier": { "type": "number", @@ -11852,14 +11859,75 @@ "title": "Score multiplier", "description": "Multiplier applied to relevance scores from this vector store. Used to weight results when querying multiple knowledge sources. Values > 1 boost this store's results; values < 1 reduce them.", "default": 1.0 + }, + "host": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "PostgreSQL host", + "description": "PostgreSQL host for remote::pgvector. Defaults to ${env.POSTGRES_HOST} when rag_type is remote::pgvector." + }, + "port": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "PostgreSQL port", + "description": "PostgreSQL port for remote::pgvector. Defaults to ${env.POSTGRES_PORT} when rag_type is remote::pgvector." + }, + "db": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "PostgreSQL database", + "description": "PostgreSQL database name for remote::pgvector. Defaults to ${env.POSTGRES_DATABASE} when rag_type is remote::pgvector." + }, + "user": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "PostgreSQL user", + "description": "PostgreSQL user for remote::pgvector. Defaults to ${env.POSTGRES_USER} when rag_type is remote::pgvector." + }, + "password": { + "anyOf": [ + { + "type": "string", + "format": "password", + "writeOnly": true + }, + { + "type": "null" + } + ], + "title": "PostgreSQL password", + "description": "PostgreSQL password for remote::pgvector. Defaults to ${env.POSTGRES_PASSWORD} when rag_type is remote::pgvector." } }, "additionalProperties": false, "type": "object", "required": [ "rag_id", - "vector_db_id", - "db_path" + "vector_db_id" ], "title": "ByokRag", "description": "BYOK (Bring Your Own Knowledge) RAG configuration." diff --git a/docs/rag_guide.md b/docs/rag_guide.md index 490e413ef..07a3938f9 100644 --- a/docs/rag_guide.md +++ b/docs/rag_guide.md @@ -88,10 +88,6 @@ See the full working [config example](../examples/lightspeed-stack-byok-okp-rag. This example shows how to configure a remote PostgreSQL database with the [pgvector](https://github.com/pgvector/pgvector) extension for storing embeddings. -> [!NOTE] -> pgvector is not yet supported via `byok_rag` in `lightspeed-stack.yaml` (see [LCORE-2437](https://redhat.atlassian.net/browse/LCORE-2437)). -> It must be configured directly in the Llama Stack configuration file. - > You will need to install PostgreSQL with a matching version to pgvector, then log in with `psql` and enable the extension with: > ```sql > CREATE EXTENSION IF NOT EXISTS vector; @@ -107,33 +103,22 @@ Each pgvector-backed table follows this schema: > The `vector_store_id` (e.g. `rhdocs`) is used to point to the table named `vector_store_rhdocs` in the specified database, which stores the vector embeddings. ```yaml -providers: - [...] - vector_io: - - provider_id: pgvector-example - provider_type: remote::pgvector - config: - host: localhost - port: 5432 - db: pgvector_example # PostgreSQL database (psql -d pgvector_example) - user: lightspeed # PostgreSQL user - password: password123 - kvstore: - type: sqlite - db_path: .llama/distributions/pgvector/pgvector_registry.db -vector_stores: -- embedding_dimension: 768 - embedding_model: sentence-transformers/all-mpnet-base-v2 - provider_id: pgvector-example - # A unique ID that becomes the PostgreSQL table name, prefixed with 'vector_store_'. - # e.g., 'rhdocs' will create the table 'vector_store_rhdocs'. - # If the table was already created, this value must match the ID used at creation. - vector_store_id: rhdocs +byok_rag: + - rag_id: pgvector-example + rag_type: remote::pgvector + embedding_model: sentence-transformers/all-mpnet-base-v2 + embedding_dimension: 768 + vector_db_id: rhdocs # becomes PostgreSQL table 'vector_store_rhdocs' + host: ${env.POSTGRES_HOST} + port: ${env.POSTGRES_PORT} + db: ${env.POSTGRES_DATABASE} + user: ${env.POSTGRES_USER} + password: ${env.POSTGRES_PASSWORD} ``` > [!NOTE] -> For pgvector, the PostgreSQL connection details (host, port, database, user, password) are configured -> in the provider configuration. Use environment variables for credentials. +> Connection fields (`host`, `port`, `db`, `user`, `password`) default to `${env.POSTGRES_*}` +> environment variable references when omitted. Use environment variables for credentials. --- diff --git a/src/llama_stack_configuration.py b/src/llama_stack_configuration.py index ca0775bcf..0c6a2ff00 100644 --- a/src/llama_stack_configuration.py +++ b/src/llama_stack_configuration.py @@ -11,12 +11,34 @@ import yaml from llama_stack.core.stack import replace_env_vars +from pydantic import SecretStr import constants from log import get_logger logger = get_logger(__name__) +VECTOR_IO_TEMPLATES: dict[str, dict[str, Any]] = { + "inline::faiss": { + "persistence_backend": "{backend_name}", + "persistence_namespace": "vector_io::faiss", + "needs_storage_backend": True, + "extra_fields": {}, + }, + "remote::pgvector": { + "persistence_backend": "kv_default", + "persistence_namespace": "vector_io::pgvector", + "needs_storage_backend": False, + "extra_fields": { + "host": "${env.POSTGRES_HOST}", + "port": "${env.POSTGRES_PORT}", + "db": "${env.POSTGRES_DATABASE}", + "user": "${env.POSTGRES_USER}", + "password": "${env.POSTGRES_PASSWORD}", + }, + }, +} + class YamlDumper(yaml.Dumper): # pylint: disable=too-many-ancestors """Custom YAML dumper with proper indentation levels.""" @@ -137,19 +159,25 @@ def construct_storage_backends_section( if "storage" in ls_config and "backends" in ls_config["storage"]: output = ls_config["storage"]["backends"].copy() - # add new backends for each BYOK RAG + # add new backends for each BYOK RAG (skip types that don't need one) + added = 0 for brag in byok_rag: if not brag.get("rag_id"): raise ValueError(f"BYOK RAG entry is missing required 'rag_id': {brag}") + rag_type = brag.get("rag_type", constants.DEFAULT_RAG_TYPE) + template = VECTOR_IO_TEMPLATES.get(rag_type, {}) + if not template.get("needs_storage_backend", True): + continue rag_id = brag["rag_id"] backend_name = f"byok_{rag_id}_storage" output[backend_name] = { "type": "kv_sqlite", "db_path": brag.get("db_path", f".llama/{rag_id}.db"), } + added += 1 logger.info( "Added %s backends into storage.backends section, total backends %s", - len(byok_rag), + added, len(output), ) return output @@ -284,6 +312,44 @@ def construct_models_section( return output +def _build_vector_io_config( + rag_type: str, backend_name: str, brag: dict[str, Any] +) -> dict[str, Any]: + """Build the provider config dict from VECTOR_IO_TEMPLATES. + + Parameters: + rag_type: Llama Stack provider type (e.g. 'inline::faiss', 'remote::pgvector'). + backend_name: Storage backend name (used when template has '{backend_name}'). + brag: BYOK RAG entry dict — extra_fields are read from here. + + Returns: + dict[str, Any]: Provider config mapping. + """ + template = VECTOR_IO_TEMPLATES.get(rag_type) + if template is None: + raise ValueError( + f"Unsupported rag_type '{rag_type}'. " + f"Supported types: {list(VECTOR_IO_TEMPLATES.keys())}" + ) + persistence_backend = template["persistence_backend"].format( + backend_name=backend_name + ) + config: dict[str, Any] = { + "persistence": { + "namespace": template["persistence_namespace"], + "backend": persistence_backend, + } + } + for field, default in template.get("extra_fields", {}).items(): + value = brag.get(field) + if isinstance(value, SecretStr): + value = value.get_secret_value() + if value is None or (isinstance(value, str) and not value.strip()): + value = default + config[field] = value + return config + + def construct_vector_io_providers_section( ls_config: dict[str, Any], byok_rag: list[dict[str, Any]] ) -> list[dict[str, Any]]: @@ -335,16 +401,13 @@ def construct_vector_io_providers_section( continue existing_ids.add(provider_id) added += 1 + rag_type = brag.get("rag_type", constants.DEFAULT_RAG_TYPE) + config = _build_vector_io_config(rag_type, backend_name, brag) output.append( { "provider_id": provider_id, - "provider_type": brag.get("rag_type", "inline::faiss"), - "config": { - "persistence": { - "namespace": "vector_io::faiss", - "backend": backend_name, - } - }, + "provider_type": rag_type, + "config": config, } ) logger.info( diff --git a/src/models/config.py b/src/models/config.py index 923d720f0..167f9ff10 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -1776,7 +1776,7 @@ class ByokRag(ConfigurationBase): constants.DEFAULT_RAG_TYPE, min_length=1, title="RAG type", - description="Type of RAG database.", + description="Type of RAG database (e.g. 'inline::faiss', 'remote::pgvector').", ) embedding_model: str = Field( @@ -1799,10 +1799,10 @@ class ByokRag(ConfigurationBase): description="Vector database identification.", ) - db_path: str = Field( - ..., + db_path: Optional[str] = Field( + default=None, title="DB path", - description="Path to RAG database.", + description="Path to RAG database. Required for inline::faiss.", ) score_multiplier: float = Field( @@ -1814,6 +1814,60 @@ class ByokRag(ConfigurationBase): "Values > 1 boost this store's results; values < 1 reduce them.", ) + host: Optional[str] = Field( + default=None, + title="PostgreSQL host", + description="PostgreSQL host for remote::pgvector. " + "Defaults to ${env.POSTGRES_HOST} when rag_type is remote::pgvector.", + ) + + port: Optional[str] = Field( + default=None, + title="PostgreSQL port", + description="PostgreSQL port for remote::pgvector. " + "Defaults to ${env.POSTGRES_PORT} when rag_type is remote::pgvector.", + ) + + db: Optional[str] = Field( + default=None, + title="PostgreSQL database", + description="PostgreSQL database name for remote::pgvector. " + "Defaults to ${env.POSTGRES_DATABASE} when rag_type is remote::pgvector.", + ) + + user: Optional[str] = Field( + default=None, + title="PostgreSQL user", + description="PostgreSQL user for remote::pgvector. " + "Defaults to ${env.POSTGRES_USER} when rag_type is remote::pgvector.", + ) + + password: Optional[SecretStr] = Field( + default=None, + title="PostgreSQL password", + description="PostgreSQL password for remote::pgvector. " + "Defaults to ${env.POSTGRES_PASSWORD} when rag_type is remote::pgvector.", + ) + + @model_validator(mode="after") + def validate_rag_type_fields(self) -> Self: + """Validate and populate fields based on rag_type.""" + if self.rag_type == "inline::faiss": + if not self.db_path: + raise ValueError("db_path is required when rag_type is 'inline::faiss'") + elif self.rag_type == "remote::pgvector": + pgvector_defaults: dict[str, str | SecretStr] = { + "host": "${env.POSTGRES_HOST}", + "port": "${env.POSTGRES_PORT}", + "db": "${env.POSTGRES_DATABASE}", + "user": "${env.POSTGRES_USER}", + "password": SecretStr("${env.POSTGRES_PASSWORD}"), + } + for field_name, default_value in pgvector_defaults.items(): + if getattr(self, field_name) is None: + object.__setattr__(self, field_name, default_value) + return self + class QuotaLimiterConfiguration(ConfigurationBase): """Configuration for one quota limiter. diff --git a/tests/unit/models/config/test_byok_rag.py b/tests/unit/models/config/test_byok_rag.py index e80e749c3..b90c72cdd 100644 --- a/tests/unit/models/config/test_byok_rag.py +++ b/tests/unit/models/config/test_byok_rag.py @@ -1,5 +1,7 @@ """Unit tests for ByokRag model.""" +# pylint: disable=no-member + import pytest from pydantic import ValidationError @@ -199,3 +201,70 @@ def test_byok_rag_configuration_score_multiplier_must_be_positive() -> None: db_path="tests/configuration/rag.txt", score_multiplier=0.0, ) + + +def test_byok_rag_faiss_requires_db_path() -> None: + """Test that inline::faiss requires db_path.""" + with pytest.raises(ValidationError, match="db_path is required"): + _ = ByokRag( + rag_id="rag_id", + rag_type="inline::faiss", + vector_db_id="vector_db_id", + ) + + +def test_byok_rag_pgvector_defaults() -> None: + """Test pgvector auto-populates connection fields with env var defaults.""" + store = ByokRag( + rag_id="pg_store", + rag_type="remote::pgvector", + vector_db_id="vs_pg", + ) + assert store.rag_type == "remote::pgvector" + assert store.host == "${env.POSTGRES_HOST}" + assert store.port == "${env.POSTGRES_PORT}" + assert store.db == "${env.POSTGRES_DATABASE}" + assert store.user == "${env.POSTGRES_USER}" + assert store.password.get_secret_value() == "${env.POSTGRES_PASSWORD}" + assert store.db_path is None + + +def test_byok_rag_pgvector_custom_connection_fields() -> None: + """Test pgvector accepts custom connection field values.""" + store = ByokRag( + rag_id="pg_store", + rag_type="remote::pgvector", + vector_db_id="vs_pg", + host="db.example.com", + port="5433", + db="my_knowledge", + user="admin", + password="secret", + ) + assert store.host == "db.example.com" + assert store.port == "5433" + assert store.db == "my_knowledge" + assert store.user == "admin" + assert store.password.get_secret_value() == "secret" + + +def test_byok_rag_pgvector_partial_overrides() -> None: + """Test pgvector fills only missing connection fields with defaults.""" + store = ByokRag( + rag_id="pg_store", + rag_type="remote::pgvector", + vector_db_id="vs_pg", + host="custom-host", + ) + assert store.host == "custom-host" + assert store.port == "${env.POSTGRES_PORT}" + + +def test_byok_rag_pgvector_does_not_require_db_path() -> None: + """Test pgvector does not require db_path.""" + store = ByokRag( + rag_id="pg_store", + rag_type="remote::pgvector", + vector_db_id="vs_pg", + ) + assert store.db_path is None diff --git a/tests/unit/models/config/test_dump_configuration.py b/tests/unit/models/config/test_dump_configuration.py index 99ec6587b..e6c57acd1 100644 --- a/tests/unit/models/config/test_dump_configuration.py +++ b/tests/unit/models/config/test_dump_configuration.py @@ -1244,6 +1244,11 @@ def test_dump_configuration_byok(tmp_path: Path) -> None: "rag_type": "inline::faiss", "vector_db_id": "vector_db_id", "score_multiplier": 1.0, + "host": None, + "port": None, + "db": None, + "user": None, + "password": None, }, ], "quota_handlers": { diff --git a/tests/unit/test_llama_stack_configuration.py b/tests/unit/test_llama_stack_configuration.py index aaa3bf53e..4b46cfa8b 100644 --- a/tests/unit/test_llama_stack_configuration.py +++ b/tests/unit/test_llama_stack_configuration.py @@ -317,6 +317,112 @@ def test_construct_vector_io_providers_section_collapses_existing_duplicates() - assert output[0]["provider_id"] == "byok_rag1" +def test_construct_vector_io_providers_section_pgvector() -> None: + """Test generates correct pgvector provider config.""" + ls_config: dict[str, Any] = {"providers": {}} + byok_rag = [ + { + "rag_id": "pg1", + "vector_db_id": "vs_pg", + "rag_type": "remote::pgvector", + "host": "${env.POSTGRES_HOST}", + "port": "${env.POSTGRES_PORT}", + "db": "${env.POSTGRES_DATABASE}", + "user": "${env.POSTGRES_USER}", + "password": "${env.POSTGRES_PASSWORD}", + }, + ] + output = construct_vector_io_providers_section(ls_config, byok_rag) + assert len(output) == 1 + provider = output[0] + assert provider["provider_id"] == "byok_pg1" + assert provider["provider_type"] == "remote::pgvector" + assert provider["config"]["persistence"]["namespace"] == "vector_io::pgvector" + assert provider["config"]["persistence"]["backend"] == "kv_default" + assert provider["config"]["host"] == "${env.POSTGRES_HOST}" + assert provider["config"]["db"] == "${env.POSTGRES_DATABASE}" + + +def test_construct_vector_io_providers_section_mixed() -> None: + """Test mixed faiss and pgvector entries generate correct configs.""" + ls_config: dict[str, Any] = {"providers": {}} + byok_rag = [ + {"rag_id": "f1", "vector_db_id": "vs_f", "rag_type": "inline::faiss"}, + { + "rag_id": "pg1", + "vector_db_id": "vs_pg", + "rag_type": "remote::pgvector", + "host": "localhost", + "port": "5432", + "db": "mydb", + "user": "user", + "password": "pass", + }, + ] + output = construct_vector_io_providers_section(ls_config, byok_rag) + assert len(output) == 2 + faiss_p = next(p for p in output if p["provider_id"] == "byok_f1") + assert faiss_p["provider_type"] == "inline::faiss" + assert faiss_p["config"]["persistence"]["backend"] == "byok_f1_storage" + pg_p = next(p for p in output if p["provider_id"] == "byok_pg1") + assert pg_p["provider_type"] == "remote::pgvector" + assert pg_p["config"]["persistence"]["backend"] == "kv_default" + assert pg_p["config"]["host"] == "localhost" + + +def test_construct_storage_backends_section_skips_pgvector() -> None: + """Test pgvector entries are skipped (they use kv_default).""" + ls_config: dict[str, Any] = {} + byok_rag = [ + {"rag_id": "pg1", "vector_db_id": "vs_pg", "rag_type": "remote::pgvector"} + ] + output = construct_storage_backends_section(ls_config, byok_rag) + assert len(output) == 0 + + +def test_construct_storage_backends_section_mixed_faiss_pgvector() -> None: + """Test only faiss entries get storage backends, pgvector is skipped.""" + ls_config: dict[str, Any] = {} + byok_rag = [ + {"rag_id": "f1", "vector_db_id": "vs_f", "db_path": "/tmp/f.db"}, + {"rag_id": "pg1", "vector_db_id": "vs_pg", "rag_type": "remote::pgvector"}, + ] + output = construct_storage_backends_section(ls_config, byok_rag) + assert len(output) == 1 + assert "byok_f1_storage" in output + assert "byok_pg1_storage" not in output + + +def test_enrich_byok_rag_pgvector_end_to_end() -> None: + """Test enrich_byok_rag with a pgvector store entry.""" + ls_config: dict[str, Any] = {} + byok_rag = [ + { + "rag_id": "pg1", + "vector_db_id": "vs_pg", + "rag_type": "remote::pgvector", + "embedding_model": "sentence-transformers/all-mpnet-base-v2", + "embedding_dimension": 768, + "host": "${env.POSTGRES_HOST}", + "port": "${env.POSTGRES_PORT}", + "db": "${env.POSTGRES_DATABASE}", + "user": "${env.POSTGRES_USER}", + "password": "${env.POSTGRES_PASSWORD}", + }, + ] + enrich_byok_rag(ls_config, byok_rag) + assert "byok_pg1_storage" not in ls_config.get("storage", {}).get("backends", {}) + providers = ls_config["providers"]["vector_io"] + pg_p = next(p for p in providers if p["provider_id"] == "byok_pg1") + assert pg_p["provider_type"] == "remote::pgvector" + assert pg_p["config"]["persistence"]["backend"] == "kv_default" + assert pg_p["config"]["host"] == "${env.POSTGRES_HOST}" + store_ids = [ + s["vector_store_id"] for s in ls_config["registered_resources"]["vector_stores"] + ] + assert "vs_pg" in store_ids + + def test_enrich_byok_rag_skipped_still_dedupes_vector_io() -> None: """When BYOK is off, existing duplicate vector_io entries are still collapsed.""" dup = { @@ -609,6 +715,44 @@ def test_generate_configuration_with_byok(tmp_path: Path) -> None: assert "byok_rag1_embedding" in model_ids +def test_generate_configuration_with_pgvector(tmp_path: Path) -> None: + """Test generate_configuration adds pgvector BYOK entries.""" + config = { + "byok_rag": [ + { + "rag_id": "pg1", + "vector_db_id": "vs_pg", + "embedding_model": "sentence-transformers/all-mpnet-base-v2", + "embedding_dimension": 768, + "rag_type": "remote::pgvector", + "host": "localhost", + "port": "5432", + "db": "knowledge_db", + "user": "admin", + "password": "secret", + }, + ], + } + outfile = tmp_path / "output.yaml" + generate_configuration("tests/configuration/run.yaml", str(outfile), config) + with open(outfile, encoding="utf-8") as f: + result = yaml.safe_load(f) + store_ids = [ + s["vector_store_id"] for s in result["registered_resources"]["vector_stores"] + ] + assert "vs_pg" in store_ids + backends = result.get("storage", {}).get("backends", {}) + assert "byok_pg1_storage" not in backends + pg_provider = next( + p + for p in result["providers"]["vector_io"] + if p.get("provider_id") == "byok_pg1" + ) + assert pg_provider["provider_type"] == "remote::pgvector" + assert pg_provider["config"]["host"] == "localhost" + assert pg_provider["config"]["persistence"]["backend"] == "kv_default" + + # ============================================================================= # Test enrich_solr # =============================================================================