ag2 0.10.2__py3-none-any.whl

autogen/oai/client_utils.py

@@ -0,0 +1,196 @@
+# 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
+"""Utilities for client classes"""
+
+import logging
+import warnings
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class FormatterProtocol(Protocol):
+    """Structured Output classes with a format method"""
+
+    def format(self) -> str: ...
+
+
+def validate_parameter(
+    params: dict[str, Any],
+    param_name: str,
+    allowed_types: tuple[Any, ...],
+    allow_None: bool,  # noqa: N803
+    default_value: Any,
+    numerical_bound: tuple[float | None, float | None] | None,
+    allowed_values: list[Any] | None,
+) -> Any:
+    """Validates a given config parameter, checking its type, values, and setting defaults
+    Parameters:
+        params (Dict[str, Any]): Dictionary containing parameters to validate.
+        param_name (str): The name of the parameter to validate.
+        allowed_types (Tuple): Tuple of acceptable types for the parameter.
+        allow_None (bool): Whether the parameter can be `None`.
+        default_value (Any): The default value to use if the parameter is invalid or missing.
+        numerical_bound (Optional[Tuple[Optional[float], Optional[float]]]):
+            A tuple specifying the lower and upper bounds for numerical parameters.
+            Each bound can be `None` if not applicable.
+        allowed_values (Optional[List[Any]]): A list of acceptable values for the parameter.
+            Can be `None` if no specific values are required.
+
+    Returns:
+        Any: The validated parameter value or the default value if validation fails.
+
+    Raises:
+        TypeError: If `allowed_values` is provided but is not a list.
+
+    Example Usage:
+    ```python
+        # Validating a numerical parameter within specific bounds
+        params = {"temperature": 0.5, "safety_model": "Meta-Llama/Llama-Guard-7b"}
+        temperature = validate_parameter(params, "temperature", (int, float), True, 0.7, (0, 1), None)
+        # Result: 0.5
+
+        # Validating a parameter that can be one of a list of allowed values
+        model = validate_parameter(
+        params, "safety_model", str, True, None, None, ["Meta-Llama/Llama-Guard-7b", "Meta-Llama/Llama-Guard-13b"]
+        )
+        # If "safety_model" is missing or invalid in params, defaults to "default"
+    ```
+    """
+    if allowed_values is not None and not isinstance(allowed_values, list):
+        raise TypeError(f"allowed_values should be a list or None, got {type(allowed_values).__name__}")
+
+    param_value = params.get(param_name, default_value)
+    warning = ""
+
+    if param_value is None and allow_None:
+        pass
+    elif param_value is None:
+        if not allow_None:
+            warning = "cannot be None"
+    elif not isinstance(param_value, allowed_types):
+        # Check types and list possible types if invalid
+        if isinstance(allowed_types, tuple):
+            formatted_types = "(" + ", ".join(f"{t.__name__}" for t in allowed_types) + ")"
+        else:
+            formatted_types = f"{allowed_types.__name__}"
+        warning = f"must be of type {formatted_types}{' or None' if allow_None else ''}"
+    elif numerical_bound:
+        # Check the value fits in possible bounds
+        lower_bound, upper_bound = numerical_bound
+        if (lower_bound is not None and param_value < lower_bound) or (
+            upper_bound is not None and param_value > upper_bound
+        ):
+            warning = "has numerical bounds"
+            if lower_bound is not None:
+                warning += f", >= {lower_bound!s}"
+            if upper_bound is not None:
+                if lower_bound is not None:
+                    warning += " and"
+                warning += f" <= {upper_bound!s}"
+            if allow_None:
+                warning += ", or can be None"
+
+    elif allowed_values:  # noqa: SIM102
+        # Check if the value matches any allowed values
+        if not (allow_None and param_value is None) and param_value not in allowed_values:
+            warning = f"must be one of these values [{allowed_values}]{', or can be None' if allow_None else ''}"
+
+    # If we failed any checks, warn and set to default value
+    if warning:
+        warnings.warn(
+            f"Config error - {param_name} {warning}, defaulting to {default_value}.",
+            UserWarning,
+        )
+        param_value = default_value
+
+    return param_value
+
+
+def merge_config_with_tools(config: dict[str, Any], client_config: dict[str, Any]) -> dict[str, Any]:
+    """Merge configuration dictionaries with proper tools and functions handling.
+
+    This function merges two configuration dictionaries while ensuring that:
+    1. Empty 'tools' arrays are not added unnecessarily
+    2. 'tools' and deprecated 'functions' parameters are not both present
+    3. Actual tool configurations are properly merged
+
+    Args:
+        config: The base configuration dictionary (e.g., from create() call)
+        client_config: The client-specific configuration dictionary (e.g., from config_list)
+
+    Returns:
+        dict[str, Any]: The merged configuration with proper tools/functions handling
+    """
+    # Start with a clean merge of both configs
+    full_config = {**config, **client_config}
+
+    # Add tools if tools contains something AND are not using deprecated functions
+    tools = config.get("tools", []) + client_config.get("tools", [])
+    if tools and "functions" not in full_config:
+        # Don't add tools if functions parameter is present (deprecated API)
+        full_config["tools"] = tools
+
+    return full_config
+
+
+def should_hide_tools(messages: list[dict[str, Any]], tools: list[dict[str, Any]], hide_tools_param: str) -> bool:
+    """Determines if tools should be hidden. This function is used to hide tools when they have been run, minimising the chance of the LLM choosing them when they shouldn't.
+    Parameters:
+        messages (List[Dict[str, Any]]): List of messages
+        tools (List[Dict[str, Any]]): List of tools
+        hide_tools_param (str): "hide_tools" parameter value. Can be "if_all_run" (hide tools if all tools have been run), "if_any_run" (hide tools if any of the tools have been run), "never" (never hide tools). Default is "never".
+
+    Returns:
+        bool: Indicates whether the tools should be excluded from the response create request
+
+    Example Usage:
+    ```python
+        # Validating a numerical parameter within specific bounds
+        messages = params.get("messages", [])
+        tools = params.get("tools", None)
+        hide_tools = should_hide_tools(messages, tools, params["hide_tools"])
+    """
+    if hide_tools_param == "never" or tools is None or len(tools) == 0:
+        return False
+    elif hide_tools_param == "if_any_run":
+        # Return True if any tool_call_id exists, indicating a tool call has been executed. False otherwise.
+        return any("tool_call_id" in dictionary for dictionary in messages)
+    elif hide_tools_param == "if_all_run":
+        # Return True if all tools have been executed at least once. False otherwise.
+
+        # Get the list of tool names
+        check_tool_names = [item["function"]["name"] for item in tools]
+
+        # Prepare a list of tool call ids and related function names
+        tool_call_ids = {}
+
+        # Loop through the messages and check if the tools have been run, removing them as we go
+        for message in messages:
+            if "tool_calls" in message:
+                # Register the tool ids and the function names (there could be multiple tool calls)
+                for tool_call in message["tool_calls"]:
+                    tool_call_ids[tool_call["id"]] = tool_call["function"]["name"]
+            elif "tool_call_id" in message:
+                # Tool called, get the name of the function based on the id
+                tool_name_called = tool_call_ids[message["tool_call_id"]]
+
+                # If we had not yet called the tool, check and remove it to indicate we have
+                if tool_name_called in check_tool_names:
+                    check_tool_names.remove(tool_name_called)
+
+        # Return True if all tools have been called at least once (accounted for)
+        return len(check_tool_names) == 0
+    else:
+        raise TypeError(
+            f"hide_tools_param is not a valid value ['if_all_run','if_any_run','never'], got '{hide_tools_param}'"
+        )
+
+
+# Logging format (originally from FLAML)
+logging_formatter = logging.Formatter(
+    "[%(name)s: %(asctime)s] {%(lineno)d} %(levelname)s - %(message)s", "%m-%d %H:%M:%S"
+)

