浏览代码

enh: PGVECTOR_CREATE_EXTENSION env var

Timothy Jaeryang Baek 1 月之前
父节点
当前提交
c2b4976c82
共有 1 个文件被更改,包括 13 次插入11 次删除
  1. 13 11
      backend/open_webui/retrieval/vector/dbs/pgvector.py

+ 13 - 11
backend/open_webui/retrieval/vector/dbs/pgvector.py

@@ -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