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