autogen/oai/cohere.py

@@ -0,0 +1,494 @@
+# 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 Cohere's API.
+
+Example:
+    ```python
+    llm_config={
+        "config_list": [{
+            "api_type": "cohere",
+            "model": "command-r-plus",
+            "api_key": os.environ.get("COHERE_API_KEY")
+            "client_name": "autogen-cohere", # Optional parameter
+            }
+    ]}
+
+    agent = autogen.AssistantAgent("my_agent", llm_config=llm_config)
+    ```
+
+Install Cohere's python library using: pip install --upgrade cohere
+
+Resources:
+- https://docs.cohere.com/reference/chat
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import sys
+import time
+import warnings
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+from typing_extensions import Unpack
+
+from autogen.oai.client_utils import FormatterProtocol, logging_formatter, validate_parameter
+
+from ..import_utils import optional_import_block, require_optional_import
+from ..llm_config.entry import LLMConfigEntry, LLMConfigEntryDict
+from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
+
+with optional_import_block():
+    from cohere import ClientV2 as CohereV2
+    from cohere.types import ToolResult
+
+logger = logging.getLogger(__name__)
+if not logger.handlers:
+    # Add the console handler.
+    _ch = logging.StreamHandler(stream=sys.stdout)
+    _ch.setFormatter(logging_formatter)
+    logger.addHandler(_ch)
+
+
+COHERE_PRICING_1K = {
+    "command-r-plus": (0.003, 0.015),
+    "command-r": (0.0005, 0.0015),
+    "command-nightly": (0.00025, 0.00125),
+    "command": (0.015, 0.075),
+    "command-light": (0.008, 0.024),
+    "command-light-nightly": (0.008, 0.024),
+}
+
+
+class CohereEntryDict(LLMConfigEntryDict, total=False):
+    api_type: Literal["cohere"]
+
+    k: int
+    seed: int | None
+    frequency_penalty: float
+    presence_penalty: float
+    client_name: str | None
+    strict_tools: bool
+    stream: bool
+    tool_choice: Literal["NONE", "REQUIRED"] | None
+
+
+class CohereLLMConfigEntry(LLMConfigEntry):
+    api_type: Literal["cohere"] = "cohere"
+
+    k: int = Field(default=0, ge=0, le=500)
+    seed: int | None = None
+    frequency_penalty: float = Field(default=0, ge=0, le=1)
+    presence_penalty: float = Field(default=0, ge=0, le=1)
+    client_name: str | None = None
+    strict_tools: bool = False
+    stream: bool = False
+    tool_choice: Literal["NONE", "REQUIRED"] | None = None
+
+    def create_client(self):
+        raise NotImplementedError("CohereLLMConfigEntry.create_client is not implemented.")
+
+
+class CohereClient:
+    """Client for Cohere's API."""
+
+    RESPONSE_USAGE_KEYS: list[str] = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]
+
+    def __init__(self, **kwargs: Unpack[CohereEntryDict]):
+        """Requires api_key or environment variable to be set
+
+        Args:
+            **kwargs: The keyword arguments to pass to the Cohere API.
+        """
+        # Ensure we have the api_key upon instantiation
+        self.api_key = kwargs.get("api_key")
+        if not self.api_key:
+            self.api_key = os.getenv("COHERE_API_KEY")
+
+        assert self.api_key, (
+            "Please include the api_key in your config list entry for Cohere or set the COHERE_API_KEY env variable."
+        )
+
+        # Store the response format, if provided (for structured outputs)
+        self._response_format: type[BaseModel] | None = None
+
+    def message_retrieval(self, response) -> 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) -> float:
+        return response.cost
+
+    @staticmethod
+    def get_usage(response) -> 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 Cohere API from the passed in parameters and returns a validated set. Checks types, ranges, and sets defaults"""
+        cohere_params = {}
+
+        # Check that we have what we need to use Cohere's API
+        # We won't enforce the available models as they are likely to change
+        cohere_params["model"] = params.get("model")
+        assert cohere_params["model"], (
+            "Please specify the 'model' in your config list entry to nominate the Cohere model to use."
+        )
+
+        # Handle structured output response format from Pydantic model
+        if "response_format" in params and params["response_format"] is not None:
+            self._response_format = params.get("response_format")
+
+            response_format = params["response_format"]
+
+            # Check if it's a Pydantic model
+            if hasattr(response_format, "model_json_schema"):
+                # Get the JSON schema from the Pydantic model
+                schema = response_format.model_json_schema()
+
+                def resolve_ref(ref: str, defs: dict) -> dict:
+                    """Resolve a $ref to its actual schema definition"""
+                    # Extract the definition name from "#/$defs/Name"
+                    def_name = ref.split("/")[-1]
+                    return defs[def_name]
+
+                def ensure_type_fields(obj: dict, defs: dict) -> dict:
+                    """Recursively ensure all objects in the schema have a type and properties field"""
+                    if isinstance(obj, dict):
+                        # If it has a $ref, replace it with the actual definition
+                        if "$ref" in obj:
+                            ref_def = resolve_ref(obj["$ref"], defs)
+                            # Merge the reference definition with any existing fields
+                            obj = {**ref_def, **obj}
+                            # Remove the $ref as we've replaced it
+                            del obj["$ref"]
+
+                        # Process each value recursively
+                        return {
+                            k: ensure_type_fields(v, defs) if isinstance(v, (dict, list)) else v for k, v in obj.items()
+                        }
+                    elif isinstance(obj, list):
+                        return [ensure_type_fields(item, defs) for item in obj]
+                    return obj
+
+                # Make a copy of $defs before processing
+                defs = schema.get("$defs", {})
+
+                # Process the schema
+                processed_schema = ensure_type_fields(schema, defs)
+
+                cohere_params["response_format"] = {"type": "json_object", "json_schema": processed_schema}
+            else:
+                raise ValueError("response_format must be a Pydantic BaseModel")
+
+        # Handle strict tools parameter for structured outputs with tools
+        if "tools" in params:
+            cohere_params["strict_tools"] = validate_parameter(params, "strict_tools", bool, False, False, None, None)
+
+        # Validate allowed Cohere parameters
+        # https://docs.cohere.com/reference/chat
+        if "temperature" in params:
+            cohere_params["temperature"] = validate_parameter(
+                params, "temperature", (int, float), False, 0.3, (0, None), None
+            )
+
+        if "max_tokens" in params:
+            cohere_params["max_tokens"] = validate_parameter(params, "max_tokens", int, True, None, (0, None), None)
+
+        if "k" in params:
+            cohere_params["k"] = validate_parameter(params, "k", int, False, 0, (0, 500), None)
+
+        if "top_p" in params:
+            cohere_params["p"] = validate_parameter(params, "top_p", (int, float), False, 0.75, (0.01, 0.99), None)
+
+        if "seed" in params:
+            cohere_params["seed"] = validate_parameter(params, "seed", int, True, None, None, None)
+
+        if "frequency_penalty" in params:
+            cohere_params["frequency_penalty"] = validate_parameter(
+                params, "frequency_penalty", (int, float), True, 0, (0, 1), None
+            )
+
+        if "presence_penalty" in params:
+            cohere_params["presence_penalty"] = validate_parameter(
+                params, "presence_penalty", (int, float), True, 0, (0, 1), None
+            )
+
+        if "tool_choice" in params:
+            cohere_params["tool_choice"] = validate_parameter(
+                params, "tool_choice", str, True, None, None, ["NONE", "REQUIRED"]
+            )
+
+        return cohere_params
+
+    @require_optional_import("cohere", "cohere")
+    def create(self, params: dict) -> ChatCompletion:
+        messages = params.get("messages", [])
+        client_name = params.get("client_name") or "AG2"
+        cohere_tool_names = set()
+        tool_calls_modified_ids = set()
+
+        # Parse parameters to the Cohere API's parameters
+        cohere_params = self.parse_params(params)
+
+        cohere_params["messages"] = messages
+
+        if "tools" in params:
+            cohere_tool_names = {tool["function"]["name"] for tool in params["tools"]}
+            cohere_params["tools"] = params["tools"]
+
+        # Strip out name
+        for message in cohere_params["messages"]:
+            message_name = message.pop("name", "")
+            # Extract and prepend name to content or tool_plan if available
+            message["content"] = (
+                f"{message_name}: {(message.get('content') or message.get('tool_plan'))}"
+                if message_name
+                else (message.get("content") or message.get("tool_plan"))
+            )
+
+            # Handle tool calls
+            if message.get("tool_calls") is not None and len(message["tool_calls"]) > 0:
+                message["tool_plan"] = message.get("tool_plan", message["content"])
+                del message["content"]  # Remove content as tool_plan is prioritized
+
+                # If tool call name is missing or not recognized, modify role and content
+                for tool_call in message["tool_calls"] or []:
+                    if (not tool_call.get("function", {}).get("name")) or tool_call.get("function", {}).get(
+                        "name"
+                    ) not in cohere_tool_names:
+                        message["role"] = "assistant"
+                        message["content"] = f"{message.pop('tool_plan', '')}{str(message['tool_calls'])}"
+                        tool_calls_modified_ids = tool_calls_modified_ids.union({
+                            tool_call.get("id") for tool_call in message["tool_calls"]
+                        })
+                        del message["tool_calls"]
+                        break
+
+            # Adjust role if message comes from a tool with a modified ID
+            if message.get("role") == "tool":
+                tool_id = message.get("tool_call_id")
+                if tool_id in tool_calls_modified_ids:
+                    message["role"] = "user"
+                    del message["tool_call_id"]  # Remove the tool call ID
+
+        # We use chat model by default
+        client = CohereV2(api_key=self.api_key, client_name=client_name)
+
+        # Token counts will be returned
+        prompt_tokens = 0
+        completion_tokens = 0
+        total_tokens = 0
+
+        # Stream if in parameters
+        streaming = params.get("stream")
+        cohere_finish = "stop"
+        tool_calls = None
+        ans = None
+        if streaming:
+            response = client.chat_stream(**cohere_params)
+            # Streaming...
+            ans = ""
+            plan = ""
+            prompt_tokens = 0
+            completion_tokens = 0
+            for chunk in response:
+                if chunk.type == "content-delta":
+                    ans = ans + chunk.delta.message.content.text
+                elif chunk.type == "tool-plan-delta":
+                    plan = plan + chunk.delta.message.tool_plan
+                elif chunk.type == "tool-call-start":
+                    cohere_finish = "tool_calls"
+
+                    # Initialize a new tool call
+                    tool_call = chunk.delta.message.tool_calls
+                    current_tool = {
+                        "id": tool_call.id,
+                        "type": "function",
+                        "function": {"name": tool_call.function.name, "arguments": ""},
+                    }
+                elif chunk.type == "tool-call-delta":
+                    # Progressively build the arguments as they stream in
+                    if current_tool is not None:
+                        current_tool["function"]["arguments"] += chunk.delta.message.tool_calls.function.arguments
+                elif chunk.type == "tool-call-end":
+                    # Append the finished tool call to the list
+                    if current_tool is not None:
+                        if tool_calls is None:
+                            tool_calls = []
+                        tool_calls.append(ChatCompletionMessageToolCall(**current_tool))
+                        current_tool = None
+                elif chunk.type == "message-start":
+                    response_id = chunk.id
+                elif chunk.type == "message-end":
+                    prompt_tokens = (
+                        chunk.delta.usage.billed_units.input_tokens
+                    )  # Note total (billed+non-billed) available with ...usage.tokens...
+                    completion_tokens = chunk.delta.usage.billed_units.output_tokens
+
+            total_tokens = prompt_tokens + completion_tokens
+        else:
+            response = client.chat(**cohere_params)
+
+            if response.message.tool_calls is not None:
+                ans = response.message.tool_plan
+                cohere_finish = "tool_calls"
+                tool_calls = []
+                for tool_call in response.message.tool_calls:
+                    # if parameters are null, clear them out (Cohere can return a string "null" if no parameter values)
+
+                    tool_calls.append(
+                        ChatCompletionMessageToolCall(
+                            id=tool_call.id,
+                            function={
+                                "name": tool_call.function.name,
+                                "arguments": (
+                                    "" if tool_call.function.arguments is None else tool_call.function.arguments
+                                ),
+                            },
+                            type="function",
+                        )
+                    )
+            else:
+                ans: str = response.message.content[0].text
+
+            # Not using billed_units, but that may be better for cost purposes
+            prompt_tokens = (
+                response.usage.billed_units.input_tokens
+            )  # Note total (billed+non-billed) available with ...usage.tokens...
+            completion_tokens = response.usage.billed_units.output_tokens
+            total_tokens = prompt_tokens + completion_tokens
+
+            response_id = response.id
+
+        # Clean up structured output if needed
+        if self._response_format:
+            # ans = clean_return_response_format(ans)
+            try:
+                parsed_response = self._convert_json_response(ans)
+                ans = _format_json_response(parsed_response, ans)
+            except ValueError as e:
+                ans = str(e)
+
+        # 3. convert output
+        message = ChatCompletionMessage(
+            role="assistant",
+            content=ans,
+            function_call=None,
+            tool_calls=tool_calls,
+        )
+        choices = [Choice(finish_reason=cohere_finish, index=0, message=message)]
+
+        response_oai = ChatCompletion(
+            id=response_id,
+            model=cohere_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,
+            ),
+            cost=calculate_cohere_cost(prompt_tokens, completion_tokens, cohere_params["model"]),
+        )
+
+        return response_oai
+
+    def _convert_json_response(self, response: str) -> Any:
+        """Extract and validate JSON response from the output for structured outputs.
+
+        Args:
+            response (str): The response from the API.
+
+        Returns:
+            Any: The parsed JSON response.
+        """
+        if not self._response_format:
+            return response
+
+        try:
+            # Parse JSON and validate against the Pydantic model
+            json_data = json.loads(response)
+            return self._response_format.model_validate(json_data)
+        except Exception as e:
+            raise ValueError(
+                f"Failed to parse response as valid JSON matching the schema for Structured Output: {str(e)}"
+            )
+
+
+def _format_json_response(response: Any, original_answer: str) -> str:
+    """Formats the JSON response for structured outputs using the format method if it exists."""
+    return (
+        response.format() if isinstance(response, FormatterProtocol) else clean_return_response_format(original_answer)
+    )
+
+
+def extract_to_cohere_tool_results(tool_call_id: str, content_output: str, all_tool_calls) -> list[dict[str, Any]]:
+    temp_tool_results = []
+
+    for tool_call in all_tool_calls:
+        if tool_call["id"] == tool_call_id:
+            call = {
+                "name": tool_call["function"]["name"],
+                "parameters": json.loads(
+                    tool_call["function"]["arguments"] if tool_call["function"]["arguments"] != "" else "{}"
+                ),
+            }
+            output = [{"value": content_output}]
+            temp_tool_results.append(ToolResult(call=call, outputs=output))
+    return temp_tool_results
+
+
+def calculate_cohere_cost(input_tokens: int, output_tokens: int, model: str) -> float:
+    """Calculate the cost of the completion using the Cohere pricing."""
+    total = 0.0
+
+    if model in COHERE_PRICING_1K:
+        input_cost_per_k, output_cost_per_k = COHERE_PRICING_1K[model]
+        input_cost = (input_tokens / 1000) * input_cost_per_k
+        output_cost = (output_tokens / 1000) * output_cost_per_k
+        total = input_cost + output_cost
+    else:
+        warnings.warn(f"Cost calculation not available for {model} model", UserWarning)
+
+    return total
+
+
+def clean_return_response_format(response_str: str) -> str:
+    """Clean up the response string by parsing through json library."""
+    # Parse the string to a JSON object to handle escapes
+    data = json.loads(response_str)
+
+    # Convert back to JSON string with minimal formatting
+    return json.dumps(data)
+
+
+class CohereError(Exception):
+    """Base class for other Cohere exceptions"""
+
+    pass
+
+
+class CohereRateLimitError(CohereError):
+    """Raised when rate limit is exceeded"""
+
+    pass