phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-30 07:39:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

pgvector fixes

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-04 16:38:02 -08:00
parent ac93dd89cf
commit d5ed71c482
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):
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: