Merge branch 'main' into remove-deprecated-embeddings

llama_stack/providers/remote/inference/cerebras/cerebras.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 from collections.abc import AsyncGenerator
+from urllib.parse import urljoin
 
 from cerebras.cloud.sdk import AsyncCerebras
 
@@ -31,42 +32,41 @@ from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
 )
 from llama_stack.providers.utils.inference.openai_compat import (
-    OpenAIChatCompletionToLlamaStackMixin,
-    OpenAICompletionToLlamaStackMixin,
     get_sampling_options,
     process_chat_completion_response,
     process_chat_completion_stream_response,
     process_completion_response,
     process_completion_stream_response,
 )
+from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 from llama_stack.providers.utils.inference.prompt_adapter import (
     chat_completion_request_to_prompt,
     completion_request_to_prompt,
 )
 
 from .config import CerebrasImplConfig
-from .models import MODEL_ENTRIES
 
 
 class CerebrasInferenceAdapter(
+    OpenAIMixin,
     ModelRegistryHelper,
     Inference,
-    OpenAIChatCompletionToLlamaStackMixin,
-    OpenAICompletionToLlamaStackMixin,
 ):
     def __init__(self, config: CerebrasImplConfig) -> None:
-        ModelRegistryHelper.__init__(
-            self,
-            model_entries=MODEL_ENTRIES,
-        )
         self.config = config
 
         # TODO: make this use provider data, etc. like other providers
-        self.client = AsyncCerebras(
+        self._cerebras_client = AsyncCerebras(
             base_url=self.config.base_url,
             api_key=self.config.api_key.get_secret_value(),
         )
 
+    def get_api_key(self) -> str:
+        return self.config.api_key.get_secret_value()
+
+    def get_base_url(self) -> str:
+        return urljoin(self.config.base_url, "v1")
+
     async def initialize(self) -> None:
         return
 
@@ -103,14 +103,14 @@ class CerebrasInferenceAdapter(
     async def _nonstream_completion(self, request: CompletionRequest) -> CompletionResponse:
         params = await self._get_params(request)
 
-        r = await self.client.completions.create(**params)
+        r = await self._cerebras_client.completions.create(**params)
 
         return process_completion_response(r)
 
     async def _stream_completion(self, request: CompletionRequest) -> AsyncGenerator:
         params = await self._get_params(request)
 
-        stream = await self.client.completions.create(**params)
+        stream = await self._cerebras_client.completions.create(**params)
 
         async for chunk in process_completion_stream_response(stream):
             yield chunk
@@ -152,14 +152,14 @@ class CerebrasInferenceAdapter(
     async def _nonstream_chat_completion(self, request: CompletionRequest) -> CompletionResponse:
         params = await self._get_params(request)
 
-        r = await self.client.completions.create(**params)
+        r = await self._cerebras_client.completions.create(**params)
 
         return process_chat_completion_response(r, request)
 
     async def _stream_chat_completion(self, request: CompletionRequest) -> AsyncGenerator:
         params = await self._get_params(request)
 
-        stream = await self.client.completions.create(**params)
+        stream = await self._cerebras_client.completions.create(**params)
 
         async for chunk in process_chat_completion_stream_response(stream, request):
             yield chunk

llama_stack/providers/remote/inference/cerebras/config.py

@@ -20,8 +20,8 @@ class CerebrasImplConfig(BaseModel):
         default=os.environ.get("CEREBRAS_BASE_URL", DEFAULT_BASE_URL),
         description="Base URL for the Cerebras API",
     )
-    api_key: SecretStr | None = Field(
-        default=os.environ.get("CEREBRAS_API_KEY"),
+    api_key: SecretStr = Field(
+        default=SecretStr(os.environ.get("CEREBRAS_API_KEY")),
         description="Cerebras API Key",
     )
 

llama_stack/providers/remote/inference/cerebras/models.py

@@ -1,28 +0,0 @@
-# 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.
-
-from llama_stack.models.llama.sku_types import CoreModelId
-from llama_stack.providers.utils.inference.model_registry import (
-    build_hf_repo_model_entry,
-)
-
-SAFETY_MODELS_ENTRIES = []
-
-# https://inference-docs.cerebras.ai/models
-MODEL_ENTRIES = [
-    build_hf_repo_model_entry(
-        "llama3.1-8b",
-        CoreModelId.llama3_1_8b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "llama-3.3-70b",
-        CoreModelId.llama3_3_70b_instruct.value,
-    ),
-    build_hf_repo_model_entry(
-        "llama-4-scout-17b-16e-instruct",
-        CoreModelId.llama4_scout_17b_16e_instruct.value,
-    ),
-] + SAFETY_MODELS_ENTRIES