enh: PGVECTOR_CREATE_EXTENSION env var

This commit is contained in:
Timothy Jaeryang Baek 2025-08-31 23:58:18 +04:00
parent e0ab5adb97
commit c2b4976c82

View file

@ -37,6 +37,7 @@ from open_webui.retrieval.vector.main import (
from open_webui.config import ( from open_webui.config import (
PGVECTOR_DB_URL, PGVECTOR_DB_URL,
PGVECTOR_INITIALIZE_MAX_VECTOR_LENGTH, PGVECTOR_INITIALIZE_MAX_VECTOR_LENGTH,
PGVECTOR_CREATE_EXTENSION,
PGVECTOR_PGCRYPTO, PGVECTOR_PGCRYPTO,
PGVECTOR_PGCRYPTO_KEY, PGVECTOR_PGCRYPTO_KEY,
PGVECTOR_POOL_SIZE, PGVECTOR_POOL_SIZE,
@ -112,18 +113,19 @@ class PgvectorClient(VectorDBBase):
try: try:
# Ensure the pgvector extension is available # Ensure the pgvector extension is available
# Use a conditional check to avoid permission issues on Azure PostgreSQL # Use a conditional check to avoid permission issues on Azure PostgreSQL
self.session.execute( if PGVECTOR_CREATE_EXTENSION:
text( self.session.execute(
""" text(
DO $$ """
BEGIN DO $$
IF NOT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector') THEN BEGIN
CREATE EXTENSION IF NOT EXISTS vector; IF NOT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector') THEN
END IF; CREATE EXTENSION IF NOT EXISTS vector;
END $$; END IF;
""" END $$;
"""
)
) )
)
if PGVECTOR_PGCRYPTO: if PGVECTOR_PGCRYPTO:
# Ensure the pgcrypto extension is available for encryption # Ensure the pgcrypto extension is available for encryption