feat: add min-connection and max-connection for pgvector (#8841)

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

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,
+    )

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 @@ def validate_config(cls, values: dict) -> dict:
             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 @@ def get_type(self) -> str:
 
     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 @@ def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings
                 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,
             ),
         )

api/tests/integration_tests/vdb/pgvector/test_pgvector.py

@@ -18,6 +18,8 @@ def __init__(self):
                 user="postgres",
                 password="difyai123456",
                 database="dify",
+                min_connection=1,
+                max_connection=5,
             ),
         )
 

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