diff --git a/llama_stack/core/routers/vector_io.py b/llama_stack/core/routers/vector_io.py index 41d32b10c..27e56fef2 100644 --- a/llama_stack/core/routers/vector_io.py +++ b/llama_stack/core/routers/vector_io.py @@ -135,6 +135,8 @@ class VectorIORouter(VectorIO): logger.debug(f"VectorIORouter.openai_create_vector_store: name={params.name}, provider_id={provider_id}") # If no embedding model is provided, use the first available one + # TODO: this branch will soon be deleted so you _must_ provide the embedding_model when + # creating a vector store if embedding_model is None: embedding_model_info = await self._get_first_embedding_model() if embedding_model_info is None: @@ -153,7 +155,14 @@ class VectorIORouter(VectorIO): ) provider = await self.routing_table.get_provider_impl(registered_vector_db.identifier) - # Pass params as-is to provider - it will extract what it needs from model_extra + # Update model_extra with registered values so provider uses the already-registered vector_db + if params.model_extra is None: + params.model_extra = {} + params.model_extra["provider_vector_db_id"] = registered_vector_db.provider_resource_id + params.model_extra["provider_id"] = registered_vector_db.provider_id + params.model_extra["embedding_model"] = embedding_model + params.model_extra["embedding_dimension"] = embedding_dimension + return await provider.openai_create_vector_store(params) async def openai_list_vector_stores( diff --git a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py index 992a29664..70bcbba32 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -10,8 +10,9 @@ import mimetypes import time import uuid from abc import ABC, abstractmethod -from typing import Any +from typing import Annotated, Any +from fastapi import Body from pydantic import TypeAdapter from llama_stack.apis.common.errors import VectorStoreNotFoundError @@ -342,7 +343,7 @@ class OpenAIVectorStoreMixin(ABC): async def openai_create_vector_store( self, - params: OpenAICreateVectorStoreRequestWithExtraBody, + params: Annotated[OpenAICreateVectorStoreRequestWithExtraBody, Body(...)], ) -> VectorStoreObject: """Creates a vector store.""" created_at = int(time.time()) @@ -978,7 +979,7 @@ class OpenAIVectorStoreMixin(ABC): async def openai_create_vector_store_file_batch( self, vector_store_id: str, - params: OpenAICreateVectorStoreFileBatchRequestWithExtraBody, + params: Annotated[OpenAICreateVectorStoreFileBatchRequestWithExtraBody, Body(...)], ) -> VectorStoreFileBatchObject: """Create a vector store file batch.""" if vector_store_id not in self.openai_vector_stores: diff --git a/llama_stack/providers/utils/memory/vector_store.py b/llama_stack/providers/utils/memory/vector_store.py index c0534a875..0375ecaaa 100644 --- a/llama_stack/providers/utils/memory/vector_store.py +++ b/llama_stack/providers/utils/memory/vector_store.py @@ -21,6 +21,7 @@ from llama_stack.apis.common.content_types import ( URL, InterleavedContent, ) +from llama_stack.apis.inference import OpenAIEmbeddingsRequestWithExtraBody from llama_stack.apis.tools import RAGDocument from llama_stack.apis.vector_dbs import VectorDB from llama_stack.apis.vector_io import Chunk, ChunkMetadata, QueryChunksResponse @@ -274,10 +275,11 @@ class VectorDBWithIndex: _validate_embedding(c.embedding, i, self.vector_db.embedding_dimension) if chunks_to_embed: - resp = await self.inference_api.openai_embeddings( - self.vector_db.embedding_model, - [c.content for c in chunks_to_embed], + params = OpenAIEmbeddingsRequestWithExtraBody( + model=self.vector_db.embedding_model, + input=[c.content for c in chunks_to_embed], ) + resp = await self.inference_api.openai_embeddings(params) for c, data in zip(chunks_to_embed, resp.data, strict=False): c.embedding = data.embedding @@ -316,7 +318,11 @@ class VectorDBWithIndex: if mode == "keyword": return await self.index.query_keyword(query_string, k, score_threshold) - embeddings_response = await self.inference_api.openai_embeddings(self.vector_db.embedding_model, [query_string]) + params = OpenAIEmbeddingsRequestWithExtraBody( + model=self.vector_db.embedding_model, + input=[query_string], + ) + embeddings_response = await self.inference_api.openai_embeddings(params) query_vector = np.array(embeddings_response.data[0].embedding, dtype=np.float32) if mode == "hybrid": return await self.index.query_hybrid(