Repository: meta-llama/llama-stack
File: llama_stack/providers/utils/inference/openai_mixin.py (Python)

Code before:
            "input": params.input,
            "encoding_format": params.encoding_format if params.encoding_format is not None else NOT_GIVEN,
            "dimensions": params.dimensions if params.dimensions is not None else NOT_GIVEN,
            "user": params.user if params.user is not None else NOT_GIVEN,
        }

        # Add extra_body if present
        extra_body = params.model_extra
        if extra_body:
            request_params["extra_body"] = extra_body

        # Call OpenAI embeddings API with properly typed parameters
        response = await self.client.embeddings.create(**request_params)

        data = []
        for i, embedding_data in enumerate(response.data):
            data.append(
                OpenAIEmbeddingData(
                    embedding=embedding_data.embedding,
                    index=i,
                )
            )

        # Handle providers that don't return usage statistics (e.g., Gemini)
        # Use default values of 0 when usage is not provided
        if response.usage:
            usage = OpenAIEmbeddingUsage(
                prompt_tokens=response.usage.prompt_tokens,
                total_tokens=response.usage.total_tokens,
            )
        else:
            usage = OpenAIEmbeddingUsage(
                prompt_tokens=0,
                total_tokens=0,
            )

        return OpenAIEmbeddingsResponse(
            data=data,
            model=params.model,
            usage=usage,
        )

    ###
    # ModelsProtocolPrivate implementation - provide model management functionality
    #
    #  async def register_model(self, model: Model) -> Model: ...
    #  async def unregister_model(self, model_id: str) -> None: ...
    #
    #  async def list_models(self) -> list[Model] | None: ...
    #  async def should_refresh_models(self) -> bool: ...
    ##

Review comment (suggestion):
```suggestion
        if hasattr(response, "usage"):
```

Code after:
            "input": params.input,
            "encoding_format": params.encoding_format if params.encoding_format is not None else NOT_GIVEN,
            "dimensions": params.dimensions if params.dimensions is not None else NOT_GIVEN,
            "user": params.user if params.user is not None else NOT_GIVEN,
        }

        # Add extra_body if present
        extra_body = params.model_extra
        if extra_body:
            request_params["extra_body"] = extra_body

        # Call OpenAI embeddings API with properly typed parameters
        response = await self.client.embeddings.create(**request_params)

        data = []
        for i, embedding_data in enumerate(response.data):
            data.append(
                OpenAIEmbeddingData(
                    embedding=embedding_data.embedding,
                    index=i,
                )
            )

        usage = OpenAIEmbeddingUsage(
            prompt_tokens=response.usage.prompt_tokens,
            total_tokens=response.usage.total_tokens,
        )

        return OpenAIEmbeddingsResponse(
            data=data,
            model=params.model,
            usage=usage,
        )

    ###
    # ModelsProtocolPrivate implementation - provide model management functionality
    #
    #  async def register_model(self, model: Model) -> Model: ...
    #  async def unregister_model(self, model_id: str) -> None: ...
    #
    #  async def list_models(self) -> list[Model] | None: ...
    #  async def should_refresh_models(self) -> bool: ...
    ##

    async def register_model(self, model: Model) -> Model:
        if not await self.check_model_availability(model.provider_model_id):
            raise ValueError(f"Model {model.provider_model_id} is not available from provider {self.__provider_id__}")  # type: ignore[attr-defined]
        return model

    async def unregister_model(self, model_id: str) -> None:
        return None