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

updating to use execute_values

Browse source 
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
Francisco Javier Arceo 2025-07-15 11:34:31 -04:00
parent dd6e6f9ac2
commit f097da7a31
1 changed files with 31 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

51
llama_stack/providers/remote/vector_io/pgvector/pgvector.py
View file

@ -217,10 +217,6 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
async def register_vector_db(self, vector_db: VectorDB) -> None: async def register_vector_db(self, vector_db: VectorDB) -> None:
# Persist vector DB metadata in the KV store # Persist vector DB metadata in the KV store
assert self.kvstore is not None
key = f"{VECTOR_DBS_PREFIX}{vector_db.identifier}"
await self.kvstore.set(key=key, value=vector_db.model_dump_json())
# Upsert model metadata in Postgres # Upsert model metadata in Postgres
upsert_models(self.conn, [(vector_db.identifier, vector_db)]) upsert_models(self.conn, [(vector_db.identifier, vector_db)])
@ -277,7 +273,7 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
if self.conn is None: if self.conn is None:
raise RuntimeError("PostgreSQL connection is not initialized") raise RuntimeError("PostgreSQL connection is not initialized")
try: try:
with self.conn.cursor() as cur: with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
cur.execute( cur.execute(
""" """
CREATE TABLE IF NOT EXISTS openai_vector_store_files ( CREATE TABLE IF NOT EXISTS openai_vector_store_files (
@ -298,16 +294,26 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
) )
""" """
) )
cur.execute( # Insert file metadata
"INSERT INTO openai_vector_store_files (store_id, file_id, metadata) VALUES (%s, %s, %s)" files_query = sql.SQL(
" ON CONFLICT (store_id, file_id) DO UPDATE SET metadata = EXCLUDED.metadata", """
(store_id, file_id, Json(file_info)), INSERT INTO openai_vector_store_files (store_id, file_id, metadata)
VALUES %s
ON CONFLICT (store_id, file_id) DO UPDATE SET metadata = EXCLUDED.metadata
"""
) )
cur.execute( files_values = [(store_id, file_id, Json(file_info))]
"INSERT INTO openai_vector_store_files_contents (store_id, file_id, contents) VALUES (%s, %s, %s)" execute_values(cur, files_query, files_values, template="(%s, %s, %s)")
" ON CONFLICT (store_id, file_id) DO UPDATE SET contents = EXCLUDED.contents", # Insert file contents
(store_id, file_id, Json(file_contents)), contents_query = sql.SQL(
"""
INSERT INTO openai_vector_store_files_contents (store_id, file_id, contents)
VALUES %s
ON CONFLICT (store_id, file_id) DO UPDATE SET contents = EXCLUDED.contents
"""
) )
contents_values = [(store_id, file_id, Json(file_contents))]
execute_values(cur, contents_query, contents_values, template="(%s, %s, %s)")
except Exception as e: except Exception as e:
log.error(f"Error saving openai vector store file {file_id} for store {store_id}: {e}") log.error(f"Error saving openai vector store file {file_id} for store {store_id}: {e}")
raise raise
@ -317,7 +323,7 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
if self.conn is None: if self.conn is None:
raise RuntimeError("PostgreSQL connection is not initialized") raise RuntimeError("PostgreSQL connection is not initialized")
try: try:
with self.conn.cursor() as cur: with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
cur.execute( cur.execute(
"SELECT metadata FROM openai_vector_store_files WHERE store_id = %s AND file_id = %s", "SELECT metadata FROM openai_vector_store_files WHERE store_id = %s AND file_id = %s",
(store_id, file_id), (store_id, file_id),
@ -333,7 +339,7 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
if self.conn is None: if self.conn is None:
raise RuntimeError("PostgreSQL connection is not initialized") raise RuntimeError("PostgreSQL connection is not initialized")
try: try:
with self.conn.cursor() as cur: with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
cur.execute( cur.execute(
"SELECT contents FROM openai_vector_store_files_contents WHERE store_id = %s AND file_id = %s", "SELECT contents FROM openai_vector_store_files_contents WHERE store_id = %s AND file_id = %s",
(store_id, file_id), (store_id, file_id),
@ -349,11 +355,16 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
if self.conn is None: if self.conn is None:
raise RuntimeError("PostgreSQL connection is not initialized") raise RuntimeError("PostgreSQL connection is not initialized")
try: try:
with self.conn.cursor() as cur: with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
cur.execute( query = sql.SQL(
"UPDATE openai_vector_store_files SET metadata = %s WHERE store_id = %s AND file_id = %s", """
(Json(file_info), store_id, file_id), INSERT INTO openai_vector_store_files (store_id, file_id, metadata)
VALUES %s
ON CONFLICT (store_id, file_id) DO UPDATE SET metadata = EXCLUDED.metadata
"""
) )
values = [(store_id, file_id, Json(file_info))]
execute_values(cur, query, values, template="(%s, %s, %s)")
except Exception as e: except Exception as e:
log.error(f"Error updating openai vector store file {file_id} for store {store_id}: {e}") log.error(f"Error updating openai vector store file {file_id} for store {store_id}: {e}")
raise raise
@ -363,7 +374,7 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
if self.conn is None: if self.conn is None:
raise RuntimeError("PostgreSQL connection is not initialized") raise RuntimeError("PostgreSQL connection is not initialized")
try: try:
with self.conn.cursor() as cur: with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
cur.execute( cur.execute(
"DELETE FROM openai_vector_store_files WHERE store_id = %s AND file_id = %s", "DELETE FROM openai_vector_store_files WHERE store_id = %s AND file_id = %s",
(store_id, file_id), (store_id, file_id),