llama-stack 0.3.5__py3-none-any.whl → 0.4.0__py3-none-any.whl

llama_stack/core/routers/vector_io.py

@@ -10,19 +10,31 @@ from typing import Annotated, Any
 
 from fastapi import Body
 
-from llama_stack.apis.common.content_types import InterleavedContent
-from llama_stack.apis.models import ModelType
-from llama_stack.apis.vector_io import (
-    Chunk,
+from llama_stack.core.datatypes import VectorStoresConfig
+from llama_stack.log import get_logger
+from llama_stack_api import (
+    EmbeddedChunk,
+    HealthResponse,
+    HealthStatus,
+    Inference,
+    InterleavedContent,
+    ModelNotFoundError,
+    ModelType,
+    ModelTypeError,
+    OpenAIChatCompletionRequestWithExtraBody,
     OpenAICreateVectorStoreFileBatchRequestWithExtraBody,
     OpenAICreateVectorStoreRequestWithExtraBody,
+    OpenAIUserMessageParam,
     QueryChunksResponse,
+    RoutingTable,
     SearchRankingOptions,
     VectorIO,
     VectorStoreChunkingStrategy,
+    VectorStoreChunkingStrategyStatic,
+    VectorStoreChunkingStrategyStaticConfig,
     VectorStoreDeleteResponse,
     VectorStoreFileBatchObject,
-    VectorStoreFileContentsResponse,
+    VectorStoreFileContentResponse,
     VectorStoreFileDeleteResponse,
     VectorStoreFileObject,
     VectorStoreFilesListInBatchResponse,
@@ -31,9 +43,6 @@ from llama_stack.apis.vector_io import (
     VectorStoreObject,
     VectorStoreSearchResponsePage,
 )
-from llama_stack.core.datatypes import VectorStoresConfig
-from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import HealthResponse, HealthStatus, RoutingTable
 
 logger = get_logger(name=__name__, category="core::routers")
 
@@ -45,10 +54,11 @@ class VectorIORouter(VectorIO):
         self,
         routing_table: RoutingTable,
         vector_stores_config: VectorStoresConfig | None = None,
+        inference_api: Inference | None = None,
     ) -> None:
-        logger.debug("Initializing VectorIORouter")
         self.routing_table = routing_table
         self.vector_stores_config = vector_stores_config
+        self.inference_api = inference_api
 
     async def initialize(self) -> None:
         logger.debug("VectorIORouter.initialize")
@@ -58,6 +68,46 @@ class VectorIORouter(VectorIO):
         logger.debug("VectorIORouter.shutdown")
         pass
 
+    async def _rewrite_query_for_search(self, query: str) -> str:
+        """Rewrite a search query using the configured LLM model for better retrieval results."""
+        if (
+            not self.vector_stores_config
+            or not self.vector_stores_config.rewrite_query_params
+            or not self.vector_stores_config.rewrite_query_params.model
+        ):
+            logger.warning(
+                "User is trying to use vector_store query rewriting, but it is not configured. Please configure rewrite_query_params.model in vector_stores config."
+            )
+            raise ValueError("Query rewriting is not available")
+
+        if not self.inference_api:
+            logger.warning("Query rewriting requires inference API but it is not available")
+            raise ValueError("Query rewriting is not available")
+
+        model = self.vector_stores_config.rewrite_query_params.model
+        model_id = f"{model.provider_id}/{model.model_id}"
+
+        prompt = self.vector_stores_config.rewrite_query_params.prompt.format(query=query)
+
+        request = OpenAIChatCompletionRequestWithExtraBody(
+            model=model_id,
+            messages=[OpenAIUserMessageParam(role="user", content=prompt)],
+            max_tokens=self.vector_stores_config.rewrite_query_params.max_tokens or 100,
+            temperature=self.vector_stores_config.rewrite_query_params.temperature or 0.3,
+        )
+
+        try:
+            response = await self.inference_api.openai_chat_completion(request)
+            content = response.choices[0].message.content
+            if content is None:
+                logger.error(f"LLM returned None content for query rewriting. Model: {model_id}")
+                raise RuntimeError("Query rewrite failed due to an internal error")
+            rewritten_query: str = content.strip()
+            return rewritten_query
+        except Exception as e:
+            logger.error(f"Query rewrite failed with LLM call error. Model: {model_id}, Error: {e}")
+            raise RuntimeError("Query rewrite failed due to an internal error") from e
+
     async def _get_embedding_model_dimension(self, embedding_model_id: str) -> int:
         """Get the embedding dimension for a specific embedding model."""
         all_models = await self.routing_table.get_all_with_type("model")
@@ -73,27 +123,25 @@ class VectorIORouter(VectorIO):
 
     async def insert_chunks(
         self,
-        vector_db_id: str,
-        chunks: list[Chunk],
+        vector_store_id: str,
+        chunks: list[EmbeddedChunk],
         ttl_seconds: int | None = None,
     ) -> None:
         doc_ids = [chunk.document_id for chunk in chunks[:3]]
         logger.debug(
-            f"VectorIORouter.insert_chunks: {vector_db_id}, {len(chunks)} chunks, "
+            f"VectorIORouter.insert_chunks: {vector_store_id}, {len(chunks)} chunks, "
             f"ttl_seconds={ttl_seconds}, chunk_ids={doc_ids}{' and more...' if len(chunks) > 3 else ''}"
         )
-        provider = await self.routing_table.get_provider_impl(vector_db_id)
-        return await provider.insert_chunks(vector_db_id, chunks, ttl_seconds)
+        return await self.routing_table.insert_chunks(vector_store_id, chunks, ttl_seconds)
 
     async def query_chunks(
         self,
-        vector_db_id: str,
+        vector_store_id: str,
         query: InterleavedContent,
         params: dict[str, Any] | None = None,
     ) -> QueryChunksResponse:
-        logger.debug(f"VectorIORouter.query_chunks: {vector_db_id}")
-        provider = await self.routing_table.get_provider_impl(vector_db_id)
-        return await provider.query_chunks(vector_db_id, query, params)
+        logger.debug(f"VectorIORouter.query_chunks: {vector_store_id}")
+        return await self.routing_table.query_chunks(vector_store_id, query, params)
 
     # OpenAI Vector Stores API endpoints
     async def openai_create_vector_store(
@@ -120,6 +168,14 @@ class VectorIORouter(VectorIO):
         if embedding_model is not None and embedding_dimension is None:
             embedding_dimension = await self._get_embedding_model_dimension(embedding_model)
 
+        # Validate that embedding model exists and is of the correct type
+        if embedding_model is not None:
+            model = await self.routing_table.get_object_by_identifier("model", embedding_model)
+            if model is None:
+                raise ModelNotFoundError(embedding_model)
+            if model.model_type != ModelType.embedding:
+                raise ModelTypeError(embedding_model, model.model_type, ModelType.embedding)
+
         # Auto-select provider if not specified
         if provider_id is None:
             num_providers = len(self.routing_table.impls_by_provider_id)
@@ -167,6 +223,13 @@ class VectorIORouter(VectorIO):
         if embedding_dimension is not None:
             params.model_extra["embedding_dimension"] = embedding_dimension
 
+        # Set chunking strategy explicitly if not provided
+        if params.chunking_strategy is None or params.chunking_strategy.type == "auto":
+            # actualize the chunking strategy to static
+            params.chunking_strategy = VectorStoreChunkingStrategyStatic(
+                static=VectorStoreChunkingStrategyStaticConfig()
+            )
+
         return await provider.openai_create_vector_store(params)
 
     async def openai_list_vector_stores(
@@ -183,9 +246,8 @@ class VectorIORouter(VectorIO):
         all_stores = []
         for vector_store in vector_stores:
             try:
-                provider = await self.routing_table.get_provider_impl(vector_store.identifier)
-                vector_store = await provider.openai_retrieve_vector_store(vector_store.identifier)
-                all_stores.append(vector_store)
+                vector_store_obj = await self.routing_table.openai_retrieve_vector_store(vector_store.identifier)
+                all_stores.append(vector_store_obj)
             except Exception as e:
                 logger.error(f"Error retrieving vector store {vector_store.identifier}: {e}")
                 continue
@@ -227,8 +289,7 @@ class VectorIORouter(VectorIO):
         vector_store_id: str,
     ) -> VectorStoreObject:
         logger.debug(f"VectorIORouter.openai_retrieve_vector_store: {vector_store_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_retrieve_vector_store(vector_store_id)
+        return await self.routing_table.openai_retrieve_vector_store(vector_store_id)
 
     async def openai_update_vector_store(
         self,
@@ -238,8 +299,14 @@ class VectorIORouter(VectorIO):
         metadata: dict[str, Any] | None = None,
     ) -> VectorStoreObject:
         logger.debug(f"VectorIORouter.openai_update_vector_store: {vector_store_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_update_vector_store(
+
+        # Check if provider_id is being changed (not supported)
+        if metadata and "provider_id" in metadata:
+            current_store = await self.routing_table.get_object_by_identifier("vector_store", vector_store_id)
+            if current_store and current_store.provider_id != metadata["provider_id"]:
+                raise ValueError("provider_id cannot be changed after vector store creation")
+
+        return await self.routing_table.openai_update_vector_store(
             vector_store_id=vector_store_id,
             name=name,
             expires_after=expires_after,
@@ -264,14 +331,23 @@ class VectorIORouter(VectorIO):
         search_mode: str | None = "vector",
     ) -> VectorStoreSearchResponsePage:
         logger.debug(f"VectorIORouter.openai_search_vector_store: {vector_store_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_search_vector_store(
+
+        # Handle query rewriting at the router level
+        search_query = query
+        if rewrite_query:
+            if isinstance(query, list):
+                original_query = " ".join(query)
+            else:
+                original_query = query
+            search_query = await self._rewrite_query_for_search(original_query)
+
+        return await self.routing_table.openai_search_vector_store(
             vector_store_id=vector_store_id,
-            query=query,
+            query=search_query,
             filters=filters,
             max_num_results=max_num_results,
             ranking_options=ranking_options,
-            rewrite_query=rewrite_query,
+            rewrite_query=False,  # Already handled at router level
             search_mode=search_mode,
         )
 
@@ -283,8 +359,9 @@ class VectorIORouter(VectorIO):
         chunking_strategy: VectorStoreChunkingStrategy | None = None,
     ) -> VectorStoreFileObject:
         logger.debug(f"VectorIORouter.openai_attach_file_to_vector_store: {vector_store_id}, {file_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_attach_file_to_vector_store(
+        if chunking_strategy is None or chunking_strategy.type == "auto":
+            chunking_strategy = VectorStoreChunkingStrategyStatic(static=VectorStoreChunkingStrategyStaticConfig())
+        return await self.routing_table.openai_attach_file_to_vector_store(
             vector_store_id=vector_store_id,
             file_id=file_id,
             attributes=attributes,
@@ -301,8 +378,7 @@ class VectorIORouter(VectorIO):
         filter: VectorStoreFileStatus | None = None,
     ) -> list[VectorStoreFileObject]:
         logger.debug(f"VectorIORouter.openai_list_files_in_vector_store: {vector_store_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_list_files_in_vector_store(
+        return await self.routing_table.openai_list_files_in_vector_store(
             vector_store_id=vector_store_id,
             limit=limit,
             order=order,
@@ -317,8 +393,7 @@ class VectorIORouter(VectorIO):
         file_id: str,
     ) -> VectorStoreFileObject:
         logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file: {vector_store_id}, {file_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_retrieve_vector_store_file(
+        return await self.routing_table.openai_retrieve_vector_store_file(
             vector_store_id=vector_store_id,
             file_id=file_id,
         )
@@ -327,12 +402,19 @@ class VectorIORouter(VectorIO):
         self,
         vector_store_id: str,
         file_id: str,
-    ) -> VectorStoreFileContentsResponse:
-        logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_contents: {vector_store_id}, {file_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_retrieve_vector_store_file_contents(
+        include_embeddings: bool | None = False,
+        include_metadata: bool | None = False,
+    ) -> VectorStoreFileContentResponse:
+        logger.debug(
+            f"VectorIORouter.openai_retrieve_vector_store_file_contents: {vector_store_id}, {file_id}, "
+            f"include_embeddings={include_embeddings}, include_metadata={include_metadata}"
+        )
+
+        return await self.routing_table.openai_retrieve_vector_store_file_contents(
             vector_store_id=vector_store_id,
             file_id=file_id,
+            include_embeddings=include_embeddings,
+            include_metadata=include_metadata,
         )
 
     async def openai_update_vector_store_file(
@@ -342,8 +424,7 @@ class VectorIORouter(VectorIO):
         attributes: dict[str, Any],
     ) -> VectorStoreFileObject:
         logger.debug(f"VectorIORouter.openai_update_vector_store_file: {vector_store_id}, {file_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_update_vector_store_file(
+        return await self.routing_table.openai_update_vector_store_file(
             vector_store_id=vector_store_id,
             file_id=file_id,
             attributes=attributes,
@@ -355,8 +436,7 @@ class VectorIORouter(VectorIO):
         file_id: str,
     ) -> VectorStoreFileDeleteResponse:
         logger.debug(f"VectorIORouter.openai_delete_vector_store_file: {vector_store_id}, {file_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_delete_vector_store_file(
+        return await self.routing_table.openai_delete_vector_store_file(
             vector_store_id=vector_store_id,
             file_id=file_id,
         )
@@ -392,8 +472,10 @@ class VectorIORouter(VectorIO):
         logger.debug(
             f"VectorIORouter.openai_create_vector_store_file_batch: {vector_store_id}, {len(params.file_ids)} files"
         )
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_create_vector_store_file_batch(vector_store_id, params)
+        return await self.routing_table.openai_create_vector_store_file_batch(
+            vector_store_id=vector_store_id,
+            params=params,
+        )
 
     async def openai_retrieve_vector_store_file_batch(
         self,
@@ -401,8 +483,7 @@ class VectorIORouter(VectorIO):
         vector_store_id: str,
     ) -> VectorStoreFileBatchObject:
         logger.debug(f"VectorIORouter.openai_retrieve_vector_store_file_batch: {batch_id}, {vector_store_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_retrieve_vector_store_file_batch(
+        return await self.routing_table.openai_retrieve_vector_store_file_batch(
             batch_id=batch_id,
             vector_store_id=vector_store_id,
         )
@@ -418,8 +499,7 @@ class VectorIORouter(VectorIO):
         order: str | None = "desc",
     ) -> VectorStoreFilesListInBatchResponse:
         logger.debug(f"VectorIORouter.openai_list_files_in_vector_store_file_batch: {batch_id}, {vector_store_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_list_files_in_vector_store_file_batch(
+        return await self.routing_table.openai_list_files_in_vector_store_file_batch(
             batch_id=batch_id,
             vector_store_id=vector_store_id,
             after=after,
@@ -435,8 +515,7 @@ class VectorIORouter(VectorIO):
         vector_store_id: str,
     ) -> VectorStoreFileBatchObject:
         logger.debug(f"VectorIORouter.openai_cancel_vector_store_file_batch: {batch_id}, {vector_store_id}")
-        provider = await self.routing_table.get_provider_impl(vector_store_id)
-        return await provider.openai_cancel_vector_store_file_batch(
+        return await self.routing_table.openai_cancel_vector_store_file_batch(
             batch_id=batch_id,
             vector_store_id=vector_store_id,
         )

llama_stack/core/routing_tables/benchmarks.py

@@ -4,13 +4,20 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Any
 
-from llama_stack.apis.benchmarks import Benchmark, Benchmarks, ListBenchmarksResponse
 from llama_stack.core.datatypes import (
     BenchmarkWithOwner,
 )
 from llama_stack.log import get_logger
+from llama_stack_api import (
+    Benchmark,
+    Benchmarks,
+    GetBenchmarkRequest,
+    ListBenchmarksRequest,
+    ListBenchmarksResponse,
+    RegisterBenchmarkRequest,
+    UnregisterBenchmarkRequest,
+)
 
 from .common import CommonRoutingTableImpl
 
@@ -18,26 +25,21 @@ logger = get_logger(name=__name__, category="core::routing_tables")
 
 
 class BenchmarksRoutingTable(CommonRoutingTableImpl, Benchmarks):
-    async def list_benchmarks(self) -> ListBenchmarksResponse:
+    async def list_benchmarks(self, request: ListBenchmarksRequest) -> ListBenchmarksResponse:
         return ListBenchmarksResponse(data=await self.get_all_with_type("benchmark"))
 
-    async def get_benchmark(self, benchmark_id: str) -> Benchmark:
-        benchmark = await self.get_object_by_identifier("benchmark", benchmark_id)
+    async def get_benchmark(self, request: GetBenchmarkRequest) -> Benchmark:
+        benchmark = await self.get_object_by_identifier("benchmark", request.benchmark_id)
         if benchmark is None:
-            raise ValueError(f"Benchmark '{benchmark_id}' not found")
+            raise ValueError(f"Benchmark '{request.benchmark_id}' not found")
         return benchmark
 
     async def register_benchmark(
         self,
-        benchmark_id: str,
-        dataset_id: str,
-        scoring_functions: list[str],
-        metadata: dict[str, Any] | None = None,
-        provider_benchmark_id: str | None = None,
-        provider_id: str | None = None,
+        request: RegisterBenchmarkRequest,
     ) -> None:
-        if metadata is None:
-            metadata = {}
+        metadata = request.metadata if request.metadata is not None else {}
+        provider_id = request.provider_id
         if provider_id is None:
             if len(self.impls_by_provider_id) == 1:
                 provider_id = list(self.impls_by_provider_id.keys())[0]
@@ -45,18 +47,20 @@ class BenchmarksRoutingTable(CommonRoutingTableImpl, Benchmarks):
                 raise ValueError(
                     "No provider specified and multiple providers available. Please specify a provider_id."
                 )
+        provider_benchmark_id = request.provider_benchmark_id
         if provider_benchmark_id is None:
-            provider_benchmark_id = benchmark_id
+            provider_benchmark_id = request.benchmark_id
         benchmark = BenchmarkWithOwner(
-            identifier=benchmark_id,
-            dataset_id=dataset_id,
-            scoring_functions=scoring_functions,
+            identifier=request.benchmark_id,
+            dataset_id=request.dataset_id,
+            scoring_functions=request.scoring_functions,
             metadata=metadata,
             provider_id=provider_id,
             provider_resource_id=provider_benchmark_id,
         )
         await self.register_object(benchmark)
 
-    async def unregister_benchmark(self, benchmark_id: str) -> None:
-        existing_benchmark = await self.get_benchmark(benchmark_id)
+    async def unregister_benchmark(self, request: UnregisterBenchmarkRequest) -> None:
+        get_request = GetBenchmarkRequest(benchmark_id=request.benchmark_id)
+        existing_benchmark = await self.get_benchmark(get_request)
         await self.unregister_object(existing_benchmark)

llama_stack/core/routing_tables/common.py

@@ -6,9 +6,6 @@
 
 from typing import Any
 
-from llama_stack.apis.common.errors import ModelNotFoundError
-from llama_stack.apis.models import Model
-from llama_stack.apis.resource import ResourceType
 from llama_stack.core.access_control.access_control import AccessDeniedError, is_action_allowed
 from llama_stack.core.access_control.datatypes import Action
 from llama_stack.core.datatypes import (
@@ -21,7 +18,7 @@ from llama_stack.core.datatypes import (
 from llama_stack.core.request_headers import get_authenticated_user
 from llama_stack.core.store import DistributionRegistry
 from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import Api, RoutingTable
+from llama_stack_api import Api, Model, ModelNotFoundError, ResourceType, RoutingTable
 
 logger = get_logger(name=__name__, category="core::routing_tables")
 

llama_stack/core/routing_tables/datasets.py

@@ -5,24 +5,26 @@
 # the root directory of this source tree.
 
 import uuid
-from typing import Any
 
-from llama_stack.apis.common.errors import DatasetNotFoundError
-from llama_stack.apis.datasets import (
+from llama_stack.core.datatypes import (
+    DatasetWithOwner,
+)
+from llama_stack.log import get_logger
+from llama_stack_api import (
     Dataset,
-    DatasetPurpose,
-    Datasets,
+    DatasetNotFoundError,
     DatasetType,
-    DataSource,
     ListDatasetsResponse,
+    ResourceType,
     RowsDataSource,
     URIDataSource,
 )
-from llama_stack.apis.resource import ResourceType
-from llama_stack.core.datatypes import (
-    DatasetWithOwner,
+from llama_stack_api.datasets.api import (
+    Datasets,
+    GetDatasetRequest,
+    RegisterDatasetRequest,
+    UnregisterDatasetRequest,
 )
-from llama_stack.log import get_logger
 
 from .common import CommonRoutingTableImpl
 
@@ -33,19 +35,17 @@ class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
     async def list_datasets(self) -> ListDatasetsResponse:
         return ListDatasetsResponse(data=await self.get_all_with_type(ResourceType.dataset.value))
 
-    async def get_dataset(self, dataset_id: str) -> Dataset:
-        dataset = await self.get_object_by_identifier("dataset", dataset_id)
+    async def get_dataset(self, request: GetDatasetRequest) -> Dataset:
+        dataset = await self.get_object_by_identifier("dataset", request.dataset_id)
         if dataset is None:
-            raise DatasetNotFoundError(dataset_id)
+            raise DatasetNotFoundError(request.dataset_id)
         return dataset
 
-    async def register_dataset(
-        self,
-        purpose: DatasetPurpose,
-        source: DataSource,
-        metadata: dict[str, Any] | None = None,
-        dataset_id: str | None = None,
-    ) -> Dataset:
+    async def register_dataset(self, request: RegisterDatasetRequest) -> Dataset:
+        purpose = request.purpose
+        source = request.source
+        metadata = request.metadata
+        dataset_id = request.dataset_id
         if isinstance(source, dict):
             if source["type"] == "uri":
                 source = URIDataSource.parse_obj(source)
@@ -86,6 +86,6 @@ class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
         await self.register_object(dataset)
         return dataset
 
-    async def unregister_dataset(self, dataset_id: str) -> None:
-        dataset = await self.get_dataset(dataset_id)
+    async def unregister_dataset(self, request: UnregisterDatasetRequest) -> None:
+        dataset = await self.get_dataset(GetDatasetRequest(dataset_id=request.dataset_id))
         await self.unregister_object(dataset)