From c9bf1d7d0bb76cd02a8f1257b79fca575e6c4b46 Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Mon, 4 Nov 2024 17:01:09 -0800 Subject: [PATCH] pgvector fixes (#369) Co-authored-by: Dinesh Yeduguru --- llama_stack/providers/adapters/memory/pgvector/pgvector.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/llama_stack/providers/adapters/memory/pgvector/pgvector.py b/llama_stack/providers/adapters/memory/pgvector/pgvector.py index 87d6dbdab..0d188d944 100644 --- a/llama_stack/providers/adapters/memory/pgvector/pgvector.py +++ b/llama_stack/providers/adapters/memory/pgvector/pgvector.py @@ -46,8 +46,7 @@ def upsert_models(cur, keys_models: List[Tuple[str, BaseModel]]): def load_models(cur, cls): - query = "SELECT key, data FROM metadata_store" - cur.execute(query) + cur.execute("SELECT key, data FROM metadata_store") rows = cur.fetchall() return [parse_obj_as(cls, row["data"]) for row in rows] @@ -116,7 +115,6 @@ class PGVectorIndex(EmbeddingIndex): class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate): def __init__(self, config: PGVectorConfig) -> None: - print(f"Initializing PGVectorMemoryAdapter -> {config.host}:{config.port}") self.config = config self.cursor = None self.conn = None @@ -131,7 +129,8 @@ class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate): user=self.config.user, password=self.config.password, ) - self.cursor = self.conn.cursor() + self.conn.autocommit = True + self.cursor = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) version = check_extension_version(self.cursor) if version: