/
OS-Worldb968155# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
#
# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
# SPDX-License-Identifier: MIT
"""Create an OpenAI-compatible client using Cerebras's API.
Example:
```python
llm_config = {
"config_list": [{"api_type": "cerebras", "model": "llama3.1-8b", "api_key": os.environ.get("CEREBRAS_API_KEY")}]
}
agent = autogen.AssistantAgent("my_agent", llm_config=llm_config)
```
Install Cerebras's python library using: pip install --upgrade cerebras_cloud_sdk
Resources:
- https://inference-docs.cerebras.ai/quickstart
"""
from __future__ import annotations
import copy
import math
import os
import time
import warnings
from typing import Any, Literal, Optional
from pydantic import Field, ValidationInfo, field_validator
from ..import_utils import optional_import_block, require_optional_import
from ..llm_config import LLMConfigEntry, register_llm_config
from .client_utils import should_hide_tools, validate_parameter
from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
with optional_import_block():
from cerebras.cloud.sdk import Cerebras, Stream
CEREBRAS_PRICING_1K = {
# Convert pricing per million to per thousand tokens.
"llama3.1-8b": (0.10 / 1000, 0.10 / 1000),
"llama-3.3-70b": (0.85 / 1000, 1.20 / 1000),
}
@register_llm_config
class CerebrasLLMConfigEntry(LLMConfigEntry):
api_type: Literal["cerebras"] = "cerebras"
max_tokens: Optional[int] = None
seed: Optional[int] = None
stream: bool = False
temperature: float = Field(default=1.0, ge=0.0, le=1.5)
top_p: Optional[float] = None
hide_tools: Literal["if_all_run", "if_any_run", "never"] = "never"
tool_choice: Optional[Literal["none", "auto", "required"]] = None
@field_validator("top_p", mode="before")
@classmethod
def check_top_p(cls, v: Any, info: ValidationInfo) -> Any:
if v is not None and info.data.get("temperature") is not None:
raise ValueError("temperature and top_p cannot be set at the same time.")
return v
def create_client(self):
raise NotImplementedError("CerebrasLLMConfigEntry.create_client is not implemented.")
class CerebrasClient:
"""Client for Cerebras's API."""
def __init__(self, api_key=None, **kwargs):
"""Requires api_key or environment variable to be set
Args:
api_key (str): The API key for using Cerebras (or environment variable CEREBRAS_API_KEY needs to be set)
**kwargs: Additional keyword arguments to pass to the Cerebras client
"""
# Ensure we have the api_key upon instantiation
self.api_key = api_key
if not self.api_key:
self.api_key = os.getenv("CEREBRAS_API_KEY")
assert self.api_key, (
"Please include the api_key in your config list entry for Cerebras or set the CEREBRAS_API_KEY env variable."
)
if "response_format" in kwargs and kwargs["response_format"] is not None:
warnings.warn("response_format is not supported for Crebras, it will be ignored.", UserWarning)
def message_retrieval(self, response: ChatCompletion) -> list:
"""Retrieve and return a list of strings or a list of Choice.Message from the response.
NOTE: if a list of Choice.Message is returned, it currently needs to contain the fields of OpenAI's ChatCompletion Message object,
since that is expected for function or tool calling in the rest of the codebase at the moment, unless a custom agent is being used.
"""
return [choice.message for choice in response.choices]
def cost(self, response: ChatCompletion) -> float:
# Note: This field isn't explicitly in `ChatCompletion`, but is injected during chat creation.
return response.cost
@staticmethod
def get_usage(response: ChatCompletion) -> dict:
"""Return usage summary of the response using RESPONSE_USAGE_KEYS."""
# ... # pragma: no cover
return {
"prompt_tokens": response.usage.prompt_tokens,
"completion_tokens": response.usage.completion_tokens,
"total_tokens": response.usage.total_tokens,
"cost": response.cost,
"model": response.model,
}
def parse_params(self, params: dict[str, Any]) -> dict[str, Any]:
"""Loads the parameters for Cerebras API from the passed in parameters and returns a validated set. Checks types, ranges, and sets defaults"""
cerebras_params = {}
# Check that we have what we need to use Cerebras's API
# We won't enforce the available models as they are likely to change
cerebras_params["model"] = params.get("model")
assert cerebras_params["model"], (
"Please specify the 'model' in your config list entry to nominate the Cerebras model to use."
)
# Validate allowed Cerebras parameters
# https://inference-docs.cerebras.ai/api-reference/chat-completions
cerebras_params["max_tokens"] = validate_parameter(params, "max_tokens", int, True, None, (0, None), None)
cerebras_params["seed"] = validate_parameter(params, "seed", int, True, None, None, None)
cerebras_params["stream"] = validate_parameter(params, "stream", bool, True, False, None, None)
cerebras_params["temperature"] = validate_parameter(
params, "temperature", (int, float), True, 1, (0, 1.5), None
)
cerebras_params["top_p"] = validate_parameter(params, "top_p", (int, float), True, None, None, None)
cerebras_params["tool_choice"] = validate_parameter(
params, "tool_choice", str, True, None, None, ["none", "auto", "required"]
)
return cerebras_params
@require_optional_import("cerebras", "cerebras")
def create(self, params: dict) -> ChatCompletion:
messages = params.get("messages", [])
# Convert AG2 messages to Cerebras messages
cerebras_messages = oai_messages_to_cerebras_messages(messages)
# Parse parameters to the Cerebras API's parameters
cerebras_params = self.parse_params(params)
# Add tools to the call if we have them and aren't hiding them
if "tools" in params:
hide_tools = validate_parameter(
params, "hide_tools", str, False, "never", None, ["if_all_run", "if_any_run", "never"]
)
if not should_hide_tools(cerebras_messages, params["tools"], hide_tools):
cerebras_params["tools"] = params["tools"]
cerebras_params["messages"] = cerebras_messages
# We use chat model by default, and set max_retries to 5 (in line with typical retries loop)
client = Cerebras(api_key=self.api_key, max_retries=5)
# Token counts will be returned
prompt_tokens = 0
completion_tokens = 0
total_tokens = 0
# Streaming tool call recommendations
streaming_tool_calls = []
ans = None
response = client.chat.completions.create(**cerebras_params)
if cerebras_params["stream"]:
# Read in the chunks as they stream, taking in tool_calls which may be across
# multiple chunks if more than one suggested
ans = ""
for chunk in response:
# Grab first choice, which _should_ always be generated.
ans = ans + (getattr(chunk.choices[0].delta, "content", None) or "")
if "tool_calls" in chunk.choices[0].delta:
# We have a tool call recommendation
for tool_call in chunk.choices[0].delta["tool_calls"]:
streaming_tool_calls.append(
ChatCompletionMessageToolCall(
id=tool_call["id"],
function={
"name": tool_call["function"]["name"],
"arguments": tool_call["function"]["arguments"],
},
type="function",
)
)
if chunk.choices[0].finish_reason:
prompt_tokens = chunk.usage.prompt_tokens
completion_tokens = chunk.usage.completion_tokens
total_tokens = chunk.usage.total_tokens
else:
# Non-streaming finished
ans: str = response.choices[0].message.content
prompt_tokens = response.usage.prompt_tokens
completion_tokens = response.usage.completion_tokens
total_tokens = response.usage.total_tokens
if response is not None:
if isinstance(response, Stream):
# Streaming response
if chunk.choices[0].finish_reason == "tool_calls":
cerebras_finish = "tool_calls"
tool_calls = streaming_tool_calls
else:
cerebras_finish = "stop"
tool_calls = None
response_content = ans
response_id = chunk.id
else:
# Non-streaming response
# If we have tool calls as the response, populate completed tool calls for our return OAI response
if response.choices[0].finish_reason == "tool_calls":
cerebras_finish = "tool_calls"
tool_calls = []
for tool_call in response.choices[0].message.tool_calls:
tool_calls.append(
ChatCompletionMessageToolCall(
id=tool_call.id,
function={"name": tool_call.function.name, "arguments": tool_call.function.arguments},
type="function",
)
)
else:
cerebras_finish = "stop"
tool_calls = None
response_content = response.choices[0].message.content
response_id = response.id
# 3. convert output
message = ChatCompletionMessage(
role="assistant",
content=response_content,
function_call=None,
tool_calls=tool_calls,
)
choices = [Choice(finish_reason=cerebras_finish, index=0, message=message)]
response_oai = ChatCompletion(
id=response_id,
model=cerebras_params["model"],
created=int(time.time()),
object="chat.completion",
choices=choices,
usage=CompletionUsage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=total_tokens,
),
# Note: This seems to be a field that isn't in the schema of `ChatCompletion`, so Pydantic
# just adds it dynamically.
cost=calculate_cerebras_cost(prompt_tokens, completion_tokens, cerebras_params["model"]),
)
return response_oai
def oai_messages_to_cerebras_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Convert messages from OAI format to Cerebras's format.
We correct for any specific role orders and types.
"""
cerebras_messages = copy.deepcopy(messages)
# Remove the name field
for message in cerebras_messages:
if "name" in message:
message.pop("name", None)
return cerebras_messages
def calculate_cerebras_cost(input_tokens: int, output_tokens: int, model: str) -> float:
"""Calculate the cost of the completion using the Cerebras pricing."""
total = 0.0
if model in CEREBRAS_PRICING_1K:
input_cost_per_k, output_cost_per_k = CEREBRAS_PRICING_1K[model]
input_cost = math.ceil((input_tokens / 1000) * input_cost_per_k * 1e6) / 1e6
output_cost = math.ceil((output_tokens / 1000) * output_cost_per_k * 1e6) / 1e6
total = math.ceil((input_cost + output_cost) * 1e6) / 1e6
else:
warnings.warn(f"Cost calculation not available for model {model}", UserWarning)
return total