diff --git a/README.md b/README.md index f985b6fc..ee2a610c 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,7 @@ The following environment variables are required to run the application: - `POSTGRES_PASSWORD`: (Optional) The password for connecting to the PostgreSQL database. - `DB_HOST`: (Optional) The hostname or IP address of the PostgreSQL database server. - `DB_PORT`: (Optional) The port number of the PostgreSQL database server. +- `PGVECTOR_CREATE_EXTENSION`: (Optional) Set to "False" to skip the `CREATE EXTENSION IF NOT EXISTS vector` call on startup. Default is "True". Use this when the `vector` extension is already installed on a managed Postgres (e.g. RDS, Azure Database for PostgreSQL) and the application user is not a superuser. - `RAG_HOST`: (Optional) The hostname or IP address where the API server will run. Defaults to "0.0.0.0" - `RAG_PORT`: (Optional) The port number where the API server will run. Defaults to port 8000. - `JWT_SECRET`: (Optional) The secret key used for verifying JWT tokens for requests. diff --git a/app/config.py b/app/config.py index 8031763a..9d92be2f 100644 --- a/app/config.py +++ b/app/config.py @@ -59,6 +59,9 @@ def get_env_variable( POSTGRES_PASSWORD = get_env_variable("POSTGRES_PASSWORD", "mypassword") DB_HOST = get_env_variable("DB_HOST", "db") DB_PORT = get_env_variable("DB_PORT", "5432") +PGVECTOR_CREATE_EXTENSION = get_env_variable( + "PGVECTOR_CREATE_EXTENSION", "True" +).lower() in ("true", "1", "yes", "on") COLLECTION_NAME = get_env_variable("COLLECTION_NAME", "testcollection") ATLAS_MONGO_DB_URI = get_env_variable( "ATLAS_MONGO_DB_URI", "mongodb://127.0.0.1:27018/LibreChat" @@ -336,6 +339,7 @@ def init_embeddings(provider, model): embeddings=embeddings, collection_name=COLLECTION_NAME, mode="async", + create_extension=PGVECTOR_CREATE_EXTENSION, ) elif VECTOR_DB_TYPE == VectorDBType.ATLAS_MONGO: # Backward compatability check diff --git a/app/services/vector_store/factory.py b/app/services/vector_store/factory.py index 854cc78a..6418ed38 100644 --- a/app/services/vector_store/factory.py +++ b/app/services/vector_store/factory.py @@ -20,11 +20,16 @@ def get_vector_store( collection_name: str, mode: str = "sync", search_index: Optional[str] = None, + create_extension: bool = True, ): """Create a vector store instance for the given mode. Note: For 'atlas-mongo' mode, the MongoClient is stored at module level so it can be closed on shutdown via close_vector_store_connections(). + + Set create_extension=False when the Postgres user lacks superuser + privileges and the `vector` extension is already installed out-of-band + (e.g. managed Postgres services like RDS, Azure Database for PostgreSQL). """ global _mongo_client @@ -34,6 +39,7 @@ def get_vector_store( embedding_function=embeddings, collection_name=collection_name, use_jsonb=True, + create_extension=create_extension, ) elif mode == "async": return AsyncPgVector( @@ -41,6 +47,7 @@ def get_vector_store( embedding_function=embeddings, collection_name=collection_name, use_jsonb=True, + create_extension=create_extension, ) elif mode == "atlas-mongo": if _mongo_client is not None: diff --git a/tests/services/test_vector_store_factory.py b/tests/services/test_vector_store_factory.py index d1139318..7f963f94 100644 --- a/tests/services/test_vector_store_factory.py +++ b/tests/services/test_vector_store_factory.py @@ -113,6 +113,35 @@ def test_get_vector_store_async_passes_use_jsonb(): assert kwargs.get("use_jsonb") is True +def test_get_vector_store_defaults_create_extension_true(): + """PgVector must default to create_extension=True for back-compat.""" + with patch("app.services.vector_store.factory.AsyncPgVector") as MockPG: + mock_embeddings = MagicMock() + factory.get_vector_store("conn", mock_embeddings, "coll", mode="async") + _, kwargs = MockPG.call_args + assert kwargs.get("create_extension") is True + + +def test_get_vector_store_propagates_create_extension_false(): + """create_extension=False must reach the underlying PGVector — the escape + hatch for managed Postgres where the app user can't CREATE EXTENSION.""" + with patch("app.services.vector_store.factory.AsyncPgVector") as MockPG: + mock_embeddings = MagicMock() + factory.get_vector_store( + "conn", mock_embeddings, "coll", mode="async", create_extension=False + ) + _, kwargs = MockPG.call_args + assert kwargs.get("create_extension") is False + + with patch("app.services.vector_store.factory.ExtendedPgVector") as MockPG: + mock_embeddings = MagicMock() + factory.get_vector_store( + "conn", mock_embeddings, "coll", mode="sync", create_extension=False + ) + _, kwargs = MockPG.call_args + assert kwargs.get("create_extension") is False + + def test_load_file_content_cleans_up_on_lazy_load_failure(): """cleanup_temp_encoding_file is called even when lazy_load() raises.""" from app.routes.document_routes import load_file_content