Merge remote-tracking branch 'origin/main' into TamiTakamiya/tool-param-definition-update

llama_stack/providers/utils/inference/inference_store.py

@@ -54,7 +54,7 @@ class InferenceStore:
 
     async def initialize(self):
         """Create the necessary tables if they don't exist."""
-        self.sql_store = AuthorizedSqlStore(sqlstore_impl(self.sql_store_config))
+        self.sql_store = AuthorizedSqlStore(sqlstore_impl(self.sql_store_config), self.policy)
         await self.sql_store.create_table(
             "chat_completions",
             {
@@ -202,7 +202,6 @@ class InferenceStore:
             order_by=[("created", order.value)],
             cursor=("id", after) if after else None,
             limit=limit,
-            policy=self.policy,
         )
 
         data = [
@@ -229,7 +228,6 @@ class InferenceStore:
         row = await self.sql_store.fetch_one(
             table="chat_completions",
             where={"id": completion_id},
-            policy=self.policy,
         )
 
         if not row:

llama_stack/providers/utils/inference/litellm_openai_mixin.py

@@ -40,7 +40,7 @@ from llama_stack.apis.inference import (
 )
 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
+from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper, ProviderModelEntry
 from llama_stack.providers.utils.inference.openai_compat import (
     b64_encode_openai_embeddings_response,
     convert_message_to_openai_dict_new,
@@ -67,10 +67,10 @@ class LiteLLMOpenAIMixin(
     #                         when calling litellm.
     def __init__(
         self,
-        model_entries,
         litellm_provider_name: str,
         api_key_from_config: str | None,
         provider_data_api_key_field: str,
+        model_entries: list[ProviderModelEntry] | None = None,
         openai_compat_api_base: str | None = None,
         download_images: bool = False,
         json_schema_strict: bool = True,
@@ -86,7 +86,7 @@ class LiteLLMOpenAIMixin(
         :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)
+        ModelRegistryHelper.__init__(self, model_entries=model_entries)
 
         self.litellm_provider_name = litellm_provider_name
         self.api_key_from_config = api_key_from_config

llama_stack/providers/utils/inference/model_registry.py

@@ -11,7 +11,6 @@ from pydantic import BaseModel, Field
 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.models.llama.sku_list import all_registered_models
 from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
 from llama_stack.providers.utils.inference import (
     ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR,
@@ -21,7 +20,7 @@ logger = get_logger(name=__name__, category="providers::utils")
 
 
 class RemoteInferenceProviderConfig(BaseModel):
-    allowed_models: list[str] | None = Field(
+    allowed_models: list[str] | None = Field(  # TODO: make this non-optional and give a list() default
         default=None,
         description="List of models that should be registered with the model registry. If None, all models are allowed.",
     )
@@ -37,13 +36,6 @@ class ProviderModelEntry(BaseModel):
     metadata: dict[str, Any] = Field(default_factory=dict)
 
 
-def get_huggingface_repo(model_descriptor: str) -> str | None:
-    for model in all_registered_models():
-        if model.descriptor() == model_descriptor:
-            return model.huggingface_repo
-    return None
-
-
 def build_hf_repo_model_entry(
     provider_model_id: str,
     model_descriptor: str,
@@ -63,25 +55,20 @@ def build_hf_repo_model_entry(
     )
 
 
-def build_model_entry(provider_model_id: str, model_descriptor: str) -> ProviderModelEntry:
-    return ProviderModelEntry(
-        provider_model_id=provider_model_id,
-        aliases=[],
-        llama_model=model_descriptor,
-        model_type=ModelType.llm,
-    )
-
-
 class ModelRegistryHelper(ModelsProtocolPrivate):
     __provider_id__: str
 
-    def __init__(self, model_entries: list[ProviderModelEntry], allowed_models: list[str] | None = None):
-        self.model_entries = model_entries
+    def __init__(
+        self,
+        model_entries: list[ProviderModelEntry] | None = None,
+        allowed_models: list[str] | None = None,
+    ):
         self.allowed_models = allowed_models
 
         self.alias_to_provider_id_map = {}
         self.provider_id_to_llama_model_map = {}
-        for entry in model_entries:
+        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
 
@@ -103,7 +90,7 @@ class ModelRegistryHelper(ModelsProtocolPrivate):
                     Model(
                         identifier=id,
                         provider_resource_id=entry.provider_model_id,
-                        model_type=ModelType.llm,
+                        model_type=entry.model_type,
                         metadata=entry.metadata,
                         provider_id=self.__provider_id__,
                     )

llama_stack/providers/utils/inference/openai_mixin.py

@@ -4,12 +4,12 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import base64
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
 from typing import Any
 
-import openai
 from openai import NOT_GIVEN, AsyncOpenAI
 
 from llama_stack.apis.inference import (
@@ -23,13 +23,16 @@ from llama_stack.apis.inference import (
     OpenAIMessageParam,
     OpenAIResponseFormatParam,
 )
+from llama_stack.apis.models import ModelType
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
 from llama_stack.providers.utils.inference.openai_compat import prepare_openai_completion_params
+from llama_stack.providers.utils.inference.prompt_adapter import localize_image_content
 
 logger = get_logger(name=__name__, category="providers::utils")
 
 
-class OpenAIMixin(ABC):
+class OpenAIMixin(ModelRegistryHelper, ABC):
     """
     Mixin class that provides OpenAI-specific functionality for inference providers.
     This class handles direct OpenAI API calls using the AsyncOpenAI client.
@@ -50,6 +53,22 @@ class OpenAIMixin(ABC):
     # This is useful for providers that do not return a unique id in the response.
     overwrite_completion_id: bool = False
 
+    # Allow subclasses to control whether to download images and convert to base64
+    # for providers that require base64 encoded images instead of URLs.
+    download_images: bool = False
+
+    # Embedding model metadata for this provider
+    # Can be set by subclasses or instances to provide embedding models
+    # Format: {"model_id": {"embedding_dimension": 1536, "context_length": 8192}}
+    embedding_model_metadata: dict[str, dict[str, int]] = {}
+
+    # Cache of available models keyed by model ID
+    # This is set in list_models() and used in check_model_availability()
+    _model_cache: dict[str, Model] = {}
+
+    # List of allowed models for this provider, if empty all models allowed
+    allowed_models: list[str] = []
+
     @abstractmethod
     def get_api_key(self) -> str:
         """
@@ -226,6 +245,24 @@ class OpenAIMixin(ABC):
         """
         Direct OpenAI chat completion API call.
         """
+        if self.download_images:
+
+            async def _localize_image_url(m: OpenAIMessageParam) -> OpenAIMessageParam:
+                if isinstance(m.content, list):
+                    for c in m.content:
+                        if c.type == "image_url" and c.image_url and c.image_url.url and "http" in c.image_url.url:
+                            localize_result = await localize_image_content(c.image_url.url)
+                            if localize_result is None:
+                                raise ValueError(
+                                    f"Failed to localize image content from {c.image_url.url[:42]}{'...' if len(c.image_url.url) > 42 else ''}"
+                                )
+                            content, format = localize_result
+                            c.image_url.url = f"data:image/{format};base64,{base64.b64encode(content).decode('utf-8')}"
+                # else it's a string and we don't need to modify it
+                return m
+
+            messages = [await _localize_image_url(m) for m in messages]
+
         resp = await self.client.chat.completions.create(
             **await prepare_openai_completion_params(
                 model=await self._get_provider_model_id(model),
@@ -292,26 +329,53 @@ class OpenAIMixin(ABC):
 
         return OpenAIEmbeddingsResponse(
             data=data,
-            model=response.model,
+            model=model,
             usage=usage,
         )
 
+    async def list_models(self) -> list[Model] | None:
+        """
+        List available models from the provider's /v1/models endpoint augmented with static embedding model metadata.
+
+        Also, caches the models in self._model_cache for use in check_model_availability().
+
+        :return: A list of Model instances representing available models.
+        """
+        self._model_cache = {}
+
+        async for m in self.client.models.list():
+            if self.allowed_models and m.id not in self.allowed_models:
+                logger.info(f"Skipping model {m.id} as it is not in the allowed models list")
+                continue
+            if metadata := self.embedding_model_metadata.get(m.id):
+                # This is an embedding model - augment with metadata
+                model = Model(
+                    provider_id=self.__provider_id__,  # type: ignore[attr-defined]
+                    provider_resource_id=m.id,
+                    identifier=m.id,
+                    model_type=ModelType.embedding,
+                    metadata=metadata,
+                )
+            else:
+                # This is an LLM
+                model = Model(
+                    provider_id=self.__provider_id__,  # type: ignore[attr-defined]
+                    provider_resource_id=m.id,
+                    identifier=m.id,
+                    model_type=ModelType.llm,
+                )
+            self._model_cache[m.id] = model
+
+        return list(self._model_cache.values())
+
     async def check_model_availability(self, model: str) -> bool:
         """
-        Check if a specific model is available from OpenAI.
+        Check if a specific model is available from the provider's /v1/models.
 
         :param model: The model identifier to check.
         :return: True if the model is available dynamically, False otherwise.
         """
-        try:
-            # Direct model lookup - returns model or raises NotFoundError
-            await self.client.models.retrieve(model)
-            return True
-        except openai.NotFoundError:
-            # Model doesn't exist - this is expected for unavailable models
-            pass
-        except Exception as e:
-            # All other errors (auth, rate limit, network, etc.)
-            logger.warning(f"Failed to check model availability for {model}: {e}")
+        if not self._model_cache:
+            await self.list_models()
 
-        return False
+        return model in self._model_cache