llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/utils/inference/litellm_openai_mixin.py

@@ -0,0 +1,336 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import base64
+import struct
+from collections.abc import AsyncIterator
+
+import litellm
+
+from llama_stack.apis.inference import (
+    ChatCompletionRequest,
+    InferenceProvider,
+    JsonSchemaResponseFormat,
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAIChatCompletionRequestWithExtraBody,
+    OpenAICompletion,
+    OpenAICompletionRequestWithExtraBody,
+    OpenAIEmbeddingData,
+    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")
+
+
+class LiteLLMOpenAIMixin(
+    ModelRegistryHelper,
+    InferenceProvider,
+    NeedsRequestProviderData,
+):
+    # TODO: avoid exposing the litellm specific model names to the user.
+    #       potential change: add a prefix param that gets added to the model name
+    #                         when calling litellm.
+    def __init__(
+        self,
+        litellm_provider_name: str,
+        api_key_from_config: str | None,
+        provider_data_api_key_field: str | None = None,
+        model_entries: list[ProviderModelEntry] | None = None,
+        openai_compat_api_base: str | None = None,
+        download_images: bool = False,
+        json_schema_strict: bool = True,
+    ):
+        """
+        Initialize the LiteLLMOpenAIMixin.
+
+        :param model_entries: The model entries to register.
+        :param api_key_from_config: The API key to use from the config.
+        :param provider_data_api_key_field: The field in the provider data that contains the API key (optional).
+        :param litellm_provider_name: The name of the provider, used for model lookups.
+        :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.
+        """
+        ModelRegistryHelper.__init__(self, model_entries=model_entries)
+
+        self.litellm_provider_name = litellm_provider_name
+        self.api_key_from_config = api_key_from_config
+        self.provider_data_api_key_field = provider_data_api_key_field
+        self.api_base = openai_compat_api_base
+        self.download_images = download_images
+        self.json_schema_strict = json_schema_strict
+
+        if openai_compat_api_base:
+            self.is_openai_compat = True
+        else:
+            self.is_openai_compat = False
+
+    async def initialize(self):
+        pass
+
+    async def shutdown(self):
+        pass
+
+    def get_litellm_model_name(self, model_id: str) -> str:
+        # users may be using openai/ prefix in their model names. the openai/models.py did this by default.
+        # model_id.startswith("openai/") is for backwards compatibility.
+        return (
+            f"{self.litellm_provider_name}/{model_id}"
+            if self.is_openai_compat and not model_id.startswith(self.litellm_provider_name)
+            else model_id
+        )
+
+    def _add_additional_properties_recursive(self, schema):
+        """
+        Recursively add additionalProperties: False to all object schemas
+        """
+        if isinstance(schema, dict):
+            if schema.get("type") == "object":
+                schema["additionalProperties"] = False
+
+                # Add required field with all property keys if properties exist
+                if "properties" in schema and schema["properties"]:
+                    schema["required"] = list(schema["properties"].keys())
+
+            if "properties" in schema:
+                for prop_schema in schema["properties"].values():
+                    self._add_additional_properties_recursive(prop_schema)
+
+            for key in ["anyOf", "allOf", "oneOf"]:
+                if key in schema:
+                    for sub_schema in schema[key]:
+                        self._add_additional_properties_recursive(sub_schema)
+
+            if "not" in schema:
+                self._add_additional_properties_recursive(schema["not"])
+
+            # Handle $defs/$ref
+            if "$defs" in schema:
+                for def_schema in schema["$defs"].values():
+                    self._add_additional_properties_recursive(def_schema)
+
+        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 not api_key:
+            raise ValueError(
+                "API key is not set. Please provide a valid API key in the "
+                "provider data header, e.g. x-llamastack-provider-data: "
+                f'{{"{key_field}": "<API_KEY>"}}, or in the provider config.'
+            )
+        return api_key
+
+    async def openai_embeddings(
+        self,
+        params: OpenAIEmbeddingsRequestWithExtraBody,
+    ) -> OpenAIEmbeddingsResponse:
+        model_obj = await self.model_store.get_model(params.model)
+
+        # Convert input to list if it's a string
+        input_list = [params.input] if isinstance(params.input, str) else params.input
+
+        # Call litellm embedding function
+        # litellm.drop_params = True
+        response = litellm.embedding(
+            model=self.get_litellm_model_name(model_obj.provider_resource_id),
+            input=input_list,
+            api_key=self.get_api_key(),
+            api_base=self.api_base,
+            dimensions=params.dimensions,
+        )
+
+        # Convert response to OpenAI format
+        data = b64_encode_openai_embeddings_response(response.data, params.encoding_format)
+
+        usage = OpenAIEmbeddingUsage(
+            prompt_tokens=response["usage"]["prompt_tokens"],
+            total_tokens=response["usage"]["total_tokens"],
+        )
+
+        return OpenAIEmbeddingsResponse(
+            data=data,
+            model=model_obj.provider_resource_id,
+            usage=usage,
+        )
+
+    async def openai_completion(
+        self,
+        params: OpenAICompletionRequestWithExtraBody,
+    ) -> OpenAICompletion:
+        model_obj = await self.model_store.get_model(params.model)
+
+        request_params = await prepare_openai_completion_params(
+            model=self.get_litellm_model_name(model_obj.provider_resource_id),
+            prompt=params.prompt,
+            best_of=params.best_of,
+            echo=params.echo,
+            frequency_penalty=params.frequency_penalty,
+            logit_bias=params.logit_bias,
+            logprobs=params.logprobs,
+            max_tokens=params.max_tokens,
+            n=params.n,
+            presence_penalty=params.presence_penalty,
+            seed=params.seed,
+            stop=params.stop,
+            stream=params.stream,
+            stream_options=params.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,
+        )
+        return await litellm.atext_completion(**request_params)
+
+    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
+
+        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}
+
+        model_obj = await self.model_store.get_model(params.model)
+
+        request_params = await prepare_openai_completion_params(
+            model=self.get_litellm_model_name(model_obj.provider_resource_id),
+            messages=params.messages,
+            frequency_penalty=params.frequency_penalty,
+            function_call=params.function_call,
+            functions=params.functions,
+            logit_bias=params.logit_bias,
+            logprobs=params.logprobs,
+            max_completion_tokens=params.max_completion_tokens,
+            max_tokens=params.max_tokens,
+            n=params.n,
+            parallel_tool_calls=params.parallel_tool_calls,
+            presence_penalty=params.presence_penalty,
+            response_format=params.response_format,
+            seed=params.seed,
+            stop=params.stop,
+            stream=params.stream,
+            stream_options=stream_options,
+            temperature=params.temperature,
+            tool_choice=params.tool_choice,
+            tools=params.tools,
+            top_logprobs=params.top_logprobs,
+            top_p=params.top_p,
+            user=params.user,
+            api_key=self.get_api_key(),
+            api_base=self.api_base,
+        )
+        return await litellm.acompletion(**request_params)
+
+    async def check_model_availability(self, model: str) -> bool:
+        """
+        Check if a specific model is available via LiteLLM for the current
+        provider (self.litellm_provider_name).
+
+        :param model: The model identifier to check.
+        :return: True if the model is available dynamically, False otherwise.
+        """
+        if self.litellm_provider_name not in litellm.models_by_provider:
+            logger.error(f"Provider {self.litellm_provider_name} is not registered in litellm.")
+            return False
+
+        return model in litellm.models_by_provider[self.litellm_provider_name]
+
+
+def b64_encode_openai_embeddings_response(
+    response_data: list[dict], encoding_format: str | None = "float"
+) -> list[OpenAIEmbeddingData]:
+    """
+    Process the OpenAI embeddings response to encode the embeddings in base64 format if specified.
+    """
+    data = []
+    for i, embedding_data in enumerate(response_data):
+        if encoding_format == "base64":
+            byte_array = bytearray()
+            for embedding_value in embedding_data["embedding"]:
+                byte_array.extend(struct.pack("f", float(embedding_value)))
+
+            response_embedding = base64.b64encode(byte_array).decode("utf-8")
+        else:
+            response_embedding = embedding_data["embedding"]
+        data.append(
+            OpenAIEmbeddingData(
+                embedding=response_embedding,
+                index=i,
+            )
+        )
+    return data

