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

chore: Migrating VectorDB to use OpenAI compatible identifier

Browse source 
Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
Francisco Javier Arceo 2025-07-03 22:02:06 -04:00
parent b094960199
commit ddb29b306c
9 changed files with 44 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

2
llama_stack/providers/remote/vector_io/chroma/chroma.py
View file

@ -217,7 +217,7 @@ class ChromaVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
embedding_model: str | None = None,
embedding_dimension: int | None = 384,
provider_id: str | None = None,
provider_vector_db_id: str | None = None,
provider_vector_db_id: str = "",
) -> VectorStoreObject:
raise NotImplementedError("OpenAI Vector Stores API is not supported in Chroma")

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

@ -247,7 +247,7 @@ class PGVectorVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
embedding_model: str | None = None,
embedding_dimension: int | None = 384,
provider_id: str | None = None,
provider_vector_db_id: str | None = None,
provider_vector_db_id: str = "",
) -> VectorStoreObject:
raise NotImplementedError("OpenAI Vector Stores API is not supported in PGVector")

24
llama_stack/providers/utils/memory/openai_vector_store_mixin.py
View file

@ -8,7 +8,6 @@ import asyncio
import logging
import mimetypes
import time
import uuid
from abc import ABC, abstractmethod
from typing import Any
@ -144,11 +143,12 @@ class OpenAIVectorStoreMixin(ABC):
embedding_model: str | None = None,
embedding_dimension: int | None = 384,
provider_id: str | None = None,
provider_vector_db_id: str | None = None,
provider_vector_db_id: str = "",
) -> VectorStoreObject:
"""Creates a vector store."""
store_id = name or str(uuid.uuid4())
created_at = int(time.time())
if provider_vector_db_id is None:
raise ValueError("Provider vector DB ID is required")
if provider_id is None:
raise ValueError("Provider ID is required")
@ -160,13 +160,13 @@ class OpenAIVectorStoreMixin(ABC):
if embedding_dimension is None:
raise ValueError("Embedding dimension is required")
provider_vector_db_id = provider_vector_db_id or store_id
vector_db = VectorDB(
identifier=store_id,
identifier=provider_vector_db_id,
embedding_dimension=embedding_dimension,
embedding_model=embedding_model,
provider_id=provider_id,
provider_resource_id=provider_vector_db_id,
provider_vector_db_name=name,
)
# Register the vector DB
await self.register_vector_db(vector_db)
@ -182,11 +182,11 @@ class OpenAIVectorStoreMixin(ABC):
in_progress=0,
total=0,
)
store_info = {
"id": store_id,
store_info: dict[str, Any] = {
"id": provider_vector_db_id,
"object": "vector_store",
"created_at": created_at,
"name": store_id,
"name": name,
"usage_bytes": 0,
"file_counts": file_counts.model_dump(),
"status": status,
@ -206,18 +206,18 @@ class OpenAIVectorStoreMixin(ABC):
store_info["metadata"] = metadata
# Save to persistent storage (provider-specific)
await self._save_openai_vector_store(store_id, store_info)
await self._save_openai_vector_store(provider_vector_db_id, store_info)
# Store in memory cache
self.openai_vector_stores[store_id] = store_info
self.openai_vector_stores[provider_vector_db_id] = store_info
# Now that our vector store is created, attach any files that were provided
file_ids = file_ids or []
tasks = [self.openai_attach_file_to_vector_store(store_id, file_id) for file_id in file_ids]
tasks = [self.openai_attach_file_to_vector_store(provider_vector_db_id, file_id) for file_id in file_ids]
await asyncio.gather(*tasks)
# Get the updated store info and return it
store_info = self.openai_vector_stores[store_id]
store_info = self.openai_vector_stores[provider_vector_db_id]
return VectorStoreObject.model_validate(store_info)
async def openai_list_vector_stores(