agno 1.8.1__py3-none-any.whl → 2.0.0a1__py3-none-any.whl

agno/models/cerebras/cerebras.py

@@ -9,20 +9,23 @@ 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.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 import ChatCompletion
     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`")
@@ -198,10 +201,12 @@ class Cerebras(Model):
     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,
-    ) -> ChatCompletion:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Send a chat completion request to the Cerebras API.
 
@@ -211,19 +216,30 @@ class Cerebras(Model):
         Returns:
             CompletionResponse: The chat completion response from the API.
         """
-        return self.get_client().chat.completions.create(
+        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) 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,
-    ) -> ChatCompletion:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Sends an asynchronous chat completion request to the Cerebras API.
 
@@ -233,19 +249,30 @@ class Cerebras(Model):
         Returns:
             ChatCompletion: The chat completion response from the API.
         """
-        return await self.get_async_client().chat.completions.create(
+        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) 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,
-    ) -> Iterator[ChatChunkResponse]:
+        run_response: Optional[RunOutput] = None,
+    ) -> Iterator[ModelResponse]:
         """
         Send a streaming chat completion request to the Cerebras API.
 
@@ -255,20 +282,30 @@ class Cerebras(Model):
         Returns:
             Iterator[ChatChunkResponse]: An iterator of chat completion chunks.
         """
-        yield from self.get_client().chat.completions.create(
+        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) for m in messages],  # type: ignore
             stream=True,
             **self.get_request_params(response_format=response_format, tools=tools),
-        )  # type: ignore
+        ):
+            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,
-    ) -> AsyncIterator[ChatChunkResponse]:
+        run_response: Optional[RunOutput] = None,
+    ) -> AsyncIterator[ModelResponse]:
         """
         Sends an asynchronous streaming chat completion request to the Cerebras API.
 
@@ -278,14 +315,22 @@ class Cerebras(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) 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 chunk  # type: ignore
+            yield self._parse_provider_response_delta(chunk)  # type: ignore
+
+        assistant_message.metrics.stop_timer()
 
     def _format_message(self, message: Message) -> Dict[str, Any]:
         """
@@ -337,7 +382,7 @@ class Cerebras(Model):
 
         return message_dict
 
-    def parse_provider_response(self, response: ChatCompletionResponse, **kwargs) -> ModelResponse:
+    def _parse_provider_response(self, response: ChatCompletionResponse, **kwargs) -> ModelResponse:
         """
         Parse the Cerebras response into a ModelResponse.
 
@@ -380,20 +425,18 @@ class Cerebras(Model):
 
         # Add usage metrics
         if response.usage:
-            model_response.response_usage = {
-                "input_tokens": response.usage.prompt_tokens,
-                "output_tokens": response.usage.completion_tokens,
-                "total_tokens": response.usage.total_tokens,
-            }
+            model_response.response_usage = self._get_metrics(response.usage)
 
         return model_response
 
-    def parse_provider_response_delta(self, response_delta: ChatChunkResponse) -> ModelResponse:
+    def _parse_provider_response_delta(
+        self, response: Union[ChatChunkResponse, ChatCompletionResponse]
+    ) -> ModelResponse:
         """
         Parse the streaming response from the Cerebras API into a ModelResponse.
 
         Args:
-            response_delta (ChatChunkResponse): The streaming response chunk.
+            response (ChatChunkResponse): The streaming response chunk.
 
         Returns:
             ModelResponse: The parsed response.
@@ -401,9 +444,9 @@ class Cerebras(Model):
         model_response = ModelResponse()
 
         # Get the first choice (assuming single response)
-        if response_delta.choices is not None:
-            choice: ChatChunkResponseChoice = response_delta.choices[0]
-            choice_delta: ChatChunkResponseChoiceDelta = choice.delta
+        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
@@ -425,11 +468,25 @@ class Cerebras(Model):
                     ]
 
         # Add usage metrics
-        if response_delta.usage:
-            model_response.response_usage = {
-                "input_tokens": response_delta.usage.prompt_tokens,
-                "output_tokens": response_delta.usage.completion_tokens,
-                "total_tokens": response_delta.usage.total_tokens,
-            }
+        if response.usage:
+            model_response.response_usage = self._get_metrics(response.usage)
 
         return model_response
+
+    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

agno/models/cohere/chat.py

@@ -5,17 +5,19 @@ from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple, Ty
 from pydantic import BaseModel
 
 from agno.exceptions import ModelProviderError
-from agno.models.base import MessageData, Model, _add_usage_metrics_to_assistant_message
+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.log import log_debug, log_error
 from agno.utils.models.cohere import format_messages
 
 try:
     from cohere import AsyncClientV2 as CohereAsyncClient
     from cohere import ClientV2 as CohereClient
-    from cohere.types.chat_response import ChatResponse
-    from cohere.types.streamed_chat_response_v2 import StreamedChatResponseV2
+    from cohere.v2.types.v2chat_response import V2ChatResponse
+    from cohere.v2.types.v2chat_stream_response import V2ChatStreamResponse
 except ImportError:
     raise ImportError("`cohere` not installed. Please install using `pip install cohere`")
 
@@ -148,18 +150,33 @@ class Cohere(Model):
     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,
-    ) -> ChatResponse:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Invoke a non-streamed chat response from the Cohere API.
         """
-
         request_kwargs = self.get_request_params(response_format=response_format, tools=tools)
 
         try:
-            return self.get_client().chat(model=self.id, messages=format_messages(messages), **request_kwargs)  # type: ignore
+            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(
+                model=self.id,
+                messages=format_messages(messages),  # type: ignore
+                **request_kwargs,
+            )  # type: ignore
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(provider_response, response_format=response_format)
+
+            return model_response
+
         except Exception as e:
             log_error(f"Unexpected error calling Cohere API: {str(e)}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -167,21 +184,35 @@ class Cohere(Model):
     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,
-    ) -> Iterator[StreamedChatResponseV2]:
+        run_response: Optional[RunOutput] = None,
+    ) -> Iterator[ModelResponse]:
         """
         Invoke a streamed chat response from the Cohere API.
         """
         request_kwargs = self.get_request_params(response_format=response_format, tools=tools)
 
         try:
-            return self.get_client().chat_stream(
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            tool_use: Dict[str, Any] = {}
+
+            assistant_message.metrics.start_timer()
+
+            for response in self.get_client().chat_stream(
                 model=self.id,
                 messages=format_messages(messages),  # type: ignore
                 **request_kwargs,
-            )
+            ):
+                model_response, tool_use = self._parse_provider_response_delta(response, tool_use=tool_use)
+                yield model_response
+
+            assistant_message.metrics.stop_timer()
+
         except Exception as e:
             log_error(f"Unexpected error calling Cohere API: {str(e)}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -189,21 +220,33 @@ class Cohere(Model):
     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,
-    ) -> ChatResponse:
+        run_response: Optional[RunOutput] = None,
+    ) -> ModelResponse:
         """
         Asynchronously invoke a non-streamed chat response from the Cohere API.
         """
         request_kwargs = self.get_request_params(response_format=response_format, tools=tools)
 
         try:
-            return await self.get_async_client().chat(
+            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(
                 model=self.id,
                 messages=format_messages(messages),  # type: ignore
                 **request_kwargs,
             )
+            assistant_message.metrics.stop_timer()
+
+            model_response = self._parse_provider_response(provider_response, response_format=response_format)
+
+            return model_response
+
         except Exception as e:
             log_error(f"Unexpected error calling Cohere API: {str(e)}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -211,32 +254,45 @@ class Cohere(Model):
     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,
-    ) -> AsyncIterator[StreamedChatResponseV2]:
+        run_response: Optional[RunOutput] = None,
+    ) -> AsyncIterator[ModelResponse]:
         """
         Asynchronously invoke a streamed chat response from the Cohere API.
         """
         request_kwargs = self.get_request_params(response_format=response_format, tools=tools)
 
         try:
+            if run_response and run_response.metrics:
+                run_response.metrics.set_time_to_first_token()
+
+            tool_use: Dict[str, Any] = {}
+
+            assistant_message.metrics.start_timer()
+
             async for response in self.get_async_client().chat_stream(
                 model=self.id,
                 messages=format_messages(messages),  # type: ignore
                 **request_kwargs,
             ):
-                yield response
+                model_response, tool_use = self._parse_provider_response_delta(response, tool_use=tool_use)
+                yield model_response
+
+            assistant_message.metrics.stop_timer()
+
         except Exception as e:
             log_error(f"Unexpected error calling Cohere API: {str(e)}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
 
-    def parse_provider_response(self, response: ChatResponse, **kwargs) -> ModelResponse:
+    def _parse_provider_response(self, response: V2ChatResponse, **kwargs) -> ModelResponse:
         """
         Parse the model provider response.
 
         Args:
-            response (ChatResponse): The response from the Cohere API.
+            response (V2ChatResponse): The response from the Cohere API.
         """
         model_response = ModelResponse()
 
@@ -246,136 +302,88 @@ class Cohere(Model):
         if response_message.content is not None:
             full_content = ""
             for item in response_message.content:
-                full_content += item.text
+                if hasattr(item, "text") and item.text is not None:  # type: ignore
+                    full_content += item.text  # type: ignore
             model_response.content = full_content
 
         if response_message.tool_calls is not None:
             model_response.tool_calls = [t.model_dump() for t in response_message.tool_calls]
 
-        if response.usage is not None and response.usage.tokens is not None:
-            model_response.response_usage = {
-                "input_tokens": int(response.usage.tokens.input_tokens) or 0,  # type: ignore
-                "output_tokens": int(response.usage.tokens.output_tokens) or 0,  # type: ignore
-                "total_tokens": int(response.usage.tokens.input_tokens + response.usage.tokens.output_tokens) or 0,  # type: ignore
-            }
+        if response.usage is not None:
+            model_response.response_usage = self._get_metrics(response.usage)
 
         return model_response
 
-    def _process_stream_response(
-        self,
-        response: StreamedChatResponseV2,
-        assistant_message: Message,
-        stream_data: MessageData,
-        tool_use: Dict[str, Any],
-    ) -> Tuple[Optional[ModelResponse], Dict[str, Any]]:
+    def _parse_provider_response_delta(
+        self, response: V2ChatStreamResponse, tool_use: Dict[str, Any]
+    ) -> Tuple[ModelResponse, Dict[str, Any]]:  # type: ignore
         """
-        Common handler for processing stream responses from Cohere.
+        Parse the streaming response from the model provider into ModelResponse objects.
 
         Args:
-            response: The streamed response from Cohere
-            assistant_message: The assistant message being built
-            stream_data: Data accumulated during streaming
-            tool_use: Current tool use data being built
+            response: Raw response chunk from the model provider
 
         Returns:
-            Tuple containing the ModelResponse to yield and updated tool_use dict
+            ModelResponse: Parsed response delta
         """
-        model_response = None
+        model_response = ModelResponse()
 
+        # 1. Add content
         if (
             response.type == "content-delta"
             and response.delta is not None
             and response.delta.message is not None
             and response.delta.message.content is not None
+            and response.delta.message.content.text is not None
         ):
-            # Update metrics
-            assistant_message.metrics.completion_tokens += 1
-            if not assistant_message.metrics.time_to_first_token:
-                assistant_message.metrics.set_time_to_first_token()
+            model_response.content = response.delta.message.content.text
 
-            # Update provider response content
-            stream_data.response_content += response.delta.message.content.text
-            model_response = ModelResponse(content=response.delta.message.content.text)
+        # 2. Add tool calls information
 
+        # 2.1 Add starting tool call
         elif response.type == "tool-call-start" and response.delta is not None:
             if response.delta.message is not None and response.delta.message.tool_calls is not None:
                 tool_use = response.delta.message.tool_calls.model_dump()
 
+        # 2.2 Add tool call delta
         elif response.type == "tool-call-delta" and response.delta is not None:
             if (
                 response.delta.message is not None
                 and response.delta.message.tool_calls is not None
                 and response.delta.message.tool_calls.function is not None
+                and response.delta.message.tool_calls.function.arguments is not None
             ):
                 tool_use["function"]["arguments"] += response.delta.message.tool_calls.function.arguments
 
+        # 2.3 Add ending tool call
         elif response.type == "tool-call-end":
-            if assistant_message.tool_calls is None:
-                assistant_message.tool_calls = []
-            assistant_message.tool_calls.append(tool_use)
-            tool_use = {}
+            model_response.tool_calls = [tool_use]
 
+        # 3. Add metrics
         elif (
             response.type == "message-end"
             and response.delta is not None
             and response.delta.usage is not None
             and response.delta.usage.tokens is not None
         ):
-            _add_usage_metrics_to_assistant_message(
-                assistant_message=assistant_message,
-                response_usage={
-                    "input_tokens": int(response.delta.usage.tokens.input_tokens) or 0,  # type: ignore
-                    "output_tokens": int(response.delta.usage.tokens.output_tokens) or 0,  # type: ignore
-                    "total_tokens": int(
-                        response.delta.usage.tokens.input_tokens + response.delta.usage.tokens.output_tokens  # type: ignore
-                    )
-                    or 0,
-                },
-            )
+            model_response.response_usage = self._get_metrics(response.delta.usage)  # type: ignore
 
         return model_response, tool_use
 
-    def process_response_stream(
-        self,
-        messages: List[Message],
-        assistant_message: Message,
-        stream_data: MessageData,
-        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
-    ) -> Iterator[ModelResponse]:
-        """Process the synchronous response stream."""
-        tool_use: Dict[str, Any] = {}
+    def _get_metrics(self, response_usage) -> Metrics:
+        """
+        Parse the given Cohere usage into an Agno Metrics object.
 
-        for response in self.invoke_stream(
-            messages=messages, response_format=response_format, tools=tools, tool_choice=tool_choice
-        ):
-            model_response, tool_use = self._process_stream_response(
-                response=response, assistant_message=assistant_message, stream_data=stream_data, tool_use=tool_use
-            )
-            if model_response is not None:
-                yield model_response
+        Args:
+            response_usage: Usage data from Cohere
 
-    async def aprocess_response_stream(
-        self,
-        messages: List[Message],
-        assistant_message: Message,
-        stream_data: MessageData,
-        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
-    ) -> AsyncIterator[ModelResponse]:
-        """Process the asynchronous response stream."""
-        tool_use: Dict[str, Any] = {}
+        Returns:
+            Metrics: Parsed metrics data
+        """
+        metrics = Metrics()
 
-        async for response in self.ainvoke_stream(
-            messages=messages, response_format=response_format, tools=tools, tool_choice=tool_choice
-        ):
-            model_response, tool_use = self._process_stream_response(
-                response=response, assistant_message=assistant_message, stream_data=stream_data, tool_use=tool_use
-            )
-            if model_response is not None:
-                yield model_response
+        metrics.input_tokens = response_usage.tokens.input_tokens or 0
+        metrics.output_tokens = response_usage.tokens.output_tokens or 0
+        metrics.total_tokens = metrics.input_tokens + metrics.output_tokens
 
-    def parse_provider_response_delta(self, response: Any) -> ModelResponse:  # type: ignore
-        pass
+        return metrics