eliott/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-04-30 11:27:08 +02:00
Code Issues Packages Projects Releases Wiki Activity

enh/refac: pgvector pool support

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2025-07-14 12:18:44 +04:00
parent 9c49d9e641
commit b4f04ff3a7
2 changed files with 62 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
backend/open_webui/retrieval/vector/dbs/pgvector.py
View File

@@ -18,7 +18,7 @@ from sqlalchemy import (
values,
)
from sqlalchemy.sql import true
from sqlalchemy.pool import NullPool
from sqlalchemy.pool import NullPool, QueuePool
from sqlalchemy.orm import declarative_base, scoped_session, sessionmaker
from sqlalchemy.dialects.postgresql import JSONB, array
@@ -37,6 +37,10 @@ from open_webui.config import (
PGVECTOR_INITIALIZE_MAX_VECTOR_LENGTH,
PGVECTOR_PGCRYPTO,
PGVECTOR_PGCRYPTO_KEY,
PGVECTOR_POOL_SIZE,
PGVECTOR_POOL_MAX_OVERFLOW,
PGVECTOR_POOL_TIMEOUT,
PGVECTOR_POOL_RECYCLE,
)
from open_webui.env import SRC_LOG_LEVELS
@@ -80,9 +84,24 @@ class PgvectorClient(VectorDBBase):
self.session = Session
else:
engine = create_engine(
PGVECTOR_DB_URL, pool_pre_ping=True, poolclass=NullPool
)
if isinstance(PGVECTOR_POOL_SIZE, int):
if PGVECTOR_POOL_SIZE > 0:
engine = create_engine(
PGVECTOR_DB_URL,
pool_size=PGVECTOR_POOL_SIZE,
max_overflow=PGVECTOR_POOL_MAX_OVERFLOW,
pool_timeout=PGVECTOR_POOL_TIMEOUT,
pool_recycle=PGVECTOR_POOL_RECYCLE,
pool_pre_ping=True,
poolclass=QueuePool,
)
else:
engine = create_engine(
PGVECTOR_DB_URL, pool_pre_ping=True, poolclass=NullPool
)
else:
engine = create_engine(PGVECTOR_DB_URL, pool_pre_ping=True)
SessionLocal = sessionmaker(
autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
)