agno 0.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/models/cerebras/__init__.py

@@ -0,0 +1,12 @@
+from agno.models.cerebras.cerebras import Cerebras
+
+try:
+    from agno.models.cerebras.cerebras_openai import CerebrasOpenAI
+except ImportError:
+
+    class CerebrasOpenAI:  # type: ignore
+        def __init__(self, *args, **kwargs):
+            raise ImportError("`openai` not installed. Please install it via `pip install openai`")
+
+
+__all__ = ["Cerebras", "CerebrasOpenAI"]

agno/models/cerebras/cerebras.py

@@ -0,0 +1,565 @@
+import json
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from os import getenv
+from typing import Any, Dict, Iterator, List, Optional, Type, Union
+
+import httpx
+from pydantic import BaseModel
+
+from agno.models.base import Model
+from agno.models.message import Message
+from agno.models.metrics import Metrics
+from agno.models.response import ModelResponse
+from agno.run.agent import RunOutput
+from agno.utils.http import get_default_async_client, get_default_sync_client
+from agno.utils.log import log_debug, log_error, log_warning
+
+try:
+    from cerebras.cloud.sdk import AsyncCerebras as AsyncCerebrasClient
+    from cerebras.cloud.sdk import Cerebras as CerebrasClient
+    from cerebras.cloud.sdk.types.chat.chat_completion import (
+        ChatChunkResponse,
+        ChatChunkResponseChoice,
+        ChatChunkResponseChoiceDelta,
+        ChatChunkResponseUsage,
+        ChatCompletionResponse,
+        ChatCompletionResponseChoice,
+        ChatCompletionResponseChoiceMessage,
+        ChatCompletionResponseUsage,
+    )
+except (ImportError, ModuleNotFoundError):
+    raise ImportError("`cerebras-cloud-sdk` not installed. Please install using `pip install cerebras-cloud-sdk`")
+
+
+@dataclass
+class Cerebras(Model):
+    """
+    A class for interacting with models using the Cerebras API.
+    """
+
+    id: str = "llama-4-scout-17b-16e-instruct"
+    name: str = "Cerebras"
+    provider: str = "Cerebras"
+
+    supports_native_structured_outputs: bool = False
+    supports_json_schema_outputs: bool = True
+
+    # Request parameters
+    parallel_tool_calls: Optional[bool] = None
+    max_completion_tokens: Optional[int] = None
+    repetition_penalty: Optional[float] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    top_k: Optional[int] = None
+    strict_output: bool = True  # When True, guarantees schema adherence for structured outputs. When False, attempts to follow schema as a guide but may occasionally deviate
+    extra_headers: Optional[Any] = None
+    extra_query: Optional[Any] = None
+    extra_body: Optional[Any] = None
+    request_params: Optional[Dict[str, Any]] = None
+
+    # Client parameters
+    api_key: Optional[str] = None
+    base_url: Optional[Union[str, httpx.URL]] = None
+    timeout: Optional[float] = None
+    max_retries: Optional[int] = None
+    default_headers: Optional[Any] = None
+    default_query: Optional[Any] = None
+    http_client: Optional[Union[httpx.Client, httpx.AsyncClient]] = None
+    client_params: Optional[Dict[str, Any]] = None
+
+    # Cerebras clients
+    client: Optional[CerebrasClient] = None
+    async_client: Optional[AsyncCerebrasClient] = None
+
+    def _get_client_params(self) -> Dict[str, Any]:
+        # Fetch API key from env if not already set
+        if not self.api_key:
+            self.api_key = getenv("CEREBRAS_API_KEY")
+            if not self.api_key:
+                log_error("CEREBRAS_API_KEY not set. Please set the CEREBRAS_API_KEY environment variable.")
+
+        # Define base client params
+        base_params = {
+            "api_key": self.api_key,
+            "base_url": self.base_url,
+            "timeout": self.timeout,
+            "max_retries": self.max_retries,
+            "default_headers": self.default_headers,
+            "default_query": self.default_query,
+        }
+
+        # Create client_params dict with non-None values
+        client_params = {k: v for k, v in base_params.items() if v is not None}
+
+        # Add additional client params if provided
+        if self.client_params:
+            client_params.update(self.client_params)
+        return client_params
+
+    def get_client(self) -> CerebrasClient:
+        """
+        Returns a Cerebras client.
+
+        Returns:
+            CerebrasClient: An instance of the Cerebras client.
+        """
+        if self.client and not self.client.is_closed():
+            return self.client
+
+        client_params: Dict[str, Any] = self._get_client_params()
+        if self.http_client is not None:
+            client_params["http_client"] = self.http_client
+        else:
+            # Use global sync client when no custom http_client is provided
+            client_params["http_client"] = get_default_sync_client()
+        self.client = CerebrasClient(**client_params)
+        return self.client
+
+    def get_async_client(self) -> AsyncCerebrasClient:
+        """
+        Returns an asynchronous Cerebras client.
+
+        Returns:
+            AsyncCerebras: An instance of the asynchronous Cerebras client.
+        """
+        if self.async_client and not self.async_client.is_closed():
+            return self.async_client
+
+        client_params: Dict[str, Any] = self._get_client_params()
+        if self.http_client and isinstance(self.http_client, httpx.AsyncClient):
+            client_params["http_client"] = self.http_client
+        else:
+            # Use global async client when no custom http_client is provided
+            client_params["http_client"] = get_default_async_client()
+        self.async_client = AsyncCerebrasClient(**client_params)
+        return self.async_client
+
+    def get_request_params(
+        self,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Returns keyword arguments for API requests.
+
+        Returns:
+            Dict[str, Any]: A dictionary of keyword arguments for API requests.
+        """
+        # Define base request parameters
+        base_params = {
+            "max_completion_tokens": self.max_completion_tokens,
+            "repetition_penalty": self.repetition_penalty,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+            "extra_headers": self.extra_headers,
+            "extra_query": self.extra_query,
+            "extra_body": self.extra_body,
+            "request_params": self.request_params,
+        }
+
+        # Filter out None values
+        request_params: Dict[str, Any] = {k: v for k, v in base_params.items() if v is not None}
+
+        # Add tools
+        if tools is not None and len(tools) > 0:
+            request_params["tools"] = [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": tool["function"]["name"],
+                        "description": tool["function"]["description"],
+                        "parameters": tool["function"]["parameters"],
+                    },
+                }
+                for tool in tools
+            ]
+            # Cerebras requires parallel_tool_calls=False for llama-4-scout-17b-16e-instruct
+            if self.id == "llama-4-scout-17b-16e-instruct":
+                request_params["parallel_tool_calls"] = False
+            elif self.parallel_tool_calls is not None:
+                request_params["parallel_tool_calls"] = self.parallel_tool_calls
+
+        # Handle response format for structured outputs
+        if response_format is not None:
+            if (
+                isinstance(response_format, dict)
+                and response_format.get("type") == "json_schema"
+                and isinstance(response_format.get("json_schema"), dict)
+            ):
+                # Ensure json_schema has strict parameter set
+                schema = response_format["json_schema"]
+                if isinstance(schema.get("schema"), dict) and "strict" not in schema:
+                    schema["strict"] = self.strict_output
+
+                request_params["response_format"] = response_format
+
+        # Add additional request params if provided
+        if self.request_params:
+            request_params.update(self.request_params)
+
+        if request_params:
+            log_debug(f"Calling {self.provider} with request parameters: {request_params}", log_level=2)
+        return request_params
+
+    def invoke(
+        self,
+        messages: List[Message],
+        assistant_message: Message,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        run_response: Optional[RunOutput] = None,
+        compress_tool_results: bool = False,
+    ) -> ModelResponse:
+        """
+        Send a chat completion request to the Cerebras API.
+
+        Args:
+            messages (List[Message]): A list of messages to send to the model.
+
+        Returns:
+            CompletionResponse: The chat completion response from the API.
+        """
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+        provider_response = self.get_client().chat.completions.create(
+            model=self.id,
+            messages=[self._format_message(m, compress_tool_results) for m in messages],  # type: ignore
+            **self.get_request_params(response_format=response_format, tools=tools),
+        )
+        assistant_message.metrics.stop_timer()
+
+        model_response = self._parse_provider_response(provider_response, response_format=response_format)  # type: ignore
+
+        return model_response
+
+    async def ainvoke(
+        self,
+        messages: List[Message],
+        assistant_message: Message,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        run_response: Optional[RunOutput] = None,
+        compress_tool_results: bool = False,
+    ) -> ModelResponse:
+        """
+        Sends an asynchronous chat completion request to the Cerebras API.
+
+        Args:
+            messages (List[Message]): A list of messages to send to the model.
+
+        Returns:
+            ChatCompletion: The chat completion response from the API.
+        """
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+        provider_response = await self.get_async_client().chat.completions.create(
+            model=self.id,
+            messages=[self._format_message(m, compress_tool_results) for m in messages],  # type: ignore
+            **self.get_request_params(response_format=response_format, tools=tools),
+        )
+        assistant_message.metrics.stop_timer()
+
+        model_response = self._parse_provider_response(provider_response, response_format=response_format)  # type: ignore
+
+        return model_response
+
+    def invoke_stream(
+        self,
+        messages: List[Message],
+        assistant_message: Message,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        run_response: Optional[RunOutput] = None,
+        compress_tool_results: bool = False,
+    ) -> Iterator[ModelResponse]:
+        """
+        Send a streaming chat completion request to the Cerebras API.
+
+        Args:
+            messages (List[Message]): A list of messages to send to the model.
+
+        Returns:
+            Iterator[ChatChunkResponse]: An iterator of chat completion chunks.
+        """
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+
+        for chunk in self.get_client().chat.completions.create(
+            model=self.id,
+            messages=[self._format_message(m, compress_tool_results) for m in messages],  # type: ignore
+            stream=True,
+            **self.get_request_params(response_format=response_format, tools=tools),
+        ):
+            yield self._parse_provider_response_delta(chunk)  # type: ignore
+
+        assistant_message.metrics.stop_timer()
+
+    async def ainvoke_stream(
+        self,
+        messages: List[Message],
+        assistant_message: Message,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        run_response: Optional[RunOutput] = None,
+        compress_tool_results: bool = False,
+    ) -> AsyncIterator[ModelResponse]:
+        """
+        Sends an asynchronous streaming chat completion request to the Cerebras API.
+
+        Args:
+            messages (List[Message]): A list of messages to send to the model.
+
+        Returns:
+            AsyncIterator[ChatChunkResponse]: An asynchronous iterator of chat completion chunks.
+        """
+        if run_response and run_response.metrics:
+            run_response.metrics.set_time_to_first_token()
+
+        assistant_message.metrics.start_timer()
+
+        async_stream = await self.get_async_client().chat.completions.create(
+            model=self.id,
+            messages=[self._format_message(m, compress_tool_results) for m in messages],  # type: ignore
+            stream=True,
+            **self.get_request_params(response_format=response_format, tools=tools),
+        )
+
+        async for chunk in async_stream:  # type: ignore
+            yield self._parse_provider_response_delta(chunk)  # type: ignore
+
+        assistant_message.metrics.stop_timer()
+
+    def _format_message(self, message: Message, compress_tool_results: bool = False) -> Dict[str, Any]:
+        """
+        Format a message into the format expected by the Cerebras API.
+
+        Args:
+            message (Message): The message to format.
+            compress_tool_results: Whether to compress tool results.
+
+        Returns:
+            Dict[str, Any]: The formatted message.
+        """
+        # Use compressed content for tool messages if compression is active
+        if message.role == "tool":
+            content = message.get_content(use_compressed_content=compress_tool_results)
+        else:
+            content = message.content if message.content is not None else ""
+
+        # Basic message content
+        message_dict: Dict[str, Any] = {
+            "role": message.role,
+            "content": content,
+        }
+
+        # Add name if present
+        if message.name:
+            message_dict["name"] = message.name
+
+        # Handle tool calls
+        if message.tool_calls:
+            # Ensure tool_calls is properly formatted
+            message_dict["tool_calls"] = [
+                {
+                    "id": tool_call["id"],
+                    "type": tool_call["type"],
+                    "function": {
+                        "name": tool_call["function"]["name"],
+                        "arguments": json.dumps(tool_call["function"]["arguments"])
+                        if isinstance(tool_call["function"]["arguments"], (dict, list))
+                        else tool_call["function"]["arguments"],
+                    },
+                }
+                for tool_call in message.tool_calls
+            ]
+
+        # Handle tool responses
+        if message.role == "tool" and message.tool_call_id:
+            message_dict = {
+                "role": "tool",
+                "tool_call_id": message.tool_call_id,
+                "content": content,
+            }
+
+        # Ensure no None values in the message
+        message_dict = {k: v for k, v in message_dict.items() if v is not None}
+
+        return message_dict
+
+    def _parse_provider_response(self, response: ChatCompletionResponse, **kwargs) -> ModelResponse:
+        """
+        Parse the Cerebras response into a ModelResponse.
+
+        Args:
+            response (CompletionResponse): The response from the Cerebras API.
+
+        Returns:
+            ModelResponse: The parsed response.
+        """
+        model_response = ModelResponse()
+
+        # Get the first choice (assuming single response)
+        choice: ChatCompletionResponseChoice = response.choices[0]
+        message: ChatCompletionResponseChoiceMessage = choice.message
+
+        # Add role
+        if message.role is not None:
+            model_response.role = message.role
+
+        # Add content
+        if message.content is not None:
+            model_response.content = message.content
+
+        # Add tool calls
+        if message.tool_calls is not None:
+            try:
+                model_response.tool_calls = [
+                    {
+                        "id": tool_call.id,
+                        "type": tool_call.type,
+                        "function": {
+                            "name": tool_call.function.name,
+                            "arguments": tool_call.function.arguments,
+                        },
+                    }
+                    for tool_call in message.tool_calls
+                ]
+            except Exception as e:
+                log_warning(f"Error processing tool calls: {e}")
+
+        # Add usage metrics
+        if response.usage:
+            model_response.response_usage = self._get_metrics(response.usage)
+
+        return model_response
+
+    def _parse_provider_response_delta(
+        self, response: Union[ChatChunkResponse, ChatCompletionResponse]
+    ) -> ModelResponse:
+        """
+        Parse the streaming response from the Cerebras API into a ModelResponse.
+
+        Args:
+            response (ChatChunkResponse): The streaming response chunk.
+
+        Returns:
+            ModelResponse: The parsed response.
+        """
+        model_response = ModelResponse()
+
+        # Get the first choice (assuming single response)
+        if response.choices is not None:
+            choice: Union[ChatChunkResponseChoice, ChatCompletionResponseChoice] = response.choices[0]
+            choice_delta: ChatChunkResponseChoiceDelta = choice.delta  # type: ignore
+
+            if choice_delta:
+                # Add content
+                if choice_delta.content:
+                    model_response.content = choice_delta.content
+
+                # Add tool calls - preserve index for proper aggregation in parse_tool_calls
+                if choice_delta.tool_calls:
+                    model_response.tool_calls = [
+                        {
+                            "index": tool_call.index if hasattr(tool_call, "index") else idx,
+                            "id": tool_call.id,
+                            "type": tool_call.type,
+                            "function": {
+                                "name": tool_call.function.name if tool_call.function else None,
+                                "arguments": tool_call.function.arguments if tool_call.function else None,
+                            },
+                        }
+                        for idx, tool_call in enumerate(choice_delta.tool_calls)
+                    ]
+
+        # Add usage metrics
+        if response.usage:
+            model_response.response_usage = self._get_metrics(response.usage)
+
+        return model_response
+
+    def parse_tool_calls(self, tool_calls_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Build complete tool calls from streamed tool call delta data.
+
+        Cerebras streams tool calls incrementally with partial data in each chunk.
+        This method aggregates those chunks by index to produce complete tool calls.
+
+        Args:
+            tool_calls_data: List of tool call deltas from streaming chunks.
+
+        Returns:
+            List[Dict[str, Any]]: List of fully-formed tool call dicts.
+        """
+        tool_calls: List[Dict[str, Any]] = []
+
+        for tool_call_delta in tool_calls_data:
+            # Get the index for this tool call (default to 0 if not present)
+            index = tool_call_delta.get("index", 0)
+
+            # Extend the list if needed
+            while len(tool_calls) <= index:
+                tool_calls.append(
+                    {
+                        "id": None,
+                        "type": None,
+                        "function": {
+                            "name": "",
+                            "arguments": "",
+                        },
+                    }
+                )
+
+            tool_call_entry = tool_calls[index]
+
+            # Update id if present
+            if tool_call_delta.get("id"):
+                tool_call_entry["id"] = tool_call_delta["id"]
+
+            # Update type if present
+            if tool_call_delta.get("type"):
+                tool_call_entry["type"] = tool_call_delta["type"]
+
+            # Update function name and arguments (concatenate for streaming)
+            if tool_call_delta.get("function"):
+                func_delta = tool_call_delta["function"]
+                if func_delta.get("name"):
+                    tool_call_entry["function"]["name"] += func_delta["name"]
+                if func_delta.get("arguments"):
+                    tool_call_entry["function"]["arguments"] += func_delta["arguments"]
+
+        # Filter out any incomplete tool calls (missing id or function name)
+        complete_tool_calls = [tc for tc in tool_calls if tc.get("id") and tc.get("function", {}).get("name")]
+
+        return complete_tool_calls
+
+    def _get_metrics(self, response_usage: Union[ChatCompletionResponseUsage, ChatChunkResponseUsage]) -> Metrics:
+        """
+        Parse the given Cerebras usage into an Agno Metrics object.
+
+        Args:
+            response_usage: Usage data from Cerebras
+
+        Returns:
+            Metrics: Parsed metrics data
+        """
+        metrics = Metrics()
+
+        metrics.input_tokens = response_usage.prompt_tokens or 0
+        metrics.output_tokens = response_usage.completion_tokens or 0
+        metrics.total_tokens = metrics.input_tokens + metrics.output_tokens
+
+        return metrics