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

patch vectordb methods

Browse source 
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
Francisco Javier Arceo 2025-07-14 14:56:06 -04:00
parent 061911b993
commit 3d86619361
1 changed files with 23 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -216,28 +216,32 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
log.info("Connection to PGVector database server closed") log.info("Connection to PGVector database server closed")
async def register_vector_db(self, vector_db: VectorDB) -> None: async def register_vector_db(self, vector_db: VectorDB) -> None:
start_key = VECTOR_DBS_PREFIX # Persist vector DB metadata in the KV store
end_key = f"{VECTOR_DBS_PREFIX}\xff" assert self.kvstore is not None
stored_vector_dbs = await self.kvstore.values_in_range(start_key, end_key) key = f"{VECTOR_DBS_PREFIX}{vector_db.identifier}"
for vector_db_data in stored_vector_dbs: await self.kvstore.set(key=key, value=vector_db.model_dump_json())
vector_db = VectorDB.model_validate_json(vector_db_data)
# Upsert model metadata in Postgres
upsert_models(self.conn, [(vector_db.identifier, vector_db)])
# Create and cache the PGVector index table for the vector DB
index = VectorDBWithIndex( index = VectorDBWithIndex(
vector_db, vector_db,
index=PGVectorIndex( index=PGVectorIndex(vector_db, vector_db.embedding_dimension, self.conn, kvstore=self.kvstore),
vector_db,
vector_db.embedding_dimension,
self.conn,
kvstore=self.kvstore,
),
inference_api=self.inference_api, inference_api=self.inference_api,
) )
upsert_models(self.conn, [(vector_db.identifier, vector_db)])
self.cache[vector_db.identifier] = index self.cache[vector_db.identifier] = index
async def unregister_vector_db(self, vector_db_id: str) -> None: async def unregister_vector_db(self, vector_db_id: str) -> None:
# Remove provider index and cache
if vector_db_id in self.cache:
await self.cache[vector_db_id].index.delete() await self.cache[vector_db_id].index.delete()
del self.cache[vector_db_id] del self.cache[vector_db_id]
# Delete vector DB metadata from KV store
assert self.kvstore is not None
await self.kvstore.delete(key=f"{VECTOR_DBS_PREFIX}{vector_db_id}")
async def insert_chunks( async def insert_chunks(
self, self,
vector_db_id: str, vector_db_id: str,