llama_stack/providers/utils/inference/model_registry.py

@@ -4,38 +4,188 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Dict, List
+from typing import Any
 
-from llama_models.sku_list import resolve_model
+from pydantic import BaseModel, Field, SecretStr
 
-from llama_stack.providers.datatypes import ModelDef, ModelsProtocolPrivate
+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,
+)
 
+logger = get_logger(name=__name__, category="providers::utils")
+
+
+class RemoteInferenceProviderConfig(BaseModel):
+    allowed_models: list[str] | None = Field(
+        default=None,
+        description="List of models that should be registered with the model registry. If None, all models are allowed.",
+    )
+    refresh_models: bool = Field(
+        default=False,
+        description="Whether to refresh models periodically from the provider",
+    )
+    auth_credential: SecretStr | None = Field(
+        default=None,
+        description="Authentication credential for the provider",
+        alias="api_key",
+    )
 
-class ModelRegistryHelper(ModelsProtocolPrivate):
 
-    def __init__(self, stack_to_provider_models_map: Dict[str, str]):
-        self.stack_to_provider_models_map = stack_to_provider_models_map
+# TODO: this class is more confusing than useful right now. We need to make it
+# more closer to the Model class.
+class ProviderModelEntry(BaseModel):
+    provider_model_id: str
+    aliases: list[str] = Field(default_factory=list)
+    llama_model: str | None = None
+    model_type: ModelType = ModelType.llm
+    metadata: dict[str, Any] = Field(default_factory=dict)
 
