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

llama_stack/providers/remote/vector_io/weaviate/weaviate.py

@@ -12,17 +12,9 @@ from numpy.typing import NDArray
 from weaviate.classes.init import Auth
 from weaviate.classes.query import Filter, HybridFusion
 
-from llama_stack.apis.common.content_types import InterleavedContent
-from llama_stack.apis.common.errors import VectorStoreNotFoundError
-from llama_stack.apis.files import Files
-from llama_stack.apis.inference import Inference
-from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO
-from llama_stack.apis.vector_stores import VectorStore
 from llama_stack.core.request_headers import NeedsRequestProviderData
+from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import VectorStoresProtocolPrivate
-from llama_stack.providers.utils.kvstore import kvstore_impl
-from llama_stack.providers.utils.kvstore.api import KVStore
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import (
     RERANKER_TYPE_RRF,
@@ -30,7 +22,20 @@ from llama_stack.providers.utils.memory.vector_store import (
     EmbeddingIndex,
     VectorStoreWithIndex,
 )
+from llama_stack.providers.utils.vector_io import load_embedded_chunk_with_backward_compat
 from llama_stack.providers.utils.vector_io.vector_utils import sanitize_collection_name
+from llama_stack_api import (
+    EmbeddedChunk,
+    Files,
+    Inference,
+    InterleavedContent,
+    QueryChunksResponse,
+    VectorIO,
+    VectorStore,
+    VectorStoreNotFoundError,
+    VectorStoresProtocolPrivate,
+)
+from llama_stack_api.internal.kvstore import KVStore
 
 from .config import WeaviateVectorIOConfig
 
@@ -53,20 +58,19 @@ class WeaviateIndex(EmbeddingIndex):
     async def initialize(self):
         pass
 
-    async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray):
-        assert len(chunks) == len(embeddings), (
-            f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
-        )
+    async def add_chunks(self, chunks: list[EmbeddedChunk]):
+        if not chunks:
+            return
 
         data_objects = []
-        for chunk, embedding in zip(chunks, embeddings, strict=False):
+        for chunk in chunks:
             data_objects.append(
                 wvc.data.DataObject(
                     properties={
                         "chunk_id": chunk.chunk_id,
                         "chunk_content": chunk.model_dump_json(),
                     },
-                    vector=embedding.tolist(),
+                    vector=chunk.embedding,  # Already a list[float]
                 )
             )
 
@@ -112,7 +116,7 @@ class WeaviateIndex(EmbeddingIndex):
             chunk_json = doc.properties["chunk_content"]
             try:
                 chunk_dict = json.loads(chunk_json)
-                chunk = Chunk(**chunk_dict)
+                chunk = load_embedded_chunk_with_backward_compat(chunk_dict)
             except Exception:
                 log.exception(f"Failed to parse document: {chunk_json}")
                 continue
@@ -172,7 +176,7 @@ class WeaviateIndex(EmbeddingIndex):
             chunk_json = doc.properties["chunk_content"]
             try:
                 chunk_dict = json.loads(chunk_json)
-                chunk = Chunk(**chunk_dict)
+                chunk = load_embedded_chunk_with_backward_compat(chunk_dict)
             except Exception:
                 log.exception(f"Failed to parse document: {chunk_json}")
                 continue
@@ -241,7 +245,7 @@ class WeaviateIndex(EmbeddingIndex):
             chunk_json = doc.properties["chunk_content"]
             try:
                 chunk_dict = json.loads(chunk_json)
-                chunk = Chunk(**chunk_dict)
+                chunk = load_embedded_chunk_with_backward_compat(chunk_dict)
             except Exception:
                 log.exception(f"Failed to parse document: {chunk_json}")
                 continue
@@ -259,9 +263,8 @@ class WeaviateIndex(EmbeddingIndex):
 
 class WeaviateVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, NeedsRequestProviderData, VectorStoresProtocolPrivate):
     def __init__(self, config: WeaviateVectorIOConfig, inference_api: Inference, files_api: Files | None) -> None:
-        super().__init__(files_api=files_api, kvstore=None)
+        super().__init__(inference_api=inference_api, files_api=files_api, kvstore=None)
         self.config = config
-        self.inference_api = inference_api
         self.client_cache = {}
         self.cache = {}
         self.vector_store_table = None
@@ -369,19 +372,21 @@ class WeaviateVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, NeedsRequestProv
         self.cache[vector_store_id] = index
         return index
 
-    async def insert_chunks(self, vector_db_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None:
-        index = await self._get_and_cache_vector_store_index(vector_db_id)
+    async def insert_chunks(
+        self, vector_store_id: str, chunks: list[EmbeddedChunk], ttl_seconds: int | None = None
+    ) -> None:
+        index = await self._get_and_cache_vector_store_index(vector_store_id)
         if not index:
-            raise VectorStoreNotFoundError(vector_db_id)
+            raise VectorStoreNotFoundError(vector_store_id)
 
         await index.insert_chunks(chunks)
 
     async def query_chunks(
-        self, vector_db_id: str, query: InterleavedContent, params: dict[str, Any] | None = None
+        self, vector_store_id: str, query: InterleavedContent, params: dict[str, Any] | None = None
     ) -> QueryChunksResponse:
-        index = await self._get_and_cache_vector_store_index(vector_db_id)
+        index = await self._get_and_cache_vector_store_index(vector_store_id)
         if not index:
-            raise VectorStoreNotFoundError(vector_db_id)
+            raise VectorStoreNotFoundError(vector_store_id)
 
         return await index.query_chunks(query, params)
 

llama_stack/providers/utils/common/data_schema_validator.py

@@ -7,12 +7,8 @@
 from enum import Enum
 from typing import Any
 
-from llama_stack.apis.common.type_system import (
-    ChatCompletionInputType,
-    CompletionInputType,
-    StringType,
-)
 from llama_stack.core.datatypes import Api
+from llama_stack_api import ChatCompletionInputType, CompletionInputType, StringType
 
 
 class ColumnName(Enum):

llama_stack/providers/utils/files/form_data.py

@@ -9,7 +9,7 @@ import json
 from fastapi import Request
 from pydantic import BaseModel, ValidationError
 
-from llama_stack.apis.files import ExpiresAfter
+from llama_stack_api import ExpiresAfter
 
 
 async def parse_pydantic_from_form[T: BaseModel](request: Request, field_name: str, model_class: type[T]) -> T | None:

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -17,7 +17,7 @@ from llama_stack.log import get_logger
 if TYPE_CHECKING:
     from sentence_transformers import SentenceTransformer
 
-from llama_stack.apis.inference import (
+from llama_stack_api import (
     ModelStore,
     OpenAIEmbeddingData,
     OpenAIEmbeddingsRequestWithExtraBody,

llama_stack/providers/utils/inference/inference_store.py

@@ -8,20 +8,19 @@ from typing import Any
 
 from sqlalchemy.exc import IntegrityError
 
-from llama_stack.apis.inference import (
+from llama_stack.core.datatypes import AccessRule
+from llama_stack.core.storage.datatypes import InferenceStoreReference, StorageBackendType
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.core.storage.sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl
+from llama_stack.log import get_logger
+from llama_stack_api import (
     ListOpenAIChatCompletionResponse,
     OpenAIChatCompletion,
     OpenAICompletionWithInputMessages,
     OpenAIMessageParam,
     Order,
 )
-from llama_stack.core.datatypes import AccessRule
-from llama_stack.core.storage.datatypes import InferenceStoreReference, StorageBackendType
-from llama_stack.log import get_logger
-
-from ..sqlstore.api import ColumnDefinition, ColumnType
-from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore
-from ..sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
 
 logger = get_logger(name=__name__, category="inference")
 

llama_stack/providers/utils/inference/litellm_openai_mixin.py

@@ -7,13 +7,20 @@
 import base64
 import struct
 from collections.abc import AsyncIterator
+from typing import Any
 
 import litellm
 
-from llama_stack.apis.inference import (
-    ChatCompletionRequest,
+from llama_stack.core.request_headers import NeedsRequestProviderData
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper, ProviderModelEntry
+from llama_stack.providers.utils.inference.openai_compat import (
+    get_stream_options_for_telemetry,
+    prepare_openai_completion_params,
+)
+from llama_stack.providers.utils.inference.stream_utils import wrap_async_stream
+from llama_stack_api import (
     InferenceProvider,
-    JsonSchemaResponseFormat,
     OpenAIChatCompletion,
     OpenAIChatCompletionChunk,
     OpenAIChatCompletionRequestWithExtraBody,
@@ -23,16 +30,6 @@ from llama_stack.apis.inference import (
     OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
-    ToolChoice,
-)
-from llama_stack.core.request_headers import NeedsRequestProviderData
-from llama_stack.log import get_logger
-from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper, ProviderModelEntry
-from llama_stack.providers.utils.inference.openai_compat import (
-    convert_message_to_openai_dict_new,
-    convert_tooldef_to_openai_tool,
-    get_sampling_options,
-    prepare_openai_completion_params,
 )
 
 logger = get_logger(name=__name__, category="providers::utils")
@@ -55,6 +52,7 @@ class LiteLLMOpenAIMixin(
         openai_compat_api_base: str | None = None,
         download_images: bool = False,
         json_schema_strict: bool = True,
+        supports_stream_options: bool = True,
     ):
         """
         Initialize the LiteLLMOpenAIMixin.
@@ -66,6 +64,7 @@ class LiteLLMOpenAIMixin(
         :param openai_compat_api_base: The base URL for OpenAI compatibility, or None if not using OpenAI compatibility.
         :param download_images: Whether to download images and convert to base64 for message conversion.
         :param json_schema_strict: Whether to use strict mode for JSON schema validation.
+        :param supports_stream_options: Whether the provider supports stream_options parameter.
         """
         ModelRegistryHelper.__init__(self, model_entries=model_entries)
 
@@ -75,6 +74,7 @@ class LiteLLMOpenAIMixin(
         self.api_base = openai_compat_api_base
         self.download_images = download_images
         self.json_schema_strict = json_schema_strict
+        self.supports_stream_options = supports_stream_options
 
         if openai_compat_api_base:
             self.is_openai_compat = True
@@ -127,59 +127,13 @@ class LiteLLMOpenAIMixin(
 
         return schema
 
-    async def _get_params(self, request: ChatCompletionRequest) -> dict:
-        input_dict = {}
-
-        input_dict["messages"] = [
-            await convert_message_to_openai_dict_new(m, download_images=self.download_images) for m in request.messages
-        ]
-        if fmt := request.response_format:
-            if not isinstance(fmt, JsonSchemaResponseFormat):
-                raise ValueError(
-                    f"Unsupported response format: {type(fmt)}. Only JsonSchemaResponseFormat is supported."
-                )
-
-            fmt = fmt.json_schema
-            name = fmt["title"]
-            del fmt["title"]
-            fmt["additionalProperties"] = False
-
-            # Apply additionalProperties: False recursively to all objects
-            fmt = self._add_additional_properties_recursive(fmt)
-
-            input_dict["response_format"] = {
-                "type": "json_schema",
-                "json_schema": {
-                    "name": name,
-                    "schema": fmt,
-                    "strict": self.json_schema_strict,
-                },
-            }
-        if request.tools:
-            input_dict["tools"] = [convert_tooldef_to_openai_tool(tool) for tool in request.tools]
-            if request.tool_config.tool_choice:
-                input_dict["tool_choice"] = (
-                    request.tool_config.tool_choice.value
-                    if isinstance(request.tool_config.tool_choice, ToolChoice)
-                    else request.tool_config.tool_choice
-                )
-
-        return {
-            "model": request.model,
-            "api_key": self.get_api_key(),
-            "api_base": self.api_base,
-            **input_dict,
-            "stream": request.stream,
-            **get_sampling_options(request.sampling_params),
-        }
-
     def get_api_key(self) -> str:
         provider_data = self.get_request_provider_data()
         key_field = self.provider_data_api_key_field
-        if provider_data and getattr(provider_data, key_field, None):
-            api_key = getattr(provider_data, key_field)
-        else:
-            api_key = self.api_key_from_config
+        if provider_data and key_field and (api_key := getattr(provider_data, key_field, None)):
+            return str(api_key)  # type: ignore[no-any-return]  # getattr returns Any, can't narrow without runtime type inspection
+
+        api_key = self.api_key_from_config
         if not api_key:
             raise ValueError(
                 "API key is not set. Please provide a valid API key in the "
@@ -192,7 +146,13 @@ class LiteLLMOpenAIMixin(
         self,
         params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
+        if not self.model_store:
+            raise ValueError("Model store is not initialized")
+
         model_obj = await self.model_store.get_model(params.model)
+        if model_obj.provider_resource_id is None:
+            raise ValueError(f"Model {params.model} has no provider_resource_id")
+        provider_resource_id = model_obj.provider_resource_id
 
         # Convert input to list if it's a string
         input_list = [params.input] if isinstance(params.input, str) else params.input
@@ -200,7 +160,7 @@ class LiteLLMOpenAIMixin(
         # Call litellm embedding function
         # litellm.drop_params = True
         response = litellm.embedding(
-            model=self.get_litellm_model_name(model_obj.provider_resource_id),
+            model=self.get_litellm_model_name(provider_resource_id),
             input=input_list,
             api_key=self.get_api_key(),
             api_base=self.api_base,
@@ -217,18 +177,29 @@ class LiteLLMOpenAIMixin(
 
         return OpenAIEmbeddingsResponse(
             data=data,
-            model=model_obj.provider_resource_id,
+            model=provider_resource_id,
             usage=usage,
         )
 
     async def openai_completion(
         self,
         params: OpenAICompletionRequestWithExtraBody,
-    ) -> OpenAICompletion:
+    ) -> OpenAICompletion | AsyncIterator[OpenAICompletion]:
+        # Inject stream_options when streaming and telemetry is active
+        stream_options = get_stream_options_for_telemetry(
+            params.stream_options, params.stream, self.supports_stream_options
+        )
+
+        if not self.model_store:
+            raise ValueError("Model store is not initialized")
+
         model_obj = await self.model_store.get_model(params.model)
+        if model_obj.provider_resource_id is None:
+            raise ValueError(f"Model {params.model} has no provider_resource_id")
+        provider_resource_id = model_obj.provider_resource_id
 
         request_params = await prepare_openai_completion_params(
-            model=self.get_litellm_model_name(model_obj.provider_resource_id),
+            model=self.get_litellm_model_name(provider_resource_id),
             prompt=params.prompt,
             best_of=params.best_of,
             echo=params.echo,
@@ -241,34 +212,42 @@ class LiteLLMOpenAIMixin(
             seed=params.seed,
             stop=params.stop,
             stream=params.stream,
-            stream_options=params.stream_options,
+            stream_options=stream_options,
             temperature=params.temperature,
             top_p=params.top_p,
             user=params.user,
             suffix=params.suffix,
             api_key=self.get_api_key(),
             api_base=self.api_base,
+            **self._litellm_extra_request_params(params),
         )
-        return await litellm.atext_completion(**request_params)
+        # LiteLLM returns compatible type but mypy can't verify external library
+        result = await litellm.atext_completion(**request_params)
+
+        if params.stream:
+            return wrap_async_stream(result)  # type: ignore[arg-type]  # LiteLLM streaming types
+
+        return result  # type: ignore[return-value]  # external lib lacks type stubs
 
     async def openai_chat_completion(
         self,
         params: OpenAIChatCompletionRequestWithExtraBody,
     ) -> OpenAIChatCompletion | AsyncIterator[OpenAIChatCompletionChunk]:
-        # Add usage tracking for streaming when telemetry is active
-        from llama_stack.providers.utils.telemetry.tracing import get_current_span
+        # Inject stream_options when streaming and telemetry is active
+        stream_options = get_stream_options_for_telemetry(
+            params.stream_options, params.stream, self.supports_stream_options
+        )
 
-        stream_options = params.stream_options
-        if params.stream and get_current_span() is not None:
-            if stream_options is None:
-                stream_options = {"include_usage": True}
-            elif "include_usage" not in stream_options:
-                stream_options = {**stream_options, "include_usage": True}
+        if not self.model_store:
+            raise ValueError("Model store is not initialized")
 
         model_obj = await self.model_store.get_model(params.model)
+        if model_obj.provider_resource_id is None:
+            raise ValueError(f"Model {params.model} has no provider_resource_id")
+        provider_resource_id = model_obj.provider_resource_id
 
         request_params = await prepare_openai_completion_params(
-            model=self.get_litellm_model_name(model_obj.provider_resource_id),
+            model=self.get_litellm_model_name(provider_resource_id),
             messages=params.messages,
             frequency_penalty=params.frequency_penalty,
             function_call=params.function_call,
@@ -293,8 +272,15 @@ class LiteLLMOpenAIMixin(
             user=params.user,
             api_key=self.get_api_key(),
             api_base=self.api_base,
+            **self._litellm_extra_request_params(params),
         )
-        return await litellm.acompletion(**request_params)
+        # LiteLLM returns compatible type but mypy can't verify external library
+        result = await litellm.acompletion(**request_params)
+
+        if params.stream:
+            return wrap_async_stream(result)  # type: ignore[arg-type]  # LiteLLM streaming types
+
+        return result  # type: ignore[return-value]  # external lib lacks type stubs
 
     async def check_model_availability(self, model: str) -> bool:
         """
@@ -310,6 +296,20 @@ class LiteLLMOpenAIMixin(
 
         return model in litellm.models_by_provider[self.litellm_provider_name]
 
+    def _litellm_extra_request_params(
+        self,
+        params: OpenAIChatCompletionRequestWithExtraBody | OpenAICompletionRequestWithExtraBody,
+    ) -> dict[str, Any]:
+        """
+        Provider hook for extra LiteLLM/OpenAI-compat request params.
+
+        This is intentionally a narrow hook so provider adapters (e.g. WatsonX)
+        can add provider-specific kwargs (timeouts, project IDs, etc.) while the
+        mixin remains the single source of truth for telemetry-driven
+        stream_options injection.
+        """
+        return {}
+
 
 def b64_encode_openai_embeddings_response(
     response_data: list[dict], encoding_format: str | None = "float"

llama_stack/providers/utils/inference/model_registry.py

@@ -8,13 +8,11 @@ from typing import Any
 
 from pydantic import BaseModel, Field, SecretStr
 
-from llama_stack.apis.common.errors import UnsupportedModelError
-from llama_stack.apis.models import ModelType
 from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 from llama_stack.providers.utils.inference import (
     ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR,
 )
+from llama_stack_api import Model, ModelsProtocolPrivate, ModelType, UnsupportedModelError
 
 logger = get_logger(name=__name__, category="providers::utils")