From 83a51c7bfb60387a74e034ebc86f6610b63d3258 Mon Sep 17 00:00:00 2001 From: Yuan Tang Date: Mon, 3 Feb 2025 00:23:13 -0500 Subject: [PATCH] Properly close PGVector DB connection during shutdown() (#931) The connection to the DB was not closed 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)])