Skip to content

cohere.calls

A module for calling Cohere's chat models.

CohereCall

Bases: BaseCall[CohereCallResponse, CohereCallResponseChunk, CohereTool]

A base class for calling Cohere's chat models.

Example:

from mirascope.cohere import CohereCall


class BookRecommender(CohereCall):
    prompt_template = "Please recommend a {genre} book"

    genre: str

response = BookRecommender(genre="fantasy").call()
print(response.content)
#> There are many great books to read, it ultimately depends...
Source code in mirascope/cohere/calls.py 
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
class CohereCall(BaseCall[CohereCallResponse, CohereCallResponseChunk, CohereTool]):
    """A base class for calling Cohere's chat models.

    Example:

    ```python
    from mirascope.cohere import CohereCall


    class BookRecommender(CohereCall):
        prompt_template = "Please recommend a {genre} book"

        genre: str

    response = BookRecommender(genre="fantasy").call()
    print(response.content)
    #> There are many great books to read, it ultimately depends...
    ```
    """

    call_params: ClassVar[CohereCallParams] = CohereCallParams()

    def messages(self) -> list[ChatMessage]:
        """Returns the template as a formatted list of messages."""
        return [
            ChatMessage(role=message["role"].upper(), message=message["content"])
            for message in self._parse_messages(
                [MessageRole.SYSTEM, MessageRole.USER, MessageRole.CHATBOT]
            )
        ]

    @retry
    def call(
        self, retries: Union[int, Retrying] = 1, **kwargs: Any
    ) -> CohereCallResponse:
        """Makes a call to the model using this `CohereCall` instance.

        Args:
            **kwargs: Additional keyword arguments parameters to pass to the call. These
                will override any existing arguments in `call_params`.

        Returns:
            A `CohereCallResponse` instance.
        """
        message, kwargs, tool_types = self._setup_cohere_kwargs(kwargs)
        co = Client(api_key=self.api_key, base_url=self.base_url)
        if self.call_params.wrapper is not None:
            co = self.call_params.wrapper(co)
        chat = co.chat
        if self.call_params.weave is not None:
            chat = self.call_params.weave(chat)  # pragma: no cover
        if self.call_params.logfire:
            chat = self.call_params.logfire(
                chat, "cohere", response_type=CohereCallResponse, tool_types=tool_types
            )  # pragma: no cover
        start_time = datetime.datetime.now().timestamp() * 1000
        response = chat(message=message, **kwargs)
        cost = None
        if response.meta and response.meta.billed_units:
            cost = cohere_api_calculate_cost(
                response.meta.billed_units, self.call_params.model
            )
        return CohereCallResponse(
            response=response,
            tool_types=tool_types,
            start_time=start_time,
            end_time=datetime.datetime.now().timestamp() * 1000,
            cost=cost,
        )

    @retry
    async def call_async(
        self, retries: Union[int, AsyncRetrying] = 1, **kwargs: Any
    ) -> CohereCallResponse:
        """Makes an asynchronous call to the model using this `CohereCall`.

        Args:
            **kwargs: Additional keyword arguments parameters to pass to the call. These
                will override any existing arguments in `call_params`.

        Returns:
            An `CohereCallResponse` instance.
        """
        message, kwargs, tool_types = self._setup_cohere_kwargs(kwargs)
        co = AsyncClient(api_key=self.api_key, base_url=self.base_url)
        if self.call_params.wrapper_async is not None:
            co = self.call_params.wrapper_async(co)
        chat = co.chat
        if self.call_params.weave is not None:
            chat = self.call_params.weave(chat)  # pragma: no cover
        if self.call_params.logfire_async:
            chat = self.call_params.logfire_async(
                chat, "cohere", response_type=CohereCallResponse, tool_types=tool_types
            )  # pragma: no cover
        start_time = datetime.datetime.now().timestamp() * 1000
        response = await chat(message=message, **kwargs)
        cost = None
        if response.meta and response.meta.billed_units:
            cost = cohere_api_calculate_cost(
                response.meta.billed_units, self.call_params.model
            )
        return CohereCallResponse(
            response=response,
            tool_types=tool_types,
            start_time=start_time,
            end_time=datetime.datetime.now().timestamp() * 1000,
            cost=cost,
        )

    @retry
    def stream(
        self, retries: Union[int, Retrying] = 1, **kwargs: Any
    ) -> Generator[CohereCallResponseChunk, None, None]:
        """Streams the response for a call using this `CohereCall`.

        Args:
            **kwargs: Additional keyword arguments parameters to pass to the call. These
                will override any existing arguments in `call_params`.

        Yields:
            A `CohereCallResponseChunk` for each chunk of the response.
        """
        message, kwargs, tool_types = self._setup_cohere_kwargs(kwargs)
        co = Client(api_key=self.api_key, base_url=self.base_url)
        if self.call_params.wrapper is not None:
            co = self.call_params.wrapper(co)
        chat_stream = co.chat_stream
        if self.call_params.weave is not None:
            chat_stream = self.call_params.weave(chat_stream)  # pragma: no cover
        if self.call_params.logfire:
            chat_stream = self.call_params.logfire(
                chat_stream, "cohere", response_chunk_type=CohereCallResponseChunk
            )  # pragma: no cover
        for event in chat_stream(message=message, **kwargs):
            yield CohereCallResponseChunk(chunk=event, tool_types=tool_types)

    @retry
    async def stream_async(
        self, retries: Union[int, AsyncRetrying] = 1, **kwargs: Any
    ) -> AsyncGenerator[CohereCallResponseChunk, None]:
        """Streams the response for an asynchronous call using this `CohereCall`.

        Args:
            **kwargs: Additional keyword arguments parameters to pass to the call. These
                will override any existing arguments in `call_params`.

        Yields:
            A `CohereCallResponseChunk` for each chunk of the response.
        """
        message, kwargs, tool_types = self._setup_cohere_kwargs(kwargs)
        co = AsyncClient(api_key=self.api_key, base_url=self.base_url)
        if self.call_params.wrapper_async is not None:
            co = self.call_params.wrapper_async(co)
        chat_stream = co.chat_stream
        if self.call_params.weave is not None:
            chat_stream = self.call_params.weave(chat_stream)  # pragma: no cover
        if self.call_params.logfire_async:
            chat_stream = self.call_params.logfire_async(
                chat_stream, "cohere", response_chunk_type=CohereCallResponseChunk
            )  # pragma: no cover
        async for event in chat_stream(message=message, **kwargs):
            yield CohereCallResponseChunk(chunk=event, tool_types=tool_types)

    ############################## PRIVATE METHODS ###################################

    def _setup_cohere_kwargs(
        self, kwargs: dict[str, Any]
    ) -> tuple[str, dict[str, Any], Optional[list[Type[CohereTool]]]]:
        """Overrides the `BaseCall._setup` for Cohere specific setup."""
        kwargs, tool_types = self._setup(kwargs, CohereTool)
        messages = self.messages()
        preamble = ""
        if "preamble" in kwargs and kwargs["preamble"] is not None:
            preamble += kwargs.pop("preamble")
        if messages[0].role == "SYSTEM":
            preamble += messages.pop(0).message
        if preamble:
            kwargs["preamble"] = preamble
        if len(messages) > 1:
            kwargs["chat_history"] = messages[:-1]
        if hasattr(self, "documents"):
            kwargs["documents"] = self.documents
        return messages[-1].message, kwargs, tool_types