-    def map_to_provider_model(self, identifier: str) -> str:
-        model = resolve_model(identifier)
-        if not model:
-            raise ValueError(f"Unknown model: `{identifier}`")
 
-        if identifier not in self.stack_to_provider_models_map:
-            raise ValueError(
-                f"Model {identifier} not found in map {self.stack_to_provider_models_map}"
-            )
+def build_hf_repo_model_entry(
+    provider_model_id: str,
+    model_descriptor: str,
+    additional_aliases: list[str] | None = None,
+) -> ProviderModelEntry:
+    aliases = [
+        # NOTE: avoid HF aliases because they _cannot_ be unique across providers
+        # get_huggingface_repo(model_descriptor),
+    ]
+    if additional_aliases:
+        aliases.extend(additional_aliases)
+    aliases = [alias for alias in aliases if alias is not None]
+    return ProviderModelEntry(
+        provider_model_id=provider_model_id,
+        aliases=aliases,
+        llama_model=model_descriptor,
+    )
+
+
+class ModelRegistryHelper(ModelsProtocolPrivate):
+    __provider_id__: str
 
-        return self.stack_to_provider_models_map[identifier]
+    def __init__(
+        self,
+        model_entries: list[ProviderModelEntry] | None = None,
+        allowed_models: list[str] | None = None,
+    ):
+        self.allowed_models = allowed_models if allowed_models else []
 
-    async def register_model(self, model: ModelDef) -> None:
-        if model.identifier not in self.stack_to_provider_models_map:
-            raise ValueError(
-                f"Unsupported model {model.identifier}. Supported models: {self.stack_to_provider_models_map.keys()}"
-            )
+        self.alias_to_provider_id_map = {}
+        self.provider_id_to_llama_model_map = {}
+        self.model_entries = model_entries or []
+        for entry in self.model_entries:
+            for alias in entry.aliases:
+                self.alias_to_provider_id_map[alias] = entry.provider_model_id
 
-    async def list_models(self) -> List[ModelDef]:
+            # also add a mapping from provider model id to itself for easy lookup
+            self.alias_to_provider_id_map[entry.provider_model_id] = entry.provider_model_id
+
+            if entry.llama_model:
+                self.alias_to_provider_id_map[entry.llama_model] = entry.provider_model_id
+                self.provider_id_to_llama_model_map[entry.provider_model_id] = entry.llama_model
+
+    async def list_models(self) -> list[Model] | None:
         models = []
-        for llama_model, provider_model in self.stack_to_provider_models_map.items():
-            models.append(ModelDef(identifier=llama_model, llama_model=llama_model))
+        for entry in self.model_entries:
+            ids = [entry.provider_model_id] + entry.aliases
+            for id in ids:
+                if self.allowed_models and id not in self.allowed_models:
+                    continue
+                models.append(
+                    Model(
+                        identifier=id,
+                        provider_resource_id=entry.provider_model_id,
+                        model_type=entry.model_type,
+                        metadata=entry.metadata,
+                        provider_id=self.__provider_id__,
+                    )
+                )
         return models
+
+    async def should_refresh_models(self) -> bool:
+        return False
+
+    def get_provider_model_id(self, identifier: str) -> str | None:
+        return self.alias_to_provider_id_map.get(identifier, None)
+
+    # TODO: why keep a separate llama model mapping?
+    def get_llama_model(self, provider_model_id: str) -> str | None:
+        return self.provider_id_to_llama_model_map.get(provider_model_id, None)
+
+    async def check_model_availability(self, model: str) -> bool:
+        """
+        Check if a specific model is available from the provider (non-static check).
+
+        This is for subclassing purposes, so providers can check if a specific
+        model is currently available for use through dynamic means (e.g., API calls).
+
+        This method should NOT check statically configured model entries in
+        `self.alias_to_provider_id_map` - that is handled separately in register_model.
+
+        Default implementation returns False (no dynamic models available).
+
+        :param model: The model identifier to check.
+        :return: True if the model is available dynamically, False otherwise.
+        """
+        logger.info(
+            f"check_model_availability is not implemented for {self.__class__.__name__}. Returning False by default."
+        )
+        return False
+
+    async def register_model(self, model: Model) -> Model:
+        # Check if model is supported in static configuration
+        supported_model_id = self.get_provider_model_id(model.provider_resource_id)
+
+        # If not found in static config, check if it's available dynamically from provider
+        if not supported_model_id:
+            if await self.check_model_availability(model.provider_resource_id):
+                supported_model_id = model.provider_resource_id
+            else:
+                # note: we cannot provide a complete list of supported models without
+                #       getting a complete list from the provider, so we return "..."
+                all_supported_models = [*self.alias_to_provider_id_map.keys(), "..."]
+                raise UnsupportedModelError(model.provider_resource_id, all_supported_models)
+
+        provider_resource_id = self.get_provider_model_id(model.model_id)
+        if model.model_type == ModelType.embedding:
+            # embedding models are always registered by their provider model id and does not need to be mapped to a llama model
+            provider_resource_id = model.provider_resource_id
+        if provider_resource_id:
+            if provider_resource_id != supported_model_id:  # be idempotent, only reject differences
+                raise ValueError(
+                    f"Model id '{model.model_id}' is already registered. Please use a different id or unregister it first."
+                )
+        else:
+            llama_model = model.metadata.get("llama_model")
+            if llama_model:
+                existing_llama_model = self.get_llama_model(model.provider_resource_id)
+                if existing_llama_model:
+                    if existing_llama_model != llama_model:
+                        raise ValueError(
+                            f"Provider model id '{model.provider_resource_id}' is already registered to a different llama model: '{existing_llama_model}'"
+                        )
+                else:
+                    if llama_model not in ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR:
+                        raise ValueError(
+                            f"Invalid llama_model '{llama_model}' specified in metadata. "
+                            f"Must be one of: {', '.join(ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR.keys())}"
+                        )
+                    self.provider_id_to_llama_model_map[model.provider_resource_id] = (
+                        ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR[llama_model]
+                    )
+
+        # Register the model alias, ensuring it maps to the correct provider model id
+        self.alias_to_provider_id_map[model.model_id] = supported_model_id
+
+        return model
+
+    async def unregister_model(self, model_id: str) -> None:
+        # model_id is the identifier, not the provider_resource_id
+        # unfortunately, this ID can be of the form provider_id/model_id which
+        # we never registered. TODO: fix this by significantly rewriting
+        # registration and registry helper
+        pass