From 55e6123db93ae7fbda06596d6cd0a163747e6559 Mon Sep 17 00:00:00 2001
From: zhuhao <37029601+hwzhuhao@users.noreply.github.com>
Date: Fri, 27 Sep 2024 18:16:20 +0800
Subject: [PATCH] feat: add min-connection and max-connection for pgvector
 (#8841)

---
 api/.env.example                                   |  2 ++
 api/configs/middleware/vdb/pgvector_config.py      | 10 ++++++++++
 api/core/rag/datasource/vdb/pgvector/pgvector.py   | 14 ++++++++++++--
 .../vdb/pgvector/test_pgvector.py                  |  2 ++
 docker/.env.example                                |  2 ++
 5 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/api/.env.example b/api/.env.example
index a3db406aea..7e5c636ff8 100644
--- a/api/.env.example
+++ b/api/.env.example
@@ -162,6 +162,8 @@ PGVECTOR_PORT=5433
 PGVECTOR_USER=postgres
 PGVECTOR_PASSWORD=postgres
 PGVECTOR_DATABASE=postgres
+PGVECTOR_MIN_CONNECTION=1
+PGVECTOR_MAX_CONNECTION=5
 
 # Tidb Vector configuration
 TIDB_VECTOR_HOST=xxx.eu-central-1.xxx.aws.tidbcloud.com
diff --git a/api/configs/middleware/vdb/pgvector_config.py b/api/configs/middleware/vdb/pgvector_config.py
index 395dcaa420..85f5dca7e2 100644
--- a/api/configs/middleware/vdb/pgvector_config.py
+++ b/api/configs/middleware/vdb/pgvector_config.py
@@ -33,3 +33,13 @@ class PGVectorConfig(BaseSettings):
         description="Name of the PostgreSQL database to connect to",
         default=None,
     )
+
+    PGVECTOR_MIN_CONNECTION: PositiveInt = Field(
+        description="Min connection of the PostgreSQL database",
+        default=1,
+    )
+
+    PGVECTOR_MAX_CONNECTION: PositiveInt = Field(
+        description="Max connection of the PostgreSQL database",
+        default=5,
+    )
diff --git a/api/core/rag/datasource/vdb/pgvector/pgvector.py b/api/core/rag/datasource/vdb/pgvector/pgvector.py
index 79879d4f63..d90707ebcd 100644
--- a/api/core/rag/datasource/vdb/pgvector/pgvector.py
+++ b/api/core/rag/datasource/vdb/pgvector/pgvector.py
@@ -23,6 +23,8 @@ class PGVectorConfig(BaseModel):
     user: str
     password: str
     database: str
+    min_connection: int
+    max_connection: int
 
     @model_validator(mode="before")
     @classmethod
@@ -37,6 +39,12 @@ class PGVectorConfig(BaseModel):
             raise ValueError("config PGVECTOR_PASSWORD is required")
         if not values["database"]:
             raise ValueError("config PGVECTOR_DATABASE is required")
+        if not values["min_connection"]:
+            raise ValueError("config PGVECTOR_MIN_CONNECTION is required")
+        if not values["max_connection"]:
+            raise ValueError("config PGVECTOR_MAX_CONNECTION is required")
+        if values["min_connection"] > values["max_connection"]:
+            raise ValueError("config PGVECTOR_MIN_CONNECTION should less than PGVECTOR_MAX_CONNECTION")
         return values
 
 
@@ -61,8 +69,8 @@ class PGVector(BaseVector):
 
     def _create_connection_pool(self, config: PGVectorConfig):
         return psycopg2.pool.SimpleConnectionPool(
-            1,
-            5,
+            config.min_connection,
+            config.max_connection,
             host=config.host,
             port=config.port,
             user=config.user,
@@ -213,5 +221,7 @@ class PGVectorFactory(AbstractVectorFactory):
                 user=dify_config.PGVECTOR_USER,
                 password=dify_config.PGVECTOR_PASSWORD,
                 database=dify_config.PGVECTOR_DATABASE,
+                min_connection=dify_config.PGVECTOR_MIN_CONNECTION,
+                max_connection=dify_config.PGVECTOR_MAX_CONNECTION,
             ),
         )
diff --git a/api/tests/integration_tests/vdb/pgvector/test_pgvector.py b/api/tests/integration_tests/vdb/pgvector/test_pgvector.py
index c5a986b747..72efdc2780 100644
--- a/api/tests/integration_tests/vdb/pgvector/test_pgvector.py
+++ b/api/tests/integration_tests/vdb/pgvector/test_pgvector.py
@@ -18,6 +18,8 @@ class PGVectorTest(AbstractVectorTest):
                 user="postgres",
                 password="difyai123456",
                 database="dify",
+                min_connection=1,
+                max_connection=5,
             ),
         )
 
diff --git a/docker/.env.example b/docker/.env.example
index d43c3edc7e..eb05f7aa4f 100644
--- a/docker/.env.example
+++ b/docker/.env.example
@@ -391,6 +391,8 @@ PGVECTOR_PORT=5432
 PGVECTOR_USER=postgres
 PGVECTOR_PASSWORD=difyai123456
 PGVECTOR_DATABASE=dify
+PGVECTOR_MIN_CONNECTION=1
+PGVECTOR_MAX_CONNECTION=5
 
 # pgvecto-rs configurations, only available when VECTOR_STORE is `pgvecto-rs`
 PGVECTO_RS_HOST=pgvecto-rs