ToolCalls

    tool_response: ContentItem = None

    @property
    def is_response_set(self) -> bool:
        """Check if the tool response is set."""
        return self.tool_response is not None

    def response_text(self, text: str) -> "MessageBuilder":
        self.tool_response = {"text": text}
        return self

    def response_image(self, image: str) -> "MessageBuilder":
        self.tool_response = {"image": image}
        return self

    def __repr__(self):
        return f"ToolCall(name={self.name}, arguments={self.arguments})"

@dataclass
class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self

    @property
    def all_responses_set(self) -> bool:
        """Check if all tool calls have responses set."""
        return all(call.is_response_set for call in self.tool_calls)

    def __len__(self) -> int:
        """Return the number of tool calls."""
        return len(self.tool_calls)

    def __iter__(self):
        """Make ToolCalls iterable."""
        return iter(self.tool_calls)

    def __bool__(self):
        """Check if there are any tool calls."""
        return len(self.tool_calls) > 0

@dataclass
class LLMOutput:
    """Serializable object for the output of a response LLM."""

    raw_response: Any = field(default=None)
    think: str = field(default="")
    action: str | None = field(default=None)  # Default action if no tool call is made
    tool_calls: ToolCalls | None = field(
        default=None
    )  # This will hold the tool call response if any

class MessageBuilder:
    def __init__(self, role: str):
        self.role = role
        self.content: List[ContentItem] = []
        self.responded_tool_calls: ToolCalls = None


LLMOutput

    @property
    def all_responses_set(self) -> bool:
        """Check if all tool calls have responses set."""
        return all(call.is_response_set for call in self.tool_calls)

    def __len__(self) -> int:
        """Return the number of tool calls."""
        return len(self.tool_calls)

    def __iter__(self):
        """Make ToolCalls iterable."""
        return iter(self.tool_calls)

    def __bool__(self):
        """Check if there are any tool calls."""
        return len(self.tool_calls) > 0

@dataclass
class LLMOutput:
    """Serializable object for the output of a response LLM."""

    raw_response: Any = field(default=None)
    think: str = field(default="")
    action: str | None = field(default=None)  # Default action if no tool call is made
    tool_calls: ToolCalls | None = field(
        default=None
    )  # This will hold the tool call response if any

class MessageBuilder:
    def __init__(self, role: str):
        self.role = role
        self.content: List[ContentItem] = []
        self.responded_tool_calls: ToolCalls = None

    @classmethod
    def system(cls) -> "MessageBuilder":
        return cls("system")

    @classmethod
    def user(cls) -> "MessageBuilder":
        return cls("user")

    @classmethod
    def assistant(cls) -> "MessageBuilder":
        return cls("assistant")


MessageBuilder

        """Make ToolCalls iterable."""
        return iter(self.tool_calls)

    def __bool__(self):
        """Check if there are any tool calls."""
        return len(self.tool_calls) > 0

@dataclass
class LLMOutput:
    """Serializable object for the output of a response LLM."""

    raw_response: Any = field(default=None)
    think: str = field(default="")
    action: str | None = field(default=None)  # Default action if no tool call is made
    tool_calls: ToolCalls | None = field(
        default=None
    )  # This will hold the tool call response if any

class MessageBuilder:
    def __init__(self, role: str):
        self.role = role
        self.content: List[ContentItem] = []
        self.responded_tool_calls: ToolCalls = None

    @classmethod
    def system(cls) -> "MessageBuilder":
        return cls("system")

    @classmethod
    def user(cls) -> "MessageBuilder":
        return cls("user")

    @classmethod
    def assistant(cls) -> "MessageBuilder":
        return cls("assistant")

    @abstractmethod
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the API call."""
        raise NotImplementedError("Subclasses must implement this method.")

    def add_text(self, text: str) -> "MessageBuilder":
        self.content.append({"text": text})
        return self

    def add_image(self, image: str) -> "MessageBuilder":
        self.content.append({"image": image})
        return self

    def to_markdown(self) -> str:
        parts = []
        for item in self.content:
            if "text" in item:
                parts.append(f"\n```\n{item['text']}\n```\n")
            elif "image" in item:
                parts.append(f"![Image]({item['image']})")

        # Tool call markdown repr
        if self.responded_tool_calls is not None:
            for i, tool_call in enumerate(self.responded_tool_calls.tool_calls, 1):
                parts.append(
                    f"\n**Tool Call {i}**: {tool_call_to_python_code(tool_call.name, tool_call.arguments)}"
                )
                response = tool_call.tool_response
                if response is not None:
                    parts.append(f"\n**Tool Response {i}:**")
                    content = (
                        f"```\n{response['text']}\n```"
                        if "text" in response
                        else f"![Tool Response Image]({response['image']})"
                    )
                    parts.append(content)

        markdown = f"### {self.role.capitalize()}\n"
        markdown += "\n".join(parts)

        return markdown

    def add_image_url(self, image_url: str) -> "MessageBuilder":
        """Add an image URL to the message content."""
        self.content.append({"image": image_to_png_base64_url(image_url)})
        return self

    def mark_all_previous_msg_for_caching(self):
        """Insert a cache breakpoint in the message content."""
        # This is a placeholder for future implementation.
        raise NotImplementedError

    @classmethod
    def add_responded_tool_calls(cls, responded_tool_calls: ToolCalls) -> "MessageBuilder":
        """Add tool calls to the message content."""
        assert responded_tool_calls.all_responses_set, "All tool calls must have a response."
        msg = cls("tool")
        msg.responded_tool_calls = responded_tool_calls
        return msg

class OpenAIResponseAPIMessageBuilder(MessageBuilder):
    @classmethod
    def system(cls) -> "OpenAIResponseAPIMessageBuilder":
        # OpenAI Responses API uses 'developer' role for system messages
        return cls("developer")

    def prepare_message(self) -> List[Message]:
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]

        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Convert the content item to the expected format for OpenAI Responses."""
        if "text" in content:
            content_type = "input_text" if self.role != "assistant" else "output_text"


OpenAIResponseAPIMessageBuilder

    def add_image_url(self, image_url: str) -> "MessageBuilder":
        """Add an image URL to the message content."""
        self.content.append({"image": image_to_png_base64_url(image_url)})
        return self

    def mark_all_previous_msg_for_caching(self):
        """Insert a cache breakpoint in the message content."""
        # This is a placeholder for future implementation.
        raise NotImplementedError

    @classmethod
    def add_responded_tool_calls(cls, responded_tool_calls: ToolCalls) -> "MessageBuilder":
        """Add tool calls to the message content."""
        assert responded_tool_calls.all_responses_set, "All tool calls must have a response."
        msg = cls("tool")
        msg.responded_tool_calls = responded_tool_calls
        return msg

class OpenAIResponseAPIMessageBuilder(MessageBuilder):
    @classmethod
    def system(cls) -> "OpenAIResponseAPIMessageBuilder":
        # OpenAI Responses API uses 'developer' role for system messages
        return cls("developer")

    def prepare_message(self) -> List[Message]:
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]

        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Convert the content item to the expected format for OpenAI Responses."""
        if "text" in content:
            content_type = "input_text" if self.role != "assistant" else "output_text"
            return {"type": content_type, "text": content["text"]}
        elif "image" in content:
            return {"type": "input_image", "image_url": content["image"]}
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:
            raise ValueError("No tool calls found in responded_tool_calls")

        output = []
        output.extend(self.responded_tool_calls.raw_calls.output)  # this contains response
        for fn_call in self.responded_tool_calls:
            call_type = fn_call.raw_call.type
            call_id = fn_call.raw_call.call_id
            call_response = fn_call.tool_response

            match call_type:
                case "function_call":
                    # image output is not supported in function calls response.
                    assert (
                        "image" not in call_response
                    ), "Image output is not supported in function calls response."
                    fn_call_response = {
                        "type": "function_call_output",
                        "call_id": call_id,
                        "output": self.convert_content_to_expected_format(call_response)["text"],
                    }
                    output.append(fn_call_response)

                case "computer_call":
                    # For computer calls, use only images are expected.
                    assert (
                        "text" not in call_response
                    ), "Text output is not supported in computer calls response."
                    computer_call_output = {
                        "type": "computer_call_output",
                        "call_id": call_id,
                        "output": self.convert_content_to_expected_format(call_response),
                    }
                    output.append(computer_call_output)  # this needs to be a screenshot

        return output

    def mark_all_previous_msg_for_caching(self):
        """Nothing special to do here for openAI. They do not have a notion of cache breakpoints."""
        pass

class AnthropicAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        content = [self.transform_content(item) for item in self.content]
        output = {"role": self.role, "content": content}

        if self.role == "tool":
            return self.handle_tool_call()

        if self.role == "assistant":
            # Strip whitespace from assistant text responses. See anthropic error code 400.
            for c in output["content"]:
                if "text" in c:
                    c["text"] = c["text"].strip()
        return [output]

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:


AnthropicAPIMessageBuilder

                case "computer_call":
                    # For computer calls, use only images are expected.
                    assert (
                        "text" not in call_response
                    ), "Text output is not supported in computer calls response."
                    computer_call_output = {
                        "type": "computer_call_output",
                        "call_id": call_id,
                        "output": self.convert_content_to_expected_format(call_response),
                    }
                    output.append(computer_call_output)  # this needs to be a screenshot

        return output

    def mark_all_previous_msg_for_caching(self):
        """Nothing special to do here for openAI. They do not have a notion of cache breakpoints."""
        pass

class AnthropicAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        content = [self.transform_content(item) for item in self.content]
        output = {"role": self.role, "content": content}

        if self.role == "tool":
            return self.handle_tool_call()

        if self.role == "assistant":
            # Strip whitespace from assistant text responses. See anthropic error code 400.
            for c in output["content"]:
                if "text" in c:
                    c["text"] = c["text"].strip()
        return [output]

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:
            raise ValueError("No tool calls found in responded_tool_calls")

        llm_tool_call = {
            "role": "assistant",
            "content": self.responded_tool_calls.raw_calls.content,
        }  # Add the toolcall block
        tool_response = {"role": "user", "content": []}  # Anthropic expects a list of messages
        for call in self.responded_tool_calls:
            assert (
                "image" not in call.tool_response
            ), "Image output is not supported in tool calls response."
            tool_response["content"].append(
                {
                    "type": "tool_result",
                    "tool_use_id": call.raw_call.id,
                    "content": self.transform_content(call.tool_response)[
                        "text"
                    ],  # needs to be str
                }
            )

        return [llm_tool_call, tool_response]

    def transform_content(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by Anthropic API."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content:
            img_str: str = content["image"]
            # make sure to get rid of the image type for anthropic
            # e.g. "data:image/png;base64"
            if img_str.startswith("data:image/png;base64,"):
                img_str = img_str[len("data:image/png;base64,") :]
            return {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png",
                    "data": img_str,
                },
            }
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def mark_all_previous_msg_for_caching(self) -> List[Message]:
        """Insert a cache breakpoint in the message content to mark all previous messages for caching."""
        self._cache_breakpoint = True

class OpenAIChatCompletionAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the OpenAI API."""
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]
        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by OpenAI ChatCompletion."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content:
            return {"type": "image_url", "image_url": {"url": content["image"]}}
        else:
            raise ValueError(f"Unsupported content type: {content}")


OpenAIChatCompletionAPIMessageBuilder

            # make sure to get rid of the image type for anthropic
            # e.g. "data:image/png;base64"
            if img_str.startswith("data:image/png;base64,"):
                img_str = img_str[len("data:image/png;base64,") :]
            return {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png",
                    "data": img_str,
                },
            }
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def mark_all_previous_msg_for_caching(self) -> List[Message]:
        """Insert a cache breakpoint in the message content to mark all previous messages for caching."""
        self._cache_breakpoint = True

class OpenAIChatCompletionAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the OpenAI API."""
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]
        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by OpenAI ChatCompletion."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content:
            return {"type": "image_url", "image_url": {"url": content["image"]}}
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:
            raise ValueError("No tool calls found in responded_tool_calls")
        output = []
        output.append(
            self.responded_tool_calls.raw_calls.choices[0].message
        )  # add raw calls to output
        for fn_call in self.responded_tool_calls:
            raw_call = fn_call.raw_call
            assert (
                "image" not in fn_call.tool_response
            ), "Image output is not supported in function calls response."
            # a function_call_output dict has keys "role", "tool_call_id" and "content"
            tool_call_reponse = {
                "name": raw_call["function"]["name"],  # required with OpenRouter
                "role": "tool",
                "tool_call_id": raw_call["id"],
                "content": self.convert_content_to_expected_format(fn_call.tool_response)["text"],
            }
            output.append(tool_call_reponse)

        return output

    def mark_all_previous_msg_for_caching(self):
        """Nothing special to do here for openAI. They do not have a notion of cache breakpoints."""
        pass

@dataclass
class APIPayload:
    messages: List[MessageBuilder] | None = None
    tools: List[Dict[str, Any]] | None = None
    tool_choice: Literal["none", "auto", "any", "required"] | None = None
    force_call_tool: str | None = (
        None  # Name of the tool to call # If set, will force the LLM to call this tool.
    )
    use_cache_breakpoints: bool = (
        False  # If True, will apply cache breakpoints to the messages. # applicable for Anthropic
    )
    cache_tool_definition: bool = (
        False  # If True, will cache the tool definition in the last message.
    )
    cache_complete_prompt: bool = (
        False  # If True, will cache the complete prompt in the last message.
    )
    reasoning_effort: Literal["low", "medium", "high"] | None = None


APIPayload

            assert (
                "image" not in fn_call.tool_response
            ), "Image output is not supported in function calls response."
            # a function_call_output dict has keys "role", "tool_call_id" and "content"
            tool_call_reponse = {
                "name": raw_call["function"]["name"],  # required with OpenRouter
                "role": "tool",
                "tool_call_id": raw_call["id"],
                "content": self.convert_content_to_expected_format(fn_call.tool_response)["text"],
            }
            output.append(tool_call_reponse)

        return output

    def mark_all_previous_msg_for_caching(self):
        """Nothing special to do here for openAI. They do not have a notion of cache breakpoints."""
        pass

@dataclass
class APIPayload:
    messages: List[MessageBuilder] | None = None
    tools: List[Dict[str, Any]] | None = None
    tool_choice: Literal["none", "auto", "any", "required"] | None = None
    force_call_tool: str | None = (
        None  # Name of the tool to call # If set, will force the LLM to call this tool.
    )
    use_cache_breakpoints: bool = (
        False  # If True, will apply cache breakpoints to the messages. # applicable for Anthropic
    )
    cache_tool_definition: bool = (
        False  # If True, will cache the tool definition in the last message.
    )
    cache_complete_prompt: bool = (
        False  # If True, will cache the complete prompt in the last message.
    )
    reasoning_effort: Literal["low", "medium", "high"] | None = None

    def __post_init__(self):
        if self.tool_choice and self.force_call_tool:
            raise ValueError("tool_choice and force_call_tool are mutually exclusive")
        if self.reasoning_effort is not None:
            logger.info(
                "In agentlab reasoning_effort is used by LiteLLM API only. We will eventually shift to LiteLLM API for all LLMs."
            )

# # Base class for all API Endpoints
class BaseResponseModel(ABC):
    def __init__(
        self,
        model_name: str,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ):
        self.model_name = model_name
        self.api_key = api_key
        self.temperature = temperature
        self.max_tokens = max_tokens
        super().__init__()

    def __call__(self, payload: APIPayload) -> LLMOutput:
        """Make a call to the model and return the parsed response."""
        response = self._call_api(payload)


BaseResponseModel

        False  # If True, will apply cache breakpoints to the messages. # applicable for Anthropic
    )
    cache_tool_definition: bool = (
        False  # If True, will cache the tool definition in the last message.
    )
    cache_complete_prompt: bool = (
        False  # If True, will cache the complete prompt in the last message.
    )
    reasoning_effort: Literal["low", "medium", "high"] | None = None

    def __post_init__(self):
        if self.tool_choice and self.force_call_tool:
            raise ValueError("tool_choice and force_call_tool are mutually exclusive")
        if self.reasoning_effort is not None:
            logger.info(
                "In agentlab reasoning_effort is used by LiteLLM API only. We will eventually shift to LiteLLM API for all LLMs."
            )

# # Base class for all API Endpoints
class BaseResponseModel(ABC):
    def __init__(
        self,
        model_name: str,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ):
        self.model_name = model_name
        self.api_key = api_key
        self.temperature = temperature
        self.max_tokens = max_tokens
        super().__init__()

    def __call__(self, payload: APIPayload) -> LLMOutput:
        """Make a call to the model and return the parsed response."""
        response = self._call_api(payload)
        return self._parse_response(response)

    @abstractmethod
    def _call_api(self, payload: APIPayload) -> Any:
        """Make a call to the model API and return the raw response."""
        pass

    @abstractmethod
    def _parse_response(self, response: Any) -> LLMOutput:
        """Parse the raw response from the model API and return a structured response."""
        pass

class AgentlabAction:
    """
    Collection of utility function to convert tool calls to Agentlab action format.
    """

    @staticmethod
    def convert_toolcall_to_agentlab_action_format(toolcall: ToolCall) -> str:
        """Convert a tool call to an Agentlab environment action string.

        Args:
            toolcall: ToolCall object containing the name and arguments of the tool call.

        Returns:
            A string representing the action in Agentlab format i.e. python function call string.
        """

        tool_name, tool_args = toolcall.name, toolcall.arguments
        return tool_call_to_python_code(tool_name, tool_args)


AgentlabAction

        self.temperature = temperature
        self.max_tokens = max_tokens
        super().__init__()

    def __call__(self, payload: APIPayload) -> LLMOutput:
        """Make a call to the model and return the parsed response."""
        response = self._call_api(payload)
        return self._parse_response(response)

    @abstractmethod
    def _call_api(self, payload: APIPayload) -> Any:
        """Make a call to the model API and return the raw response."""
        pass

    @abstractmethod
    def _parse_response(self, response: Any) -> LLMOutput:
        """Parse the raw response from the model API and return a structured response."""
        pass

class AgentlabAction:
    """
    Collection of utility function to convert tool calls to Agentlab action format.
    """

    @staticmethod
    def convert_toolcall_to_agentlab_action_format(toolcall: ToolCall) -> str:
        """Convert a tool call to an Agentlab environment action string.

        Args:
            toolcall: ToolCall object containing the name and arguments of the tool call.

        Returns:
            A string representing the action in Agentlab format i.e. python function call string.
        """

        tool_name, tool_args = toolcall.name, toolcall.arguments
        return tool_call_to_python_code(tool_name, tool_args)

    @staticmethod
    def convert_multiactions_to_agentlab_action_format(actions: list[str]) -> str | None:
        """Convert multiple actions list to a format that env supports.

        Args:
            actions: List of action strings to be joined.

        Returns:
            Joined actions separated by newlines, or None if empty.
        """
        return "\n".join(actions) if actions else None

class BaseModelWithPricing(TrackAPIPricingMixin, BaseResponseModel):
    pass

class OpenAIResponseModel(BaseModelWithPricing):
    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = 100,
    ):
        self.action_space_as_tools = True  # this should be a config
        super().__init__(  # This is passed to BaseModel
            model_name=model_name, api_key=api_key, temperature=temperature, max_tokens=max_tokens
        )
        client_args = {}


BaseModelWithPricing

        Returns:
            A string representing the action in Agentlab format i.e. python function call string.
        """

        tool_name, tool_args = toolcall.name, toolcall.arguments
        return tool_call_to_python_code(tool_name, tool_args)

    @staticmethod
    def convert_multiactions_to_agentlab_action_format(actions: list[str]) -> str | None:
        """Convert multiple actions list to a format that env supports.

        Args:
            actions: List of action strings to be joined.

        Returns:
            Joined actions separated by newlines, or None if empty.
        """
        return "\n".join(actions) if actions else None

class BaseModelWithPricing(TrackAPIPricingMixin, BaseResponseModel):
    pass

class OpenAIResponseModel(BaseModelWithPricing):
    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = 100,
    ):
        self.action_space_as_tools = True  # this should be a config
        super().__init__(  # This is passed to BaseModel
            model_name=model_name, api_key=api_key, temperature=temperature, max_tokens=max_tokens
        )
        client_args = {}
        if base_url is not None:
            client_args["base_url"] = base_url
        if api_key is not None:
            client_args["api_key"] = api_key


OpenAIResponseModel

        tool_name, tool_args = toolcall.name, toolcall.arguments
        return tool_call_to_python_code(tool_name, tool_args)

    @staticmethod
    def convert_multiactions_to_agentlab_action_format(actions: list[str]) -> str | None:
        """Convert multiple actions list to a format that env supports.

        Args:
            actions: List of action strings to be joined.

        Returns:
            Joined actions separated by newlines, or None if empty.
        """
        return "\n".join(actions) if actions else None

class BaseModelWithPricing(TrackAPIPricingMixin, BaseResponseModel):
    pass

class OpenAIResponseModel(BaseModelWithPricing):
    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = 100,
    ):
        self.action_space_as_tools = True  # this should be a config
        super().__init__(  # This is passed to BaseModel
            model_name=model_name, api_key=api_key, temperature=temperature, max_tokens=max_tokens
        )
        client_args = {}
        if base_url is not None:
            client_args["base_url"] = base_url
        if api_key is not None:
            client_args["api_key"] = api_key
        self.client = OpenAI(**client_args)
        # Init pricing tracker after super() so that all attributes have been set.
        self.init_pricing_tracker(pricing_api="openai")  # Use the PricingMixin

    def _call_api(self, payload: APIPayload) -> "OpenAIResponseObject":

        input = []
        for msg in payload.messages:
            input.extend(msg.prepare_message())
        api_params: Dict[str, Any] = {
            "model": self.model_name,
            "input": input,
        }
        # Not all Open AI models support these parameters (example: o3), so we check if they are set.
        if self.temperature is not None:
            api_params["temperature"] = self.temperature
        if self.max_tokens is not None:
            api_params["max_output_tokens"] = self.max_tokens
        if payload.tools is not None:
            api_params["tools"] = payload.tools
        if payload.tool_choice is not None and payload.force_call_tool is None:
            api_params["tool_choice"] = (
                "required" if payload.tool_choice in ("required", "any") else payload.tool_choice
            )
        if payload.force_call_tool is not None:
            api_params["tool_choice"] = {"type": "function", "name": payload.force_call_tool}

        response = call_openai_api_with_retries(
            self.client.responses.create,
            api_params,
        )

        return response

    def _parse_response(self, response: "OpenAIResponseObject") -> LLMOutput:
        """Parse the raw response from the OpenAI Responses API."""

        think_output = self._extract_thinking_content_from_response(response)
        toolcalls = self._extract_tool_calls_from_response(response)

        if self.action_space_as_tools:
            env_action = self._extract_env_actions_from_toolcalls(toolcalls)
        else:
            env_action = self._extract_env_actions_from_text_response(response)

        return LLMOutput(
            raw_response=response,
            think=think_output,
            action=env_action if env_action is not None else None,
            tool_calls=toolcalls if toolcalls is not None else None,
        )

    def _extract_tool_calls_from_response(self, response: "OpenAIResponseObject") -> ToolCalls:
        """Extracts tool calls from the response."""
        tool_calls = []
        for output in response.output:
            if output.type == "function_call":
                tool_name = output.name
                tool_args = json.loads(output.arguments)
            elif output.type == "computer_call":
                tool_name, tool_args = self.cua_action_to_env_tool_name_and_args(output.action)
            else:
                continue
            tool_calls.append(ToolCall(name=tool_name, arguments=tool_args, raw_call=output))

        return ToolCalls(tool_calls=tool_calls, raw_calls=response)

    def _extract_env_actions_from_toolcalls(self, toolcalls: ToolCalls) -> Any | None:
        """Extracts actions from the response."""
        if not toolcalls:
            return None

        actions = [
            AgentlabAction.convert_toolcall_to_agentlab_action_format(call) for call in toolcalls
        ]
        actions = (
            AgentlabAction.convert_multiactions_to_agentlab_action_format(actions)
            if len(actions) > 1
            else actions[0]
        )
        return actions

    def _extract_thinking_content_from_response(self, response: "OpenAIResponseObject") -> str:
        """Extracts the thinking content from the response."""
        thinking_content = ""
        for output in response.output:
            if output.type == "reasoning":
                if len(output.summary) > 0:
                    thinking_content += output.summary[0].text + "\n"
            elif output.type == "message" and output.content:
                thinking_content += output.content[0].text + "\n"
            elif hasattr(output, "output_text") and output.output_text:
                thinking_content += f"{output.output_text}\n"
        return thinking_content

    def cua_action_to_env_tool_name_and_args(self, action: str) -> tuple[str, Dict[str, Any]]:
        """ "Overwrite this method to convert a computer action to agentlab action string"""
        raise NotImplementedError(
            "This method should be implemented in the subclass to convert a computer action to agentlab action string."
        )

    def _extract_env_actions_from_text_response(
        self, response: "OpenAIResponseObject"
    ) -> str | None:
        """Extracts environment actions from the text response."""
        # Use when action space is not given as tools.
        pass

class OpenAIChatCompletionModel(BaseModelWithPricing):
    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = 100,
    ):
        super().__init__(
            model_name=model_name,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        self.action_space_as_tools = True  # this should be a config
        client_args = {}
        if base_url is not None:
            client_args["base_url"] = base_url


OpenAIChatCompletionModel

            elif output.type == "message" and output.content:
                thinking_content += output.content[0].text + "\n"
            elif hasattr(output, "output_text") and output.output_text:
                thinking_content += f"{output.output_text}\n"
        return thinking_content

    def cua_action_to_env_tool_name_and_args(self, action: str) -> tuple[str, Dict[str, Any]]:
        """ "Overwrite this method to convert a computer action to agentlab action string"""
        raise NotImplementedError(
            "This method should be implemented in the subclass to convert a computer action to agentlab action string."
        )

    def _extract_env_actions_from_text_response(
        self, response: "OpenAIResponseObject"
    ) -> str | None:
        """Extracts environment actions from the text response."""
        # Use when action space is not given as tools.
        pass

class OpenAIChatCompletionModel(BaseModelWithPricing):
    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = 100,
    ):
        super().__init__(
            model_name=model_name,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        self.action_space_as_tools = True  # this should be a config
        client_args = {}
        if base_url is not None:
            client_args["base_url"] = base_url
        if api_key is not None:
            client_args["api_key"] = api_key
        self.client = OpenAI(**client_args)
        self.init_pricing_tracker(pricing_api="openai")  # Use the PricingMixin

    def _call_api(self, payload: APIPayload) -> "openai.types.chat.ChatCompletion":
        input = []
        for msg in payload.messages:
            input.extend(msg.prepare_message())
        api_params: Dict[str, Any] = {
            "model": self.model_name,
            "messages": input,
        }
        if self.temperature is not None:
            api_params["temperature"] = self.temperature

        if self.max_tokens is not None:
            api_params["max_completion_tokens"] = self.max_tokens

        if payload.tools is not None:
            # tools format is OpenAI Response API format.
            api_params["tools"] = self.format_tools_for_chat_completion(payload.tools)

        if payload.tool_choice is not None and payload.force_call_tool is None:
            api_params["tool_choice"] = (
                "required" if payload.tool_choice in ("required", "any") else payload.tool_choice
            )

        if payload.force_call_tool is not None:
            api_params["tool_choice"] = {
                "type": "function",
                "function": {"name": payload.force_call_tool},
            }

        response = call_openai_api_with_retries(self.client.chat.completions.create, api_params)

        return response

    def _parse_response(self, response: "openai.types.chat.ChatCompletion") -> LLMOutput:
        think_output = self._extract_thinking_content_from_response(response)
        tool_calls = self._extract_tool_calls_from_response(response)

        if self.action_space_as_tools:
            env_action = self._extract_env_actions_from_toolcalls(tool_calls)
        else:
            env_action = self._extract_env_actions_from_text_response(response)
        return LLMOutput(
            raw_response=response,
            think=think_output,
            action=env_action if env_action is not None else None,
            tool_calls=tool_calls if tool_calls is not None else None,
        )

    def _extract_thinking_content_from_response(
        self, response: openai.types.chat.ChatCompletion, wrap_tag="think"
    ):
        """Extracts the content from the message, including reasoning if available.
        It wraps the reasoning around <think>...</think> for easy identification of reasoning content,
        When LLM produces 'text' and 'reasoning' in the same message.
        Note: The wrapping of 'thinking' content may not be nedeed and may be reconsidered.

        Args:
            response: The message object or dict containing content and reasoning.
            wrap_tag: The tag name to wrap reasoning content (default: "think").

        Returns:
            str: The extracted content with reasoning wrapped in specified tags.
        """
        message = response.choices[0].message
        if not isinstance(message, dict):
            message = message.to_dict()

        reasoning_content = message.get("reasoning", None)
        msg_content = message.get("text", "")  # works for Open-router
        if reasoning_content:
            # Wrap reasoning in <think> tags with newlines for clarity
            reasoning_content = f"<{wrap_tag}>{reasoning_content}</{wrap_tag}>\n"
            logging.debug("Extracting content from response.choices[i].message.reasoning")
        else:
            reasoning_content = ""
        return f"{reasoning_content}{msg_content}{message.get('content', '')}"

    def _extract_tool_calls_from_response(
        self, response: openai.types.chat.ChatCompletion
    ) -> ToolCalls | None:
        """Extracts tool calls from the response."""
        message = response.choices[0].message.to_dict()
        tool_calls = message.get("tool_calls", None)
        if tool_calls is None:
            return None
        tool_call_list = []
        for tc in tool_calls:
            tool_call_list.append(
                ToolCall(
                    name=tc["function"]["name"],
                    arguments=json.loads(tc["function"]["arguments"]),
                    raw_call=tc,
                )
            )
        return ToolCalls(tool_calls=tool_call_list, raw_calls=response)

    def _extract_env_actions_from_toolcalls(self, toolcalls: ToolCalls) -> Any | None:
        """Extracts actions from the response."""
        if not toolcalls:
            return None

        actions = [
            AgentlabAction.convert_toolcall_to_agentlab_action_format(call) for call in toolcalls
        ]
        actions = (
            AgentlabAction.convert_multiactions_to_agentlab_action_format(actions)
            if len(actions) > 1
            else actions[0]
        )
        return actions

    def _extract_env_actions_from_text_response(
        self, response: "openai.types.chat.ChatCompletion"
    ) -> str | None:
        """Extracts environment actions from the text response."""
        # Use when action space is not given as tools.
        pass

    @staticmethod
    def format_tools_for_chat_completion(tools):
        """Formats response tools format for OpenAI Chat Completion API.

        Why we need this?
        Ans: actionset.to_tool_description() in bgym only returns description
        format valid for OpenAI Response API.

        Args:
            tools: List of tool descriptions to format for Chat Completion API.

        Returns:
            Formatted tools list compatible with OpenAI Chat Completion API, or None if tools is None.
        """
        formatted_tools = None
        if tools is not None:
            formatted_tools = [
                {
                    "type": tool["type"],
                    "function": {k: tool[k] for k in ("name", "description", "parameters")},
                }
                for tool in tools
            ]
        return formatted_tools

class ClaudeResponseModel(BaseModelWithPricing):
    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = 100,
    ):
        self.action_space_as_tools = True  # this should be a config

        super().__init__(
            model_name=model_name,
            api_key=api_key,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        client_args = {}


ClaudeResponseModel

        format valid for OpenAI Response API.

        Args:
            tools: List of tool descriptions to format for Chat Completion API.

        Returns:
            Formatted tools list compatible with OpenAI Chat Completion API, or None if tools is None.
        """
        formatted_tools = None
        if tools is not None:
            formatted_tools = [
                {
                    "type": tool["type"],
                    "function": {k: tool[k] for k in ("name", "description", "parameters")},
                }
                for tool in tools
            ]
        return formatted_tools

class ClaudeResponseModel(BaseModelWithPricing):
    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = 100,
    ):
        self.action_space_as_tools = True  # this should be a config

        super().__init__(
            model_name=model_name,
            api_key=api_key,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        client_args = {}
        if base_url is not None:
            client_args["base_url"] = base_url
        if api_key is not None:
            client_args["api_key"] = api_key
        self.client = Anthropic(**client_args)
        self.init_pricing_tracker(pricing_api="anthropic")  # Use the PricingMixin

    def _call_api(self, payload: APIPayload) -> Completion:
        sys_msg, other_msgs = self.filter_system_messages(payload.messages)
        sys_msg_text = "\n".join(c["text"] for m in sys_msg for c in m.content)
        input = []
        for msg in other_msgs:
            temp = msg.prepare_message()
            if payload.use_cache_breakpoints:
                temp = self.apply_cache_breakpoints(msg, temp)
            input.extend(temp)

        api_params: Dict[str, Any] = {
            "model": self.model_name,
            "messages": input,
            "system": sys_msg_text,
        }  # Anthropic API expects system message as a string

        if self.temperature is not None:
            api_params["temperature"] = self.temperature
        if self.max_tokens is not None:
            api_params["max_tokens"] = self.max_tokens

        if payload.tools is not None:
            api_params["tools"] = payload.tools
        if payload.tool_choice is not None and payload.force_call_tool is None:
            api_params["tool_choice"] = (
                {"type": "any"}
                if payload.tool_choice in ("required", "any")
                else {"type": payload.tool_choice}
            )
        if payload.force_call_tool is not None:
            api_params["tool_choice"] = {"type": "tool", "name": payload.force_call_tool}
        if payload.cache_tool_definition:
            # Indicating cache control for the last message enables caching of the last message.
            api_params["tools"][-1]["cache_control"] = {"type": "ephemeral"}
        if payload.cache_complete_prompt:
            # Indicating cache control for the last message enables caching of the complete prompt.
            api_params["messages"][-1]["content"][-1]["cache_control"] = {"type": "ephemeral"}

        response = call_anthropic_api_with_retries(self.client.messages.create, api_params)

        return response

    @staticmethod
    def filter_system_messages(messages: list[dict | MessageBuilder]) -> tuple[MessageBuilder]:
        """Filter system messages from the list of messages."""
        # System message cannot have an image in the middle of the text sequences.
        # Images can be appended in the end of the system message.

        sys_msgs, other_msgs = [], []
        for msg in messages:
            if isinstance(msg, MessageBuilder) and msg.role == "system":
                sys_msgs.append(msg)
                for c in msg.content:
                    if c.get("type") == "image":
                        raise TypeError("System messages cannot contain images.")
            else:
                other_msgs.append(msg)
        return sys_msgs, other_msgs

    def _parse_response(self, response: "AnthrophicMessage") -> LLMOutput:

        toolcalls = self._extract_tool_calls_from_response(response)
        think_output = self._extract_thinking_content_from_response(response)
        if self.action_space_as_tools:
            env_action = self._extract_env_actions_from_toolcalls(toolcalls)
        else:
            env_action = self._extract_env_actions_from_text_response(response)
        return LLMOutput(
            raw_response=response,
            think=think_output,
            action=env_action if env_action is not None else None,
            tool_calls=toolcalls if toolcalls is not None else None,
        )

    def _extract_tool_calls_from_response(self, response: "AnthrophicMessage") -> ToolCalls:
        """Extracts tool calls from the response."""
        tool_calls = []
        for output in response.content:
            if output.type == "tool_use":
                tool_calls.append(
                    ToolCall(
                        name=output.name,
                        arguments=output.input,
                        raw_call=output,
                    )
                )
        return ToolCalls(tool_calls=tool_calls, raw_calls=response)

    def _extract_thinking_content_from_response(self, response: "AnthrophicMessage"):
        """Extracts the thinking content from the response."""
        return "".join(output.text for output in response.content if output.type == "text")

    def _extract_env_actions_from_toolcalls(self, toolcalls: ToolCalls) -> Any | None:
        """Extracts actions from the response."""
        if not toolcalls:
            return None

        actions = [
            AgentlabAction.convert_toolcall_to_agentlab_action_format(call) for call in toolcalls
        ]
        actions = (
            AgentlabAction.convert_multiactions_to_agentlab_action_format(actions)
            if len(actions) > 1
            else actions[0]
        )
        return actions

    def _extract_env_actions_from_text_response(self, response: "AnthrophicMessage") -> str | None:
        """Extracts environment actions from the text response."""
        # Use when action space is not given as tools.
        pass

    def apply_cache_breakpoints(self, msg: Message, prepared_msg: dict) -> List[Message]:
        """Apply cache breakpoints to the messages."""
        if getattr(msg, "_cache_breakpoint", False):
            prepared_msg[-1]["content"][-1]["cache_control"] = {"type": "ephemeral"}
        return prepared_msg

# Factory classes to create the appropriate model based on the API endpoint.

@dataclass
class OpenAIResponseModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenAI
    model."""

    api = "openai"

    def make_model(self):
        return OpenAIResponseModel(
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:


OpenAIResponseModelArgs

            else actions[0]
        )
        return actions

    def _extract_env_actions_from_text_response(self, response: "AnthrophicMessage") -> str | None:
        """Extracts environment actions from the text response."""
        # Use when action space is not given as tools.
        pass

    def apply_cache_breakpoints(self, msg: Message, prepared_msg: dict) -> List[Message]:
        """Apply cache breakpoints to the messages."""
        if getattr(msg, "_cache_breakpoint", False):
            prepared_msg[-1]["content"][-1]["cache_control"] = {"type": "ephemeral"}
        return prepared_msg

# Factory classes to create the appropriate model based on the API endpoint.

@dataclass
class OpenAIResponseModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenAI
    model."""

    api = "openai"

    def make_model(self):
        return OpenAIResponseModel(
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return OpenAIResponseAPIMessageBuilder

@dataclass
class ClaudeResponseModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenAI
    model."""

    api = "anthropic"

    def make_model(self):
        return ClaudeResponseModel(
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return AnthropicAPIMessageBuilder


ClaudeResponseModelArgs

@dataclass
class OpenAIResponseModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenAI
    model."""

    api = "openai"

    def make_model(self):
        return OpenAIResponseModel(
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return OpenAIResponseAPIMessageBuilder

@dataclass
class ClaudeResponseModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenAI
    model."""

    api = "anthropic"

    def make_model(self):
        return ClaudeResponseModel(
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return AnthropicAPIMessageBuilder

@dataclass
class OpenAIChatModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenAI
    model."""

    api = "openai"

    def make_model(self):
        return OpenAIChatCompletionModel(
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return OpenAIChatCompletionAPIMessageBuilder


OpenAIChatModelArgs

@dataclass
class ClaudeResponseModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenAI
    model."""

    api = "anthropic"

    def make_model(self):
        return ClaudeResponseModel(
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return AnthropicAPIMessageBuilder

@dataclass
class OpenAIChatModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenAI
    model."""

    api = "openai"

    def make_model(self):
        return OpenAIChatCompletionModel(
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return OpenAIChatCompletionAPIMessageBuilder

@dataclass
class OpenRouterModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenRouter
    model."""

    api: str = "openai"  # tool description format used by actionset.to_tool_description() in bgym

    def make_model(self):
        return OpenAIChatCompletionModel(
            base_url="https://openrouter.ai/api/v1",
            api_key=os.getenv("OPENROUTER_API_KEY"),
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return OpenAIChatCompletionAPIMessageBuilder


OpenRouterModelArgs

@dataclass
class OpenAIChatModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenAI
    model."""

    api = "openai"

    def make_model(self):
        return OpenAIChatCompletionModel(
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return OpenAIChatCompletionAPIMessageBuilder

@dataclass
class OpenRouterModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenRouter
    model."""

    api: str = "openai"  # tool description format used by actionset.to_tool_description() in bgym

    def make_model(self):
        return OpenAIChatCompletionModel(
            base_url="https://openrouter.ai/api/v1",
            api_key=os.getenv("OPENROUTER_API_KEY"),
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return OpenAIChatCompletionAPIMessageBuilder

def tool_call_to_python_code(func_name, kwargs):
    """Format a function name and kwargs dict into a Python function call string."""
    if kwargs is None:
        kwargs = {}

    if not kwargs:
        return f"{func_name}()"

    args_str = ", ".join(f"{key}={repr(value)}" for key, value in kwargs.items())
    return f"{func_name}({args_str})"

# ___Not__Tested__#

# class VLLMModelArgs(BaseModelArgs):
#     """Serializable object for instantiating a generic chat model with a VLLM
#     model."""


@dataclass
class OpenRouterModelArgs(BaseModelArgs):
    """Serializable object for instantiating a generic chat model with an OpenRouter
    model."""

    api: str = "openai"  # tool description format used by actionset.to_tool_description() in bgym

    def make_model(self):
        return OpenAIChatCompletionModel(
            base_url="https://openrouter.ai/api/v1",
            api_key=os.getenv("OPENROUTER_API_KEY"),
            model_name=self.model_name,
            temperature=self.temperature,
            max_tokens=self.max_new_tokens,
        )

    def get_message_builder(self) -> MessageBuilder:
        return OpenAIChatCompletionAPIMessageBuilder

def tool_call_to_python_code(func_name, kwargs):
    """Format a function name and kwargs dict into a Python function call string."""
    if kwargs is None:
        kwargs = {}

    if not kwargs:
        return f"{func_name}()"

    args_str = ", ".join(f"{key}={repr(value)}" for key, value in kwargs.items())
    return f"{func_name}({args_str})"

# ___Not__Tested__#

# class VLLMModelArgs(BaseModelArgs):
#     """Serializable object for instantiating a generic chat model with a VLLM
#     model."""

#     api = "openai"  # tool description format used by actionset.to_tool_description() in bgym

#     def make_model(self, extra_kwargs=None, **kwargs):
#         return OpenAIChatCompletionModel(
#             client_args={
#                 "base_url": "http://localhost:8000/v1",
#                 "api_key": os.getenv("VLLM_API_KEY", "EMPTY"),
#             },
#             model_name=self.model_name,  # this needs to be set
#             temperature=self.temperature,
#             max_tokens=self.max_new_tokens,
#             extra_kwargs=extra_kwargs,


ContentItem = Dict[str, Any]
Message = Dict[str, Union[str, List[ContentItem]]]

@dataclass
class ToolCall:
    """Represents a tool call made by the LLM.
    Attributes:
    name: Name of the tool called.
    arguments: Arguments passed to the tool.
    raw_call: The raw call object from the LLM API.
    tool_response: Output of the tool call goes here. It can be only one content item.
    """

    name: str = field(default=None)
    arguments: Dict[str, Any] = field(default_factory=dict)
    raw_call: Any = field(default=None)
    tool_response: ContentItem = None

    @property
    def is_response_set(self) -> bool:
        """Check if the tool response is set."""
        return self.tool_response is not None

    def response_text(self, text: str) -> "MessageBuilder":
        self.tool_response = {"text": text}
        return self

    def response_image(self, image: str) -> "MessageBuilder":
        self.tool_response = {"image": image}
        return self

    def __repr__(self):
        return f"ToolCall(name={self.name}, arguments={self.arguments})"

@dataclass
class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.


@dataclass
class ToolCall:
    """Represents a tool call made by the LLM.
    Attributes:
    name: Name of the tool called.
    arguments: Arguments passed to the tool.
    raw_call: The raw call object from the LLM API.
    tool_response: Output of the tool call goes here. It can be only one content item.
    """

    name: str = field(default=None)
    arguments: Dict[str, Any] = field(default_factory=dict)
    raw_call: Any = field(default=None)
    tool_response: ContentItem = None

    @property
    def is_response_set(self) -> bool:
        """Check if the tool response is set."""
        return self.tool_response is not None

    def response_text(self, text: str) -> "MessageBuilder":
        self.tool_response = {"text": text}
        return self

    def response_image(self, image: str) -> "MessageBuilder":
        self.tool_response = {"image": image}
        return self

    def __repr__(self):
        return f"ToolCall(name={self.name}, arguments={self.arguments})"

@dataclass
class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)


    name: Name of the tool called.
    arguments: Arguments passed to the tool.
    raw_call: The raw call object from the LLM API.
    tool_response: Output of the tool call goes here. It can be only one content item.
    """

    name: str = field(default=None)
    arguments: Dict[str, Any] = field(default_factory=dict)
    raw_call: Any = field(default=None)
    tool_response: ContentItem = None

    @property
    def is_response_set(self) -> bool:
        """Check if the tool response is set."""
        return self.tool_response is not None

    def response_text(self, text: str) -> "MessageBuilder":
        self.tool_response = {"text": text}
        return self

    def response_image(self, image: str) -> "MessageBuilder":
        self.tool_response = {"image": image}
        return self

    def __repr__(self):
        return f"ToolCall(name={self.name}, arguments={self.arguments})"

@dataclass
class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self


    """

    name: str = field(default=None)
    arguments: Dict[str, Any] = field(default_factory=dict)
    raw_call: Any = field(default=None)
    tool_response: ContentItem = None

    @property
    def is_response_set(self) -> bool:
        """Check if the tool response is set."""
        return self.tool_response is not None

    def response_text(self, text: str) -> "MessageBuilder":
        self.tool_response = {"text": text}
        return self

    def response_image(self, image: str) -> "MessageBuilder":
        self.tool_response = {"image": image}
        return self

    def __repr__(self):
        return f"ToolCall(name={self.name}, arguments={self.arguments})"

@dataclass
class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self

    @property
    def all_responses_set(self) -> bool:


    def response_image(self, image: str) -> "MessageBuilder":
        self.tool_response = {"image": image}
        return self

    def __repr__(self):
        return f"ToolCall(name={self.name}, arguments={self.arguments})"

@dataclass
class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self

    @property
    def all_responses_set(self) -> bool:
        """Check if all tool calls have responses set."""
        return all(call.is_response_set for call in self.tool_calls)

    def __len__(self) -> int:
        """Return the number of tool calls."""
        return len(self.tool_calls)

    def __iter__(self):
        """Make ToolCalls iterable."""
        return iter(self.tool_calls)

    def __bool__(self):
        """Check if there are any tool calls."""
        return len(self.tool_calls) > 0

@dataclass


        return f"ToolCall(name={self.name}, arguments={self.arguments})"

@dataclass
class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self

    @property
    def all_responses_set(self) -> bool:
        """Check if all tool calls have responses set."""
        return all(call.is_response_set for call in self.tool_calls)

    def __len__(self) -> int:
        """Return the number of tool calls."""
        return len(self.tool_calls)

    def __iter__(self):
        """Make ToolCalls iterable."""
        return iter(self.tool_calls)

    def __bool__(self):
        """Check if there are any tool calls."""
        return len(self.tool_calls) > 0

@dataclass
class LLMOutput:
    """Serializable object for the output of a response LLM."""

    raw_response: Any = field(default=None)
    think: str = field(default="")


class ToolCalls:
    """A collection of tool calls made by the LLM.

    Attributes:
    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self

    @property
    def all_responses_set(self) -> bool:
        """Check if all tool calls have responses set."""
        return all(call.is_response_set for call in self.tool_calls)

    def __len__(self) -> int:
        """Return the number of tool calls."""
        return len(self.tool_calls)

    def __iter__(self):
        """Make ToolCalls iterable."""
        return iter(self.tool_calls)

    def __bool__(self):
        """Check if there are any tool calls."""
        return len(self.tool_calls) > 0

@dataclass
class LLMOutput:
    """Serializable object for the output of a response LLM."""

    raw_response: Any = field(default=None)
    think: str = field(default="")
    action: str | None = field(default=None)  # Default action if no tool call is made
    tool_calls: ToolCalls | None = field(
        default=None
    )  # This will hold the tool call response if any


    tool_calls: List of ToolCall objects.
    raw_calls: Represents raw tool calls object returned by a LLM API, may contain one or more tool calls.
    """

    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self

    @property
    def all_responses_set(self) -> bool:
        """Check if all tool calls have responses set."""
        return all(call.is_response_set for call in self.tool_calls)

    def __len__(self) -> int:
        """Return the number of tool calls."""
        return len(self.tool_calls)

    def __iter__(self):
        """Make ToolCalls iterable."""
        return iter(self.tool_calls)

    def __bool__(self):
        """Check if there are any tool calls."""
        return len(self.tool_calls) > 0

@dataclass
class LLMOutput:
    """Serializable object for the output of a response LLM."""

    raw_response: Any = field(default=None)
    think: str = field(default="")
    action: str | None = field(default=None)  # Default action if no tool call is made
    tool_calls: ToolCalls | None = field(
        default=None
    )  # This will hold the tool call response if any

class MessageBuilder:
    def __init__(self, role: str):


    tool_calls: List[ToolCall] = field(default_factory=list)
    raw_calls: List[Any] = field(default_factory=list)

    def add_tool_call(self, tool_call: ToolCall) -> "ToolCalls":
        self.tool_calls.append(tool_call)
        return self

    @property
    def all_responses_set(self) -> bool:
        """Check if all tool calls have responses set."""
        return all(call.is_response_set for call in self.tool_calls)

    def __len__(self) -> int:
        """Return the number of tool calls."""
        return len(self.tool_calls)

    def __iter__(self):
        """Make ToolCalls iterable."""
        return iter(self.tool_calls)

    def __bool__(self):
        """Check if there are any tool calls."""
        return len(self.tool_calls) > 0

@dataclass
class LLMOutput:
    """Serializable object for the output of a response LLM."""

    raw_response: Any = field(default=None)
    think: str = field(default="")
    action: str | None = field(default=None)  # Default action if no tool call is made
    tool_calls: ToolCalls | None = field(
        default=None
    )  # This will hold the tool call response if any

class MessageBuilder:
    def __init__(self, role: str):
        self.role = role
        self.content: List[ContentItem] = []
        self.responded_tool_calls: ToolCalls = None


        Args:
            tools: List of tool descriptions to format for Chat Completion API.

        Returns:
            Formatted tools list compatible with OpenAI Chat Completion API, or None if tools is None.
        """
        formatted_tools = None
        if tools is not None:
            formatted_tools = [
                {
                    "type": tool["type"],
                    "function": {k: tool[k] for k in ("name", "description", "parameters")},
                }
                for tool in tools
            ]
        return formatted_tools

class ClaudeResponseModel(BaseModelWithPricing):
    def __init__(
        self,
        model_name: str,
        base_url: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: float | None = None,
        max_tokens: int | None = 100,
    ):
        self.action_space_as_tools = True  # this should be a config

        super().__init__(
            model_name=model_name,
            api_key=api_key,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        client_args = {}
        if base_url is not None:
            client_args["base_url"] = base_url
        if api_key is not None:
            client_args["api_key"] = api_key
        self.client = Anthropic(**client_args)
        self.init_pricing_tracker(pricing_api="anthropic")  # Use the PricingMixin

    def _call_api(self, payload: APIPayload) -> Completion:
        sys_msg, other_msgs = self.filter_system_messages(payload.messages)
        sys_msg_text = "\n".join(c["text"] for m in sys_msg for c in m.content)
        input = []
        for msg in other_msgs:
            temp = msg.prepare_message()
            if payload.use_cache_breakpoints:
                temp = self.apply_cache_breakpoints(msg, temp)
            input.extend(temp)

        api_params: Dict[str, Any] = {
            "model": self.model_name,
            "messages": input,
            "system": sys_msg_text,
        }  # Anthropic API expects system message as a string

        if self.temperature is not None:
            api_params["temperature"] = self.temperature
        if self.max_tokens is not None:


system

        """Add an image URL to the message content."""
        self.content.append({"image": image_to_png_base64_url(image_url)})
        return self

    def mark_all_previous_msg_for_caching(self):
        """Insert a cache breakpoint in the message content."""
        # This is a placeholder for future implementation.
        raise NotImplementedError

    @classmethod
    def add_responded_tool_calls(cls, responded_tool_calls: ToolCalls) -> "MessageBuilder":
        """Add tool calls to the message content."""
        assert responded_tool_calls.all_responses_set, "All tool calls must have a response."
        msg = cls("tool")
        msg.responded_tool_calls = responded_tool_calls
        return msg

class OpenAIResponseAPIMessageBuilder(MessageBuilder):
    @classmethod
    def system(cls) -> "OpenAIResponseAPIMessageBuilder":
        # OpenAI Responses API uses 'developer' role for system messages
        return cls("developer")

    def prepare_message(self) -> List[Message]:
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]

        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Convert the content item to the expected format for OpenAI Responses."""
        if "text" in content:
            content_type = "input_text" if self.role != "assistant" else "output_text"
            return {"type": content_type, "text": content["text"]}
        elif "image" in content:
            return {"type": "input_image", "image_url": content["image"]}
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def handle_tool_call(self) -> List[Message]:


user

    raw_response: Any = field(default=None)
    think: str = field(default="")
    action: str | None = field(default=None)  # Default action if no tool call is made
    tool_calls: ToolCalls | None = field(
        default=None
    )  # This will hold the tool call response if any

class MessageBuilder:
    def __init__(self, role: str):
        self.role = role
        self.content: List[ContentItem] = []
        self.responded_tool_calls: ToolCalls = None

    @classmethod
    def system(cls) -> "MessageBuilder":
        return cls("system")

    @classmethod
    def user(cls) -> "MessageBuilder":
        return cls("user")

    @classmethod
    def assistant(cls) -> "MessageBuilder":
        return cls("assistant")

    @abstractmethod
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the API call."""
        raise NotImplementedError("Subclasses must implement this method.")

    def add_text(self, text: str) -> "MessageBuilder":
        self.content.append({"text": text})
        return self

    def add_image(self, image: str) -> "MessageBuilder":
        self.content.append({"image": image})
        return self

    def to_markdown(self) -> str:
        parts = []


assistant

    tool_calls: ToolCalls | None = field(
        default=None
    )  # This will hold the tool call response if any

class MessageBuilder:
    def __init__(self, role: str):
        self.role = role
        self.content: List[ContentItem] = []
        self.responded_tool_calls: ToolCalls = None

    @classmethod
    def system(cls) -> "MessageBuilder":
        return cls("system")

    @classmethod
    def user(cls) -> "MessageBuilder":
        return cls("user")

    @classmethod
    def assistant(cls) -> "MessageBuilder":
        return cls("assistant")

    @abstractmethod
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the API call."""
        raise NotImplementedError("Subclasses must implement this method.")

    def add_text(self, text: str) -> "MessageBuilder":
        self.content.append({"text": text})
        return self

    def add_image(self, image: str) -> "MessageBuilder":
        self.content.append({"image": image})
        return self

    def to_markdown(self) -> str:
        parts = []
        for item in self.content:
            if "text" in item:
                parts.append(f"\n```\n{item['text']}\n```\n")
            elif "image" in item:


            # e.g. "data:image/png;base64"
            if img_str.startswith("data:image/png;base64,"):
                img_str = img_str[len("data:image/png;base64,") :]
            return {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png",
                    "data": img_str,
                },
            }
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def mark_all_previous_msg_for_caching(self) -> List[Message]:
        """Insert a cache breakpoint in the message content to mark all previous messages for caching."""
        self._cache_breakpoint = True

class OpenAIChatCompletionAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the OpenAI API."""
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]
        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by OpenAI ChatCompletion."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content:
            return {"type": "image_url", "image_url": {"url": content["image"]}}
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:
            raise ValueError("No tool calls found in responded_tool_calls")
        output = []
        output.append(
            self.responded_tool_calls.raw_calls.choices[0].message
        )  # add raw calls to output
        for fn_call in self.responded_tool_calls:
            raw_call = fn_call.raw_call


        self.content: List[ContentItem] = []
        self.responded_tool_calls: ToolCalls = None

    @classmethod
    def system(cls) -> "MessageBuilder":
        return cls("system")

    @classmethod
    def user(cls) -> "MessageBuilder":
        return cls("user")

    @classmethod
    def assistant(cls) -> "MessageBuilder":
        return cls("assistant")

    @abstractmethod
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the API call."""
        raise NotImplementedError("Subclasses must implement this method.")

    def add_text(self, text: str) -> "MessageBuilder":
        self.content.append({"text": text})
        return self

    def add_image(self, image: str) -> "MessageBuilder":
        self.content.append({"image": image})
        return self

    def to_markdown(self) -> str:
        parts = []
        for item in self.content:
            if "text" in item:
                parts.append(f"\n```\n{item['text']}\n```\n")
            elif "image" in item:
                parts.append(f"![Image]({item['image']})")

        # Tool call markdown repr
        if self.responded_tool_calls is not None:
            for i, tool_call in enumerate(self.responded_tool_calls.tool_calls, 1):
                parts.append(
                    f"\n**Tool Call {i}**: {tool_call_to_python_code(tool_call.name, tool_call.arguments)}"
                )
                response = tool_call.tool_response


    def system(cls) -> "MessageBuilder":
        return cls("system")

    @classmethod
    def user(cls) -> "MessageBuilder":
        return cls("user")

    @classmethod
    def assistant(cls) -> "MessageBuilder":
        return cls("assistant")

    @abstractmethod
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the API call."""
        raise NotImplementedError("Subclasses must implement this method.")

    def add_text(self, text: str) -> "MessageBuilder":
        self.content.append({"text": text})
        return self

    def add_image(self, image: str) -> "MessageBuilder":
        self.content.append({"image": image})
        return self

    def to_markdown(self) -> str:
        parts = []
        for item in self.content:
            if "text" in item:
                parts.append(f"\n```\n{item['text']}\n```\n")
            elif "image" in item:
                parts.append(f"![Image]({item['image']})")

        # Tool call markdown repr
        if self.responded_tool_calls is not None:
            for i, tool_call in enumerate(self.responded_tool_calls.tool_calls, 1):
                parts.append(
                    f"\n**Tool Call {i}**: {tool_call_to_python_code(tool_call.name, tool_call.arguments)}"
                )
                response = tool_call.tool_response
                if response is not None:
                    parts.append(f"\n**Tool Response {i}:**")
                    content = (
                        f"```\n{response['text']}\n```"


    def user(cls) -> "MessageBuilder":
        return cls("user")

    @classmethod
    def assistant(cls) -> "MessageBuilder":
        return cls("assistant")

    @abstractmethod
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the API call."""
        raise NotImplementedError("Subclasses must implement this method.")

    def add_text(self, text: str) -> "MessageBuilder":
        self.content.append({"text": text})
        return self

    def add_image(self, image: str) -> "MessageBuilder":
        self.content.append({"image": image})
        return self

    def to_markdown(self) -> str:
        parts = []
        for item in self.content:
            if "text" in item:
                parts.append(f"\n```\n{item['text']}\n```\n")
            elif "image" in item:
                parts.append(f"![Image]({item['image']})")

        # Tool call markdown repr
        if self.responded_tool_calls is not None:
            for i, tool_call in enumerate(self.responded_tool_calls.tool_calls, 1):
                parts.append(
                    f"\n**Tool Call {i}**: {tool_call_to_python_code(tool_call.name, tool_call.arguments)}"
                )
                response = tool_call.tool_response
                if response is not None:
                    parts.append(f"\n**Tool Response {i}:**")
                    content = (
                        f"```\n{response['text']}\n```"
                        if "text" in response
                        else f"![Tool Response Image]({response['image']})"
                    )
                    parts.append(content)

        markdown = f"### {self.role.capitalize()}\n"
        markdown += "\n".join(parts)

        return markdown

    def add_image_url(self, image_url: str) -> "MessageBuilder":
        """Add an image URL to the message content."""
        self.content.append({"image": image_to_png_base64_url(image_url)})
        return self

    def mark_all_previous_msg_for_caching(self):
        """Insert a cache breakpoint in the message content."""
        # This is a placeholder for future implementation.
        raise NotImplementedError

    @classmethod
    def add_responded_tool_calls(cls, responded_tool_calls: ToolCalls) -> "MessageBuilder":
        """Add tool calls to the message content."""
        assert responded_tool_calls.all_responses_set, "All tool calls must have a response."
        msg = cls("tool")
        msg.responded_tool_calls = responded_tool_calls
        return msg


        if self.responded_tool_calls is not None:
            for i, tool_call in enumerate(self.responded_tool_calls.tool_calls, 1):
                parts.append(
                    f"\n**Tool Call {i}**: {tool_call_to_python_code(tool_call.name, tool_call.arguments)}"
                )
                response = tool_call.tool_response
                if response is not None:
                    parts.append(f"\n**Tool Response {i}:**")
                    content = (
                        f"```\n{response['text']}\n```"
                        if "text" in response
                        else f"![Tool Response Image]({response['image']})"
                    )
                    parts.append(content)

        markdown = f"### {self.role.capitalize()}\n"
        markdown += "\n".join(parts)

        return markdown

    def add_image_url(self, image_url: str) -> "MessageBuilder":
        """Add an image URL to the message content."""
        self.content.append({"image": image_to_png_base64_url(image_url)})
        return self

    def mark_all_previous_msg_for_caching(self):
        """Insert a cache breakpoint in the message content."""
        # This is a placeholder for future implementation.
        raise NotImplementedError

    @classmethod
    def add_responded_tool_calls(cls, responded_tool_calls: ToolCalls) -> "MessageBuilder":
        """Add tool calls to the message content."""
        assert responded_tool_calls.all_responses_set, "All tool calls must have a response."
        msg = cls("tool")
        msg.responded_tool_calls = responded_tool_calls
        return msg

class OpenAIResponseAPIMessageBuilder(MessageBuilder):
    @classmethod
    def system(cls) -> "OpenAIResponseAPIMessageBuilder":
        # OpenAI Responses API uses 'developer' role for system messages
        return cls("developer")


        output = []
        output.append(
            self.responded_tool_calls.raw_calls.choices[0].message
        )  # add raw calls to output
        for fn_call in self.responded_tool_calls:
            raw_call = fn_call.raw_call
            assert (
                "image" not in fn_call.tool_response
            ), "Image output is not supported in function calls response."
            # a function_call_output dict has keys "role", "tool_call_id" and "content"
            tool_call_reponse = {
                "name": raw_call["function"]["name"],  # required with OpenRouter
                "role": "tool",
                "tool_call_id": raw_call["id"],
                "content": self.convert_content_to_expected_format(fn_call.tool_response)["text"],
            }
            output.append(tool_call_reponse)

        return output

    def mark_all_previous_msg_for_caching(self):
        """Nothing special to do here for openAI. They do not have a notion of cache breakpoints."""
        pass

@dataclass
class APIPayload:
    messages: List[MessageBuilder] | None = None
    tools: List[Dict[str, Any]] | None = None
    tool_choice: Literal["none", "auto", "any", "required"] | None = None
    force_call_tool: str | None = (
        None  # Name of the tool to call # If set, will force the LLM to call this tool.
    )
    use_cache_breakpoints: bool = (
        False  # If True, will apply cache breakpoints to the messages. # applicable for Anthropic
    )
    cache_tool_definition: bool = (
        False  # If True, will cache the tool definition in the last message.
    )
    cache_complete_prompt: bool = (
        False  # If True, will cache the complete prompt in the last message.
    )
    reasoning_effort: Literal["low", "medium", "high"] | None = None


                        else f"![Tool Response Image]({response['image']})"
                    )
                    parts.append(content)

        markdown = f"### {self.role.capitalize()}\n"
        markdown += "\n".join(parts)

        return markdown

    def add_image_url(self, image_url: str) -> "MessageBuilder":
        """Add an image URL to the message content."""
        self.content.append({"image": image_to_png_base64_url(image_url)})
        return self

    def mark_all_previous_msg_for_caching(self):
        """Insert a cache breakpoint in the message content."""
        # This is a placeholder for future implementation.
        raise NotImplementedError

    @classmethod
    def add_responded_tool_calls(cls, responded_tool_calls: ToolCalls) -> "MessageBuilder":
        """Add tool calls to the message content."""
        assert responded_tool_calls.all_responses_set, "All tool calls must have a response."
        msg = cls("tool")
        msg.responded_tool_calls = responded_tool_calls
        return msg

class OpenAIResponseAPIMessageBuilder(MessageBuilder):
    @classmethod
    def system(cls) -> "OpenAIResponseAPIMessageBuilder":
        # OpenAI Responses API uses 'developer' role for system messages
        return cls("developer")

    def prepare_message(self) -> List[Message]:
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]

        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Convert the content item to the expected format for OpenAI Responses."""
        if "text" in content:
            content_type = "input_text" if self.role != "assistant" else "output_text"


                    "data": img_str,
                },
            }
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def mark_all_previous_msg_for_caching(self) -> List[Message]:
        """Insert a cache breakpoint in the message content to mark all previous messages for caching."""
        self._cache_breakpoint = True

class OpenAIChatCompletionAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the OpenAI API."""
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]
        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by OpenAI ChatCompletion."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content:
            return {"type": "image_url", "image_url": {"url": content["image"]}}
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:
            raise ValueError("No tool calls found in responded_tool_calls")
        output = []
        output.append(
            self.responded_tool_calls.raw_calls.choices[0].message
        )  # add raw calls to output
        for fn_call in self.responded_tool_calls:
            raw_call = fn_call.raw_call
            assert (
                "image" not in fn_call.tool_response
            ), "Image output is not supported in function calls response."
            # a function_call_output dict has keys "role", "tool_call_id" and "content"
            tool_call_reponse = {
                "name": raw_call["function"]["name"],  # required with OpenRouter
                "role": "tool",
                "tool_call_id": raw_call["id"],
                "content": self.convert_content_to_expected_format(fn_call.tool_response)["text"],


class OpenAIChatCompletionAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the OpenAI API."""
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]
        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by OpenAI ChatCompletion."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content:
            return {"type": "image_url", "image_url": {"url": content["image"]}}
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def handle_tool_call(self) -> List[Message]:
        """Handle the tool call response from the last raw response."""
        if self.responded_tool_calls is None:
            raise ValueError("No tool calls found in responded_tool_calls")
        output = []
        output.append(
            self.responded_tool_calls.raw_calls.choices[0].message
        )  # add raw calls to output
        for fn_call in self.responded_tool_calls:
            raw_call = fn_call.raw_call
            assert (
                "image" not in fn_call.tool_response
            ), "Image output is not supported in function calls response."
            # a function_call_output dict has keys "role", "tool_call_id" and "content"
            tool_call_reponse = {
                "name": raw_call["function"]["name"],  # required with OpenRouter
                "role": "tool",
                "tool_call_id": raw_call["id"],
                "content": self.convert_content_to_expected_format(fn_call.tool_response)["text"],
            }
            output.append(tool_call_reponse)

        return output

    def mark_all_previous_msg_for_caching(self):
        """Nothing special to do here for openAI. They do not have a notion of cache breakpoints."""
        pass

@dataclass
class APIPayload:
    messages: List[MessageBuilder] | None = None
    tools: List[Dict[str, Any]] | None = None
    tool_choice: Literal["none", "auto", "any", "required"] | None = None
    force_call_tool: str | None = (
        None  # Name of the tool to call # If set, will force the LLM to call this tool.
    )
    use_cache_breakpoints: bool = (
        False  # If True, will apply cache breakpoints to the messages. # applicable for Anthropic
    )
    cache_tool_definition: bool = (
        False  # If True, will cache the tool definition in the last message.
    )


            "role": "assistant",
            "content": self.responded_tool_calls.raw_calls.content,
        }  # Add the toolcall block
        tool_response = {"role": "user", "content": []}  # Anthropic expects a list of messages
        for call in self.responded_tool_calls:
            assert (
                "image" not in call.tool_response
            ), "Image output is not supported in tool calls response."
            tool_response["content"].append(
                {
                    "type": "tool_result",
                    "tool_use_id": call.raw_call.id,
                    "content": self.transform_content(call.tool_response)[
                        "text"
                    ],  # needs to be str
                }
            )

        return [llm_tool_call, tool_response]

    def transform_content(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by Anthropic API."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content:
            img_str: str = content["image"]
            # make sure to get rid of the image type for anthropic
            # e.g. "data:image/png;base64"
            if img_str.startswith("data:image/png;base64,"):
                img_str = img_str[len("data:image/png;base64,") :]
            return {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png",
                    "data": img_str,
                },
            }
        else:
            raise ValueError(f"Unsupported content type: {content}")

    def mark_all_previous_msg_for_caching(self) -> List[Message]:
        """Insert a cache breakpoint in the message content to mark all previous messages for caching."""
        self._cache_breakpoint = True

class OpenAIChatCompletionAPIMessageBuilder(MessageBuilder):
    def prepare_message(self) -> List[Message]:
        """Prepare the message for the OpenAI API."""
        content = []
        for item in self.content:
            content.append(self.convert_content_to_expected_format(item))
        output = [{"role": self.role, "content": content}]
        return output if self.role != "tool" else self.handle_tool_call()

    def convert_content_to_expected_format(self, content: ContentItem) -> ContentItem:
        """Transform content item to the format expected by OpenAI ChatCompletion."""
        if "text" in content:
            return {"type": "text", "text": content["text"]}
        elif "image" in content: