phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-06-28 10:54:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

pgvector fixes (#369)

Browse source 
Co-authored-by: Dinesh Yeduguru <dineshyv@fb.com>
This commit is contained in:
Dinesh Yeduguru 2024-11-04 17:01:09 -08:00 • committed by GitHub
parent c810a4184d
commit c9bf1d7d0b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

7
llama_stack/providers/adapters/memory/pgvector/pgvector.py
View file

@ -46,8 +46,7 @@ def upsert_models(cur, keys_models: List[Tuple[str, BaseModel]]):
def load_models(cur, cls): def load_models(cur, cls):
query = "SELECT key, data FROM metadata_store" cur.execute("SELECT key, data FROM metadata_store")
cur.execute(query)
rows = cur.fetchall() rows = cur.fetchall()
return [parse_obj_as(cls, row["data"]) for row in rows] return [parse_obj_as(cls, row["data"]) for row in rows]
@ -116,7 +115,6 @@ class PGVectorIndex(EmbeddingIndex):
class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate): class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
def __init__(self, config: PGVectorConfig) -> None: def __init__(self, config: PGVectorConfig) -> None:
print(f"Initializing PGVectorMemoryAdapter -> {config.host}:{config.port}")
self.config = config self.config = config
self.cursor = None self.cursor = None
self.conn = None self.conn = None
@ -131,7 +129,8 @@ class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
user=self.config.user, user=self.config.user,
password=self.config.password, 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) version = check_extension_version(self.cursor)
if version: if version: