PGVector bound to a single database connection
Labels: name: area: vector store, name: auto:improvement

### System Info

python: 3.11
langchain: 0.0.306

### Who can help?

Although PGVector  implementation utilizes SQLAlchemy that supports connection pooling features, it gets one connection from the pool, saves it as a class attribute and reuses in all requests. 

```python
class PGVector(VectorStore):
    # ...

    def __init__(
        self,
        # ...
    ) -> None:
        # ...
        self.__post_init__()


    def __post_init__(
        self,
    ) -> None:
        self._conn = self.connect()
        # ...


    def create_vector_extension(self) -> None:
        try:
            with Session(self._conn) as session:
                # ...
        except Exception as e:
            self.logger.exception(e)


    def create_tables_if_not_exists(self) -> None:
        with self._conn.begin():
            Base.metadata.create_all(self._conn)


    @contextlib.contextmanager
    def _make_session(self) -> Generator[Session, None, None]:
        yield Session(self._conn)

    # ...
```

This means that:

- Since the same connection is user for handle all requests, they will not be executed in parallel but rather enqueued. In a scenario with a large number of requests, this may not scale well.

- Great pool features like connection recycle and test for liveness can not be used.

@hwchase17 
@eyurtsev


### Information

- [ ] The official example notebooks/scripts
- [X] My own modified scripts

### Related Components

- [ ] LLMs/Chat Models
- [ ] Embedding Models
- [ ] Prompts / Prompt Templates / Prompt Selectors
- [ ] Output Parsers
- [ ] Document Loaders
- [X] Vector Stores / Retrievers
- [ ] Memory
- [ ] Agents / Agent Executors
- [ ] Tools / Toolkits
- [ ] Chains
- [ ] Callbacks/Tracing
- [ ] Async

### Reproduction

1. Configure Postgres to kill idle sessions after 10 seconds (I'm using the `ankane/pgvector` docker image):

```sql
ALTER SYSTEM SET idle_session_timeout TO '10000';
ALTER SYSTEM SET idle_in_transaction_session_timeout TO '10000';
```

2. Restart the database

3. Confirm the new settings are active

```sql
SHOW ALL;
``` 

4. Create a sample program that instantiate a `PGVector` class and search by similarity every 30 seconds, sleeping between each iteration

```python
if __name__ == '__main__':
    # Logging level
    logging.basicConfig(level=logging.DEBUG)

    load_dotenv()

    # PG Vector
    driver = os.environ["PGVECTOR_DRIVER"]
    host = os.environ["PGVECTOR_HOST"]
    port = os.environ["PGVECTOR_PORT"]
    database = os.environ["PGVECTOR_DATABASE"]
    user = os.environ["PGVECTOR_USER"]
    password = os.environ["PGVECTOR_PASSWORD"]

    collection_name = "state_of_the_union_test"
    connection_string = f'postgresql+{driver}://{user}:{password}@{host}:{port}/{database}'

    embeddings = OpenAIEmbeddings()

    loader = TextLoader("./state_of_the_union.txt")
    documents = loader.load()

    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    chunks = text_splitter.split_documents(documents)

    db = PGVector.from_documents(
        embedding=embeddings,
        documents=chunks,
        collection_name=collection_name,
        connection_string=connection_string,
    )

    for i in range(100):
        try:
            print('--> *** Searching by similarity ***')

            result = db.similarity_search_with_score("foo")
            print('--> *** Documents retrieved successfully ***')

            # for doc, score in result:
            #     print(f'{score}:{doc.page_content}')
        except Exception as e:
            print('--> *** Fail ***')
            print(str(e))

        print('--> *** Sleeping ***')
        time.sleep(30)

        print('\n\n')
```

5. Confirm in the console output that some requests fail because the idle connection was closed

```
(poc-pgvector-py3.11) ➜  poc-pgvector python -m poc_pgvector

DEBUG:openai:message='Request to OpenAI API' method=post path=https://.../v1/embeddings
DEBUG:urllib3.util.retry:Converted retries value: 2 -> Retry(total=2, connect=None, read=None, redirect=None, status=None)
DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): ...:443
DEBUG:urllib3.connectionpool:https://...:443 "POST /v1/embeddings HTTP/1.1" 200 None
DEBUG:openai:message='OpenAI API response' path=https://.../v1/embeddings processing_ms=211 request_id=ec4ef62b-da79-4055-a8da-66fce881a6e9 response_code=200

--> *** Searching by similarity ***

DEBUG:openai:message='Request to OpenAI API' method=post path=https://.../v1/embeddings
DEBUG:openai:api_version=None data='{"input": [[8134]], "model": "text-embedding-ada-002", "encoding_format": "base64"}' message='Post details'
DEBUG:urllib3.connectionpool:https://...:443 "POST /v1/embeddings HTTP/1.1" 200 None
DEBUG:openai:message='OpenAI API response' path=https://.../v1/embeddings processing_ms=40 request_id=f2b791ac-0c88-4362-b14c-d6a36084a932 response_code=200

--> *** Documents retrieved successfully ***
--> *** Sleeping ***



--> *** Searching by similarity ***
DEBUG:openai:message='Request to OpenAI API' method=post path=https://.../v1/embeddings
DEBUG:openai:api_version=None data='{"input": [[8134]], "model": "text-embedding-ada-002", "encoding_format": "base64"}' message='Post details'
DEBUG:urllib3.connectionpool:https://...:443 "POST /v1/embeddings HTTP/1.1" 200 None
DEBUG:openai:message='OpenAI API response' path=https://.../v1/embeddings processing_ms=30 request_id=de593b87-29ad-4342-81c1-3c495d5b4f56 response_code=200

--> *** Fail ***
(psycopg2.OperationalError) server closed the connection unexpectedly
        This probably means the server terminated abnormally
        before or while processing the request.

[SQL: SELECT langchain_pg_collection.name AS langchain_pg_collection_name, langchain_pg_collection.cmetadata AS langchain_pg_collection_cmetadata, langchain_pg_collection.uuid AS langchain_pg_collection_uuid 
FROM langchain_pg_collection 
WHERE langchain_pg_collection.name = %(name_1)s 
 LIMIT %(param_1)s]
[parameters: {'name_1': 'state_of_the_union_test', 'param_1': 1}]
(Background on this error at: https://sqlalche.me/e/20/e3q8)

--> *** Sleeping ***



--> *** Searching by similarity ***
DEBUG:openai:message='Request to OpenAI API' method=post path=https://.../v1/embeddings
DEBUG:openai:api_version=None data='{"input": [[8134]], "model": "text-embedding-ada-002", "encoding_format": "base64"}' message='Post details'
DEBUG:urllib3.connectionpool:https://...:443 "POST /v1/embeddings HTTP/1.1" 200 None
DEBUG:openai:message='OpenAI API response' path=https://.../v1/embeddings processing_ms=42 request_id=e2cb6145-8ceb-4127-b4ee-c5130ba42e8e response_code=200

--> *** Documents retrieved successfully ***
--> *** Sleeping ***



--> *** Searching by similarity ***
DEBUG:openai:message='Request to OpenAI API' method=post path=https://.../v1/embeddings
DEBUG:openai:api_version=None data='{"input": [[8134]], "model": "text-embedding-ada-002", "encoding_format": "base64"}' message='Post details'
DEBUG:urllib3.connectionpool:https://...:443 "POST /v1/embeddings HTTP/1.1" 200 None
DEBUG:openai:message='OpenAI API response' path=https://.../v1/embeddings processing_ms=38 request_id=05e15050-0c37-4f7c-a118-a3fafc2bf979 response_code=200

--> *** Fail ***
(psycopg2.OperationalError) server closed the connection unexpectedly
        This probably means the server terminated abnormally
        before or while processing the request.

[SQL: SELECT langchain_pg_collection.name AS langchain_pg_collection_name, langchain_pg_collection.cmetadata AS langchain_pg_collection_cmetadata, langchain_pg_collection.uuid AS langchain_pg_collection_uuid 
FROM langchain_pg_collection 
WHERE langchain_pg_collection.name = %(name_1)s 
 LIMIT %(param_1)s]
[parameters: {'name_1': 'state_of_the_union_test', 'param_1': 1}]
(Background on this error at: https://sqlalche.me/e/20/e3q8)
--> *** Sleeping ***



--> *** Searching by similarity ***
DEBUG:openai:message='Request to OpenAI API' method=post path=https://.../v1/embeddings
DEBUG:openai:api_version=None data='{"input": [[8134]], "model": "text-embedding-ada-002", "encoding_format": "base64"}' message='Post details'
DEBUG:urllib3.connectionpool:https://...:443 "POST /v1/embeddings HTTP/1.1" 200 None
DEBUG:openai:message='OpenAI API response' path=https://.../v1/embeddings processing_ms=28 request_id=2deecd8b-42de-42dd-a242-004d9d5811eb response_code=200
--> *** Documents retrieved successfully ***
--> *** Sleeping ***

...
```
 

### Expected behavior

- Use pooled connections to make feasible executing requests in parallel.
- Allow fine tuning the connection pool used by PGVector, like configuring the its maximum size, number of seconds to wait before giving up on getting a connection from the pool, and enabling [pessimist](https://docs.sqlalchemy.org/en/20/core/pooling.html#pool-disconnects-pessimistic) (liveness testing connections when borrowing them from the pool) or optimistic (connection recycle time) strategies for disconnection dandling.