call(retries=1, **kwargs)

Makes a call to the model using this CohereCall instance.

Parameters:

Name Type Description Default
**kwargs Any

Additional keyword arguments parameters to pass to the call. These will override any existing arguments in call_params.

 {}

Returns:

Type Description
CohereCallResponse

A CohereCallResponse instance.
Source code in mirascope/cohere/calls.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@retry
def call(
    self, retries: Union[int, Retrying] = 1, **kwargs: Any
) -> CohereCallResponse:
    """Makes a call to the model using this `CohereCall` instance.

    Args:
        **kwargs: Additional keyword arguments parameters to pass to the call. These
            will override any existing arguments in `call_params`.

    Returns:
        A `CohereCallResponse` instance.
    """
    message, kwargs, tool_types = self._setup_cohere_kwargs(kwargs)
    co = Client(api_key=self.api_key, base_url=self.base_url)
    if self.call_params.wrapper is not None:
        co = self.call_params.wrapper(co)
    chat = co.chat
    if self.call_params.weave is not None:
        chat = self.call_params.weave(chat)  # pragma: no cover
    if self.call_params.logfire:
        chat = self.call_params.logfire(
            chat, "cohere", response_type=CohereCallResponse, tool_types=tool_types
        )  # pragma: no cover
    start_time = datetime.datetime.now().timestamp() * 1000
    response = chat(message=message, **kwargs)
    cost = None
    if response.meta and response.meta.billed_units:
        cost = cohere_api_calculate_cost(
            response.meta.billed_units, self.call_params.model
        )
    return CohereCallResponse(
        response=response,
        tool_types=tool_types,
        start_time=start_time,
        end_time=datetime.datetime.now().timestamp() * 1000,
        cost=cost,
    )

call_async(retries=1, **kwargs) async

Makes an asynchronous call to the model using this CohereCall.

Parameters:

Name Type Description Default
**kwargs Any

Additional keyword arguments parameters to pass to the call. These will override any existing arguments in call_params.

 {}

Returns:

Type Description
CohereCallResponse

An CohereCallResponse instance.
Source code in mirascope/cohere/calls.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
@retry
async def call_async(
    self, retries: Union[int, AsyncRetrying] = 1, **kwargs: Any
) -> CohereCallResponse:
    """Makes an asynchronous call to the model using this `CohereCall`.

    Args:
        **kwargs: Additional keyword arguments parameters to pass to the call. These
            will override any existing arguments in `call_params`.

    Returns:
        An `CohereCallResponse` instance.
    """
    message, kwargs, tool_types = self._setup_cohere_kwargs(kwargs)
    co = AsyncClient(api_key=self.api_key, base_url=self.base_url)
    if self.call_params.wrapper_async is not None:
        co = self.call_params.wrapper_async(co)
    chat = co.chat
    if self.call_params.weave is not None:
        chat = self.call_params.weave(chat)  # pragma: no cover
    if self.call_params.logfire_async:
        chat = self.call_params.logfire_async(
            chat, "cohere", response_type=CohereCallResponse, tool_types=tool_types
        )  # pragma: no cover
    start_time = datetime.datetime.now().timestamp() * 1000
    response = await chat(message=message, **kwargs)
    cost = None
    if response.meta and response.meta.billed_units:
        cost = cohere_api_calculate_cost(
            response.meta.billed_units, self.call_params.model
        )
    return CohereCallResponse(
        response=response,
        tool_types=tool_types,
        start_time=start_time,
        end_time=datetime.datetime.now().timestamp() * 1000,
        cost=cost,
    )

messages()

Returns the template as a formatted list of messages.

Source code in mirascope/cohere/calls.py 
38
39
40
41
42
43
44
45
def messages(self) -> list[ChatMessage]:
    """Returns the template as a formatted list of messages."""
    return [
        ChatMessage(role=message["role"].upper(), message=message["content"])
        for message in self._parse_messages(
            [MessageRole.SYSTEM, MessageRole.USER, MessageRole.CHATBOT]
        )
    ]

stream(retries=1, **kwargs)

Streams the response for a call using this CohereCall.

Parameters:

Name Type Description Default
**kwargs Any

Additional keyword arguments parameters to pass to the call. These will override any existing arguments in call_params.

 {}

Yields:

Type Description
CohereCallResponseChunk

A CohereCallResponseChunk for each chunk of the response.
Source code in mirascope/cohere/calls.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
@retry
def stream(
    self, retries: Union[int, Retrying] = 1, **kwargs: Any
) -> Generator[CohereCallResponseChunk, None, None]:
    """Streams the response for a call using this `CohereCall`.

    Args:
        **kwargs: Additional keyword arguments parameters to pass to the call. These
            will override any existing arguments in `call_params`.

    Yields:
        A `CohereCallResponseChunk` for each chunk of the response.
    """
    message, kwargs, tool_types = self._setup_cohere_kwargs(kwargs)
    co = Client(api_key=self.api_key, base_url=self.base_url)
    if self.call_params.wrapper is not None:
        co = self.call_params.wrapper(co)
    chat_stream = co.chat_stream
    if self.call_params.weave is not None:
        chat_stream = self.call_params.weave(chat_stream)  # pragma: no cover
    if self.call_params.logfire:
        chat_stream = self.call_params.logfire(
            chat_stream, "cohere", response_chunk_type=CohereCallResponseChunk
        )  # pragma: no cover
    for event in chat_stream(message=message, **kwargs):
        yield CohereCallResponseChunk(chunk=event, tool_types=tool_types)

stream_async(retries=1, **kwargs) async

Streams the response for an asynchronous call using this CohereCall.

Parameters:

Name Type Description Default
**kwargs Any

Additional keyword arguments parameters to pass to the call. These will override any existing arguments in call_params.

 {}

Yields:

Type Description
AsyncGenerator[CohereCallResponseChunk, None]

A CohereCallResponseChunk for each chunk of the response.
Source code in mirascope/cohere/calls.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
@retry
async def stream_async(
    self, retries: Union[int, AsyncRetrying] = 1, **kwargs: Any
) -> AsyncGenerator[CohereCallResponseChunk, None]:
    """Streams the response for an asynchronous call using this `CohereCall`.

    Args:
        **kwargs: Additional keyword arguments parameters to pass to the call. These
            will override any existing arguments in `call_params`.

    Yields:
        A `CohereCallResponseChunk` for each chunk of the response.
    """
    message, kwargs, tool_types = self._setup_cohere_kwargs(kwargs)
    co = AsyncClient(api_key=self.api_key, base_url=self.base_url)
    if self.call_params.wrapper_async is not None:
        co = self.call_params.wrapper_async(co)
    chat_stream = co.chat_stream
    if self.call_params.weave is not None:
        chat_stream = self.call_params.weave(chat_stream)  # pragma: no cover
    if self.call_params.logfire_async:
        chat_stream = self.call_params.logfire_async(
            chat_stream, "cohere", response_chunk_type=CohereCallResponseChunk
        )  # pragma: no cover
    async for event in chat_stream(message=message, **kwargs):
        yield CohereCallResponseChunk(chunk=event, tool_types=tool_types)