From 3444b32091ce82c86b7ec8f21a045fda4657359e Mon Sep 17 00:00:00 2001 From: Yuan Tang Date: Sun, 2 Feb 2025 20:39:18 -0500 Subject: [PATCH] Properly close PGVector DB connection during shutdown() Signed-off-by: Yuan Tang --- llama_stack/providers/remote/vector_io/pgvector/pgvector.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index f6c724648..e63eb1ccf 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py @@ -158,7 +158,9 @@ class PGVectorVectorDBAdapter(VectorIO, VectorDBsProtocolPrivate): raise RuntimeError("Could not connect to PGVector database server") from e async def shutdown(self) -> None: - pass + if self.conn is not None: + self.conn.close() + log.info("Connection to PGVector database server closed") async def register_vector_db(self, vector_db: VectorDB) -> None: upsert_models(self.cursor, [(vector_db.identifier, vector_db)])