Sfoglia il codice sorgente

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

tags/0.9.0
zhuhao 1 anno fa
parent
commit
55e6123db9
Nessun account collegato all'indirizzo email del committer

+ 2
- 0
api/.env.example Vedi File

PGVECTOR_USER=postgres PGVECTOR_USER=postgres
PGVECTOR_PASSWORD=postgres PGVECTOR_PASSWORD=postgres
PGVECTOR_DATABASE=postgres PGVECTOR_DATABASE=postgres
PGVECTOR_MIN_CONNECTION=1
PGVECTOR_MAX_CONNECTION=5


# Tidb Vector configuration # Tidb Vector configuration
TIDB_VECTOR_HOST=xxx.eu-central-1.xxx.aws.tidbcloud.com TIDB_VECTOR_HOST=xxx.eu-central-1.xxx.aws.tidbcloud.com

+ 10
- 0
api/configs/middleware/vdb/pgvector_config.py Vedi File

description="Name of the PostgreSQL database to connect to", description="Name of the PostgreSQL database to connect to",
default=None, 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,
)

+ 12
- 2
api/core/rag/datasource/vdb/pgvector/pgvector.py Vedi File

user: str user: str
password: str password: str
database: str database: str
min_connection: int
max_connection: int


@model_validator(mode="before") @model_validator(mode="before")
@classmethod @classmethod
raise ValueError("config PGVECTOR_PASSWORD is required") raise ValueError("config PGVECTOR_PASSWORD is required")
if not values["database"]: if not values["database"]:
raise ValueError("config PGVECTOR_DATABASE is required") 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 return values






def _create_connection_pool(self, config: PGVectorConfig): def _create_connection_pool(self, config: PGVectorConfig):
return psycopg2.pool.SimpleConnectionPool( return psycopg2.pool.SimpleConnectionPool(
1,
5,
config.min_connection,
config.max_connection,
host=config.host, host=config.host,
port=config.port, port=config.port,
user=config.user, user=config.user,
user=dify_config.PGVECTOR_USER, user=dify_config.PGVECTOR_USER,
password=dify_config.PGVECTOR_PASSWORD, password=dify_config.PGVECTOR_PASSWORD,
database=dify_config.PGVECTOR_DATABASE, database=dify_config.PGVECTOR_DATABASE,
min_connection=dify_config.PGVECTOR_MIN_CONNECTION,
max_connection=dify_config.PGVECTOR_MAX_CONNECTION,
), ),
) )

+ 2
- 0
api/tests/integration_tests/vdb/pgvector/test_pgvector.py Vedi File

user="postgres", user="postgres",
password="difyai123456", password="difyai123456",
database="dify", database="dify",
min_connection=1,
max_connection=5,
), ),
) )



+ 2
- 0
docker/.env.example Vedi File

PGVECTOR_USER=postgres PGVECTOR_USER=postgres
PGVECTOR_PASSWORD=difyai123456 PGVECTOR_PASSWORD=difyai123456
PGVECTOR_DATABASE=dify PGVECTOR_DATABASE=dify
PGVECTOR_MIN_CONNECTION=1
PGVECTOR_MAX_CONNECTION=5


# pgvecto-rs configurations, only available when VECTOR_STORE is `pgvecto-rs` # pgvecto-rs configurations, only available when VECTOR_STORE is `pgvecto-rs`
PGVECTO_RS_HOST=pgvecto-rs PGVECTO_RS_HOST=pgvecto-rs

Loading…
Annulla
Salva