mirror of https://github.com/hwchase17/langchain
refactor: enable connection pool usage in PGVector (#11514)
- **Description:** `PGVector` refactored to use connection pool. - **Issue:** #11433, - **Tag maintainer:** @hwchase17 @eyurtsev, --------- Co-authored-by: Diego Rani Mazine <diego.mazine@mercadolivre.com> Co-authored-by: Nuno Campos <nuno@langchain.dev>pull/15293/head
parent
507c195a4b
commit
ec72225265
@ -0,0 +1,17 @@
|
||||
version: "3.8"
|
||||
|
||||
services:
|
||||
pgvector:
|
||||
image: ankane/pgvector:latest
|
||||
environment:
|
||||
POSTGRES_DB: ${PGVECTOR_DB:-postgres}
|
||||
POSTGRES_USER: ${PGVECTOR_USER:-postgres}
|
||||
POSTGRES_PASSWORD: ${PGVECTOR_PASSWORD:-postgres}
|
||||
ports:
|
||||
- ${PGVECTOR_PORT:-5432}:5432
|
||||
restart: unless-stopped
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-f", "http://localhost:5432"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 5
|
@ -0,0 +1,73 @@
|
||||
"""Test PGVector functionality."""
|
||||
from unittest import mock
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from langchain_community.embeddings import FakeEmbeddings
|
||||
from langchain_community.vectorstores import pgvector
|
||||
|
||||
_CONNECTION_STRING = pgvector.PGVector.connection_string_from_db_params(
|
||||
driver="psycopg2",
|
||||
host="localhost",
|
||||
port=5432,
|
||||
database="postgres",
|
||||
user="postgres",
|
||||
password="postgres",
|
||||
)
|
||||
|
||||
_EMBEDDING_FUNCTION = FakeEmbeddings(size=1536)
|
||||
|
||||
|
||||
@pytest.mark.requires("pgvector")
|
||||
@mock.patch("sqlalchemy.create_engine")
|
||||
def test_given_a_connection_is_provided_then_no_engine_should_be_created(
|
||||
create_engine: Mock,
|
||||
) -> None:
|
||||
"""When a connection is provided then no engine should be created."""
|
||||
pgvector.PGVector(
|
||||
connection_string=_CONNECTION_STRING,
|
||||
embedding_function=_EMBEDDING_FUNCTION,
|
||||
connection=mock.MagicMock(),
|
||||
)
|
||||
create_engine.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.requires("pgvector")
|
||||
@mock.patch("sqlalchemy.create_engine")
|
||||
def test_given_no_connection_or_engine_args_provided_default_engine_should_be_used(
|
||||
create_engine: Mock,
|
||||
) -> None:
|
||||
"""When no connection or engine arguments are provided then the default configuration must be used.""" # noqa: E501
|
||||
pgvector.PGVector(
|
||||
connection_string=_CONNECTION_STRING,
|
||||
embedding_function=_EMBEDDING_FUNCTION,
|
||||
)
|
||||
create_engine.assert_called_with(
|
||||
url=_CONNECTION_STRING,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.requires("pgvector")
|
||||
@mock.patch("sqlalchemy.create_engine")
|
||||
def test_given_engine_args_are_provided_then_they_should_be_used(
|
||||
create_engine: Mock,
|
||||
) -> None:
|
||||
"""When engine arguments are provided then they must be used to create the underlying engine.""" # noqa: E501
|
||||
engine_args = {
|
||||
"pool_size": 5,
|
||||
"max_overflow": 10,
|
||||
"pool_recycle": -1,
|
||||
"pool_use_lifo": False,
|
||||
"pool_pre_ping": False,
|
||||
"pool_timeout": 30,
|
||||
}
|
||||
pgvector.PGVector(
|
||||
connection_string=_CONNECTION_STRING,
|
||||
embedding_function=_EMBEDDING_FUNCTION,
|
||||
engine_args=engine_args,
|
||||
)
|
||||
create_engine.assert_called_with(
|
||||
url=_CONNECTION_STRING,
|
||||
**engine_args,
|
||||
)
|
Loading…
Reference in New Issue