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