Python: cleanup and refactoring of chat clients (#2937)

* refactoring and unifying naming schemes of internal methods of chat clients

* set tool_choice to auto

* fix for mypy

* added note on naming and fix #2951

* fix responses

* fixes in azure ai agents client
This commit is contained in:
Eduard van Valkenburg
2025-12-18 12:02:23 +00:00
committed by GitHub
parent a71f768331
commit e5c11d38d6
26 changed files with 1128 additions and 1068 deletions
@@ -25,7 +25,6 @@ from agent_framework import (
TextContent,
TextReasoningContent,
TextSpanRegion,
ToolProtocol,
UsageContent,
UsageDetails,
get_logger,
@@ -214,9 +213,11 @@ class AnthropicClient(BaseChatClient):
chat_options: ChatOptions,
**kwargs: Any,
) -> ChatResponse:
# Extract necessary state from messages and options
run_options = self._create_run_options(messages, chat_options, **kwargs)
# prepare
run_options = self._prepare_options(messages, chat_options, **kwargs)
# execute
message = await self.anthropic_client.beta.messages.create(**run_options, stream=False)
# process
return self._process_message(message)
async def _inner_get_streaming_response(
@@ -226,16 +227,17 @@ class AnthropicClient(BaseChatClient):
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
# Extract necessary state from messages and options
run_options = self._create_run_options(messages, chat_options, **kwargs)
# prepare
run_options = self._prepare_options(messages, chat_options, **kwargs)
# execute and process
async for chunk in await self.anthropic_client.beta.messages.create(**run_options, stream=True):
parsed_chunk = self._process_stream_event(chunk)
if parsed_chunk:
yield parsed_chunk
# region Create Run Options and Helpers
# region Prep methods
def _create_run_options(
def _prepare_options(
self,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
@@ -251,78 +253,91 @@ class AnthropicClient(BaseChatClient):
Returns:
A dictionary of run options for the Anthropic client.
"""
if chat_options.additional_properties and "additional_beta_flags" in chat_options.additional_properties:
betas = chat_options.additional_properties.pop("additional_beta_flags")
else:
betas = []
run_options: dict[str, Any] = {
"model": chat_options.model_id or self.model_id,
"messages": self._convert_messages_to_anthropic_format(messages),
"max_tokens": chat_options.max_tokens or ANTHROPIC_DEFAULT_MAX_TOKENS,
"extra_headers": {"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
"betas": {*BETA_FLAGS, *self.additional_beta_flags, *betas},
}
run_options: dict[str, Any] = chat_options.to_dict(
exclude={
"type",
"instructions", # handled via system message
"tool_choice", # handled separately
"allow_multiple_tool_calls", # handled via tool_choice
"additional_properties", # handled separately
}
)
# Add any additional options from chat_options or kwargs
if chat_options.temperature is not None:
run_options["temperature"] = chat_options.temperature
if chat_options.top_p is not None:
run_options["top_p"] = chat_options.top_p
if chat_options.stop is not None:
run_options["stop_sequences"] = chat_options.stop
# translations between ChatOptions and Anthropic API
translations = {
"model_id": "model",
"stop": "stop_sequences",
}
for old_key, new_key in translations.items():
if old_key in run_options and old_key != new_key:
run_options[new_key] = run_options.pop(old_key)
# model id
if not run_options.get("model"):
if not self.model_id:
raise ValueError("model_id must be a non-empty string")
run_options["model"] = self.model_id
# max_tokens - Anthropic requires this, default if not provided
if not run_options.get("max_tokens"):
run_options["max_tokens"] = ANTHROPIC_DEFAULT_MAX_TOKENS
# messages
run_options["messages"] = self._prepare_messages_for_anthropic(messages)
# system message - first system message is passed as instructions
if messages and isinstance(messages[0], ChatMessage) and messages[0].role == Role.SYSTEM:
# first system message is passed as instructions
run_options["system"] = messages[0].text
if chat_options.tool_choice is not None:
match (
chat_options.tool_choice if isinstance(chat_options.tool_choice, str) else chat_options.tool_choice.mode
):
case "auto":
run_options["tool_choice"] = {"type": "auto"}
if chat_options.allow_multiple_tool_calls is not None:
run_options["tool_choice"][ # type:ignore[reportArgumentType]
"disable_parallel_tool_use"
] = not chat_options.allow_multiple_tool_calls
case "required":
if chat_options.tool_choice.required_function_name:
run_options["tool_choice"] = {
"type": "tool",
"name": chat_options.tool_choice.required_function_name,
}
if chat_options.allow_multiple_tool_calls is not None:
run_options["tool_choice"][ # type:ignore[reportArgumentType]
"disable_parallel_tool_use"
] = not chat_options.allow_multiple_tool_calls
else:
run_options["tool_choice"] = {"type": "any"}
if chat_options.allow_multiple_tool_calls is not None:
run_options["tool_choice"][ # type:ignore[reportArgumentType]
"disable_parallel_tool_use"
] = not chat_options.allow_multiple_tool_calls
case "none":
run_options["tool_choice"] = {"type": "none"}
case _:
logger.debug(f"Ignoring unsupported tool choice mode: {chat_options.tool_choice.mode} for now")
if tools_and_mcp := self._convert_tools_to_anthropic_format(chat_options.tools):
run_options.update(tools_and_mcp)
if chat_options.additional_properties:
run_options.update(chat_options.additional_properties)
# betas
run_options["betas"] = self._prepare_betas(chat_options)
# extra headers
run_options["extra_headers"] = {"User-Agent": AGENT_FRAMEWORK_USER_AGENT}
# tools, mcp servers and tool choice
if tools_config := self._prepare_tools_for_anthropic(chat_options):
run_options.update(tools_config)
# additional properties
additional_options = {
key: value
for key, value in chat_options.additional_properties.items()
if value is not None and key != "additional_beta_flags"
}
if additional_options:
run_options.update(additional_options)
run_options.update(kwargs)
return run_options
def _convert_messages_to_anthropic_format(self, messages: MutableSequence[ChatMessage]) -> list[dict[str, Any]]:
"""Convert a list of ChatMessages to the format expected by the Anthropic client.
def _prepare_betas(self, chat_options: ChatOptions) -> set[str]:
"""Prepare the beta flags for the Anthropic API request.
Args:
chat_options: The chat options that may contain additional beta flags.
Returns:
A set of beta flag strings to include in the request.
"""
return {
*BETA_FLAGS,
*self.additional_beta_flags,
*chat_options.additional_properties.get("additional_beta_flags", []),
}
def _prepare_messages_for_anthropic(self, messages: MutableSequence[ChatMessage]) -> list[dict[str, Any]]:
"""Prepare a list of ChatMessages for the Anthropic client.
This skips the first message if it is a system message,
as Anthropic expects system instructions as a separate parameter.
"""
# first system message is passed as instructions
if messages and isinstance(messages[0], ChatMessage) and messages[0].role == Role.SYSTEM:
return [self._convert_message_to_anthropic_format(msg) for msg in messages[1:]]
return [self._convert_message_to_anthropic_format(msg) for msg in messages]
return [self._prepare_message_for_anthropic(msg) for msg in messages[1:]]
return [self._prepare_message_for_anthropic(msg) for msg in messages]
def _convert_message_to_anthropic_format(self, message: ChatMessage) -> dict[str, Any]:
"""Convert a ChatMessage to the format expected by the Anthropic client.
def _prepare_message_for_anthropic(self, message: ChatMessage) -> dict[str, Any]:
"""Prepare a ChatMessage for the Anthropic client.
Args:
message: The ChatMessage to convert.
@@ -376,58 +391,96 @@ class AnthropicClient(BaseChatClient):
"content": a_content,
}
def _convert_tools_to_anthropic_format(
self, tools: list[ToolProtocol | MutableMapping[str, Any]] | None
) -> dict[str, Any] | None:
if not tools:
return None
tool_list: list[MutableMapping[str, Any]] = []
mcp_server_list: list[MutableMapping[str, Any]] = []
for tool in tools:
match tool:
case MutableMapping():
tool_list.append(tool)
case AIFunction():
tool_list.append({
"type": "custom",
"name": tool.name,
"description": tool.description,
"input_schema": tool.parameters(),
})
case HostedWebSearchTool():
search_tool: dict[str, Any] = {
"type": "web_search_20250305",
"name": "web_search",
}
if tool.additional_properties:
search_tool.update(tool.additional_properties)
tool_list.append(search_tool)
case HostedCodeInterpreterTool():
code_tool: dict[str, Any] = {
"type": "code_execution_20250825",
"name": "code_execution",
}
tool_list.append(code_tool)
case HostedMCPTool():
server_def: dict[str, Any] = {
"type": "url",
"name": tool.name,
"url": str(tool.url),
}
if tool.allowed_tools:
server_def["tool_configuration"] = {"allowed_tools": list(tool.allowed_tools)}
if tool.headers and (auth := tool.headers.get("authorization")):
server_def["authorization_token"] = auth
mcp_server_list.append(server_def)
case _:
logger.debug(f"Ignoring unsupported tool type: {type(tool)} for now")
def _prepare_tools_for_anthropic(self, chat_options: ChatOptions) -> dict[str, Any] | None:
"""Prepare tools and tool choice configuration for the Anthropic API request.
all_tools: dict[str, list[MutableMapping[str, Any]]] = {}
if tool_list:
all_tools["tools"] = tool_list
if mcp_server_list:
all_tools["mcp_servers"] = mcp_server_list
return all_tools
Args:
chat_options: The chat options containing tools and tool choice settings.
Returns:
A dictionary with tools, mcp_servers, and tool_choice configuration, or None if empty.
"""
result: dict[str, Any] = {}
# Process tools
if chat_options.tools:
tool_list: list[MutableMapping[str, Any]] = []
mcp_server_list: list[MutableMapping[str, Any]] = []
for tool in chat_options.tools:
match tool:
case MutableMapping():
tool_list.append(tool)
case AIFunction():
tool_list.append({
"type": "custom",
"name": tool.name,
"description": tool.description,
"input_schema": tool.parameters(),
})
case HostedWebSearchTool():
search_tool: dict[str, Any] = {
"type": "web_search_20250305",
"name": "web_search",
}
if tool.additional_properties:
search_tool.update(tool.additional_properties)
tool_list.append(search_tool)
case HostedCodeInterpreterTool():
code_tool: dict[str, Any] = {
"type": "code_execution_20250825",
"name": "code_execution",
}
tool_list.append(code_tool)
case HostedMCPTool():
server_def: dict[str, Any] = {
"type": "url",
"name": tool.name,
"url": str(tool.url),
}
if tool.allowed_tools:
server_def["tool_configuration"] = {"allowed_tools": list(tool.allowed_tools)}
if tool.headers and (auth := tool.headers.get("authorization")):
server_def["authorization_token"] = auth
mcp_server_list.append(server_def)
case _:
logger.debug(f"Ignoring unsupported tool type: {type(tool)} for now")
if tool_list:
result["tools"] = tool_list
if mcp_server_list:
result["mcp_servers"] = mcp_server_list
# Process tool choice
if chat_options.tool_choice is not None:
tool_choice_mode = (
chat_options.tool_choice if isinstance(chat_options.tool_choice, str) else chat_options.tool_choice.mode
)
match tool_choice_mode:
case "auto":
tool_choice: dict[str, Any] = {"type": "auto"}
if chat_options.allow_multiple_tool_calls is not None:
tool_choice["disable_parallel_tool_use"] = not chat_options.allow_multiple_tool_calls
result["tool_choice"] = tool_choice
case "required":
if (
not isinstance(chat_options.tool_choice, str)
and chat_options.tool_choice.required_function_name
):
tool_choice = {
"type": "tool",
"name": chat_options.tool_choice.required_function_name,
}
else:
tool_choice = {"type": "any"}
if chat_options.allow_multiple_tool_calls is not None:
tool_choice["disable_parallel_tool_use"] = not chat_options.allow_multiple_tool_calls
result["tool_choice"] = tool_choice
case "none":
result["tool_choice"] = {"type": "none"}
case _:
logger.debug(f"Ignoring unsupported tool choice mode: {tool_choice_mode} for now")
return result or None
# region Response Processing Methods
@@ -445,11 +498,11 @@ class AnthropicClient(BaseChatClient):
messages=[
ChatMessage(
role=Role.ASSISTANT,
contents=self._parse_message_contents(message.content),
contents=self._parse_contents_from_anthropic(message.content),
raw_representation=message,
)
],
usage_details=self._parse_message_usage(message.usage),
usage_details=self._parse_usage_from_anthropic(message.usage),
model_id=message.model,
finish_reason=FINISH_REASON_MAP.get(message.stop_reason) if message.stop_reason else None,
raw_response=message,
@@ -467,12 +520,12 @@ class AnthropicClient(BaseChatClient):
match event.type:
case "message_start":
usage_details: list[UsageContent] = []
if event.message.usage and (details := self._parse_message_usage(event.message.usage)):
if event.message.usage and (details := self._parse_usage_from_anthropic(event.message.usage)):
usage_details.append(UsageContent(details=details))
return ChatResponseUpdate(
response_id=event.message.id,
contents=[*self._parse_message_contents(event.message.content), *usage_details],
contents=[*self._parse_contents_from_anthropic(event.message.content), *usage_details],
model_id=event.message.model,
finish_reason=FINISH_REASON_MAP.get(event.message.stop_reason)
if event.message.stop_reason
@@ -480,7 +533,7 @@ class AnthropicClient(BaseChatClient):
raw_response=event,
)
case "message_delta":
usage = self._parse_message_usage(event.usage)
usage = self._parse_usage_from_anthropic(event.usage)
return ChatResponseUpdate(
contents=[UsageContent(details=usage, raw_representation=event.usage)] if usage else [],
raw_response=event,
@@ -488,13 +541,13 @@ class AnthropicClient(BaseChatClient):
case "message_stop":
logger.debug("Received message_stop event; no content to process.")
case "content_block_start":
contents = self._parse_message_contents([event.content_block])
contents = self._parse_contents_from_anthropic([event.content_block])
return ChatResponseUpdate(
contents=contents,
raw_response=event,
)
case "content_block_delta":
contents = self._parse_message_contents([event.delta])
contents = self._parse_contents_from_anthropic([event.delta])
return ChatResponseUpdate(
contents=contents,
raw_response=event,
@@ -505,7 +558,7 @@ class AnthropicClient(BaseChatClient):
logger.debug(f"Ignoring unsupported event type: {event.type}")
return None
def _parse_message_usage(self, usage: BetaUsage | BetaMessageDeltaUsage | None) -> UsageDetails | None:
def _parse_usage_from_anthropic(self, usage: BetaUsage | BetaMessageDeltaUsage | None) -> UsageDetails | None:
"""Parse usage details from the Anthropic message usage."""
if not usage:
return None
@@ -518,7 +571,7 @@ class AnthropicClient(BaseChatClient):
usage_details.additional_counts["anthropic.cache_read_input_tokens"] = usage.cache_read_input_tokens
return usage_details
def _parse_message_contents(
def _parse_contents_from_anthropic(
self, content: Sequence[BetaContentBlock | BetaRawContentBlockDelta | BetaTextBlock]
) -> list[Contents]:
"""Parse contents from the Anthropic message."""
@@ -530,7 +583,7 @@ class AnthropicClient(BaseChatClient):
TextContent(
text=content_block.text,
raw_representation=content_block,
annotations=self._parse_citations(content_block),
annotations=self._parse_citations_from_anthropic(content_block),
)
)
case "tool_use" | "mcp_tool_use" | "server_tool_use":
@@ -549,7 +602,7 @@ class AnthropicClient(BaseChatClient):
FunctionResultContent(
call_id=content_block.tool_use_id,
name=name if name and call_id == content_block.tool_use_id else "mcp_tool",
result=self._parse_message_contents(content_block.content)
result=self._parse_contents_from_anthropic(content_block.content)
if isinstance(content_block.content, list)
else content_block.content,
raw_representation=content_block,
@@ -608,7 +661,7 @@ class AnthropicClient(BaseChatClient):
logger.debug(f"Ignoring unsupported content type: {content_block.type} for now")
return contents
def _parse_citations(
def _parse_citations_from_anthropic(
self, content_block: BetaContentBlock | BetaRawContentBlockDelta | BetaTextBlock
) -> list[Annotations] | None:
content_citations = getattr(content_block, "citations", None)
@@ -151,12 +151,12 @@ def test_anthropic_client_service_url(mock_anthropic_client: MagicMock) -> None:
# Message Conversion Tests
def test_convert_message_to_anthropic_format_text(mock_anthropic_client: MagicMock) -> None:
def test_prepare_message_for_anthropic_text(mock_anthropic_client: MagicMock) -> None:
"""Test converting text message to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
message = ChatMessage(role=Role.USER, text="Hello, world!")
result = chat_client._convert_message_to_anthropic_format(message)
result = chat_client._prepare_message_for_anthropic(message)
assert result["role"] == "user"
assert len(result["content"]) == 1
@@ -164,7 +164,7 @@ def test_convert_message_to_anthropic_format_text(mock_anthropic_client: MagicMo
assert result["content"][0]["text"] == "Hello, world!"
def test_convert_message_to_anthropic_format_function_call(mock_anthropic_client: MagicMock) -> None:
def test_prepare_message_for_anthropic_function_call(mock_anthropic_client: MagicMock) -> None:
"""Test converting function call message to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
message = ChatMessage(
@@ -178,7 +178,7 @@ def test_convert_message_to_anthropic_format_function_call(mock_anthropic_client
],
)
result = chat_client._convert_message_to_anthropic_format(message)
result = chat_client._prepare_message_for_anthropic(message)
assert result["role"] == "assistant"
assert len(result["content"]) == 1
@@ -188,7 +188,7 @@ def test_convert_message_to_anthropic_format_function_call(mock_anthropic_client
assert result["content"][0]["input"] == {"location": "San Francisco"}
def test_convert_message_to_anthropic_format_function_result(mock_anthropic_client: MagicMock) -> None:
def test_prepare_message_for_anthropic_function_result(mock_anthropic_client: MagicMock) -> None:
"""Test converting function result message to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
message = ChatMessage(
@@ -202,7 +202,7 @@ def test_convert_message_to_anthropic_format_function_result(mock_anthropic_clie
],
)
result = chat_client._convert_message_to_anthropic_format(message)
result = chat_client._prepare_message_for_anthropic(message)
assert result["role"] == "user"
assert len(result["content"]) == 1
@@ -214,7 +214,7 @@ def test_convert_message_to_anthropic_format_function_result(mock_anthropic_clie
assert result["content"][0]["is_error"] is False
def test_convert_message_to_anthropic_format_text_reasoning(mock_anthropic_client: MagicMock) -> None:
def test_prepare_message_for_anthropic_text_reasoning(mock_anthropic_client: MagicMock) -> None:
"""Test converting text reasoning message to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
message = ChatMessage(
@@ -222,7 +222,7 @@ def test_convert_message_to_anthropic_format_text_reasoning(mock_anthropic_clien
contents=[TextReasoningContent(text="Let me think about this...")],
)
result = chat_client._convert_message_to_anthropic_format(message)
result = chat_client._prepare_message_for_anthropic(message)
assert result["role"] == "assistant"
assert len(result["content"]) == 1
@@ -230,7 +230,7 @@ def test_convert_message_to_anthropic_format_text_reasoning(mock_anthropic_clien
assert result["content"][0]["thinking"] == "Let me think about this..."
def test_convert_messages_to_anthropic_format_with_system(mock_anthropic_client: MagicMock) -> None:
def test_prepare_messages_for_anthropic_with_system(mock_anthropic_client: MagicMock) -> None:
"""Test converting messages list with system message."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [
@@ -238,7 +238,7 @@ def test_convert_messages_to_anthropic_format_with_system(mock_anthropic_client:
ChatMessage(role=Role.USER, text="Hello!"),
]
result = chat_client._convert_messages_to_anthropic_format(messages)
result = chat_client._prepare_messages_for_anthropic(messages)
# System message should be skipped
assert len(result) == 1
@@ -246,7 +246,7 @@ def test_convert_messages_to_anthropic_format_with_system(mock_anthropic_client:
assert result[0]["content"][0]["text"] == "Hello!"
def test_convert_messages_to_anthropic_format_without_system(mock_anthropic_client: MagicMock) -> None:
def test_prepare_messages_for_anthropic_without_system(mock_anthropic_client: MagicMock) -> None:
"""Test converting messages list without system message."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [
@@ -254,7 +254,7 @@ def test_convert_messages_to_anthropic_format_without_system(mock_anthropic_clie
ChatMessage(role=Role.ASSISTANT, text="Hi there!"),
]
result = chat_client._convert_messages_to_anthropic_format(messages)
result = chat_client._prepare_messages_for_anthropic(messages)
assert len(result) == 2
assert result[0]["role"] == "user"
@@ -264,7 +264,7 @@ def test_convert_messages_to_anthropic_format_without_system(mock_anthropic_clie
# Tool Conversion Tests
def test_convert_tools_to_anthropic_format_ai_function(mock_anthropic_client: MagicMock) -> None:
def test_prepare_tools_for_anthropic_ai_function(mock_anthropic_client: MagicMock) -> None:
"""Test converting AIFunction to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
@@ -273,9 +273,8 @@ def test_convert_tools_to_anthropic_format_ai_function(mock_anthropic_client: Ma
"""Get weather for a location."""
return f"Weather for {location}"
tools = [get_weather]
result = chat_client._convert_tools_to_anthropic_format(tools)
chat_options = ChatOptions(tools=[get_weather])
result = chat_client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "tools" in result
@@ -285,12 +284,12 @@ def test_convert_tools_to_anthropic_format_ai_function(mock_anthropic_client: Ma
assert "Get weather for a location" in result["tools"][0]["description"]
def test_convert_tools_to_anthropic_format_web_search(mock_anthropic_client: MagicMock) -> None:
def test_prepare_tools_for_anthropic_web_search(mock_anthropic_client: MagicMock) -> None:
"""Test converting HostedWebSearchTool to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
tools = [HostedWebSearchTool()]
chat_options = ChatOptions(tools=[HostedWebSearchTool()])
result = chat_client._convert_tools_to_anthropic_format(tools)
result = chat_client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "tools" in result
@@ -299,12 +298,12 @@ def test_convert_tools_to_anthropic_format_web_search(mock_anthropic_client: Mag
assert result["tools"][0]["name"] == "web_search"
def test_convert_tools_to_anthropic_format_code_interpreter(mock_anthropic_client: MagicMock) -> None:
def test_prepare_tools_for_anthropic_code_interpreter(mock_anthropic_client: MagicMock) -> None:
"""Test converting HostedCodeInterpreterTool to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
tools = [HostedCodeInterpreterTool()]
chat_options = ChatOptions(tools=[HostedCodeInterpreterTool()])
result = chat_client._convert_tools_to_anthropic_format(tools)
result = chat_client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "tools" in result
@@ -313,12 +312,12 @@ def test_convert_tools_to_anthropic_format_code_interpreter(mock_anthropic_clien
assert result["tools"][0]["name"] == "code_execution"
def test_convert_tools_to_anthropic_format_mcp_tool(mock_anthropic_client: MagicMock) -> None:
def test_prepare_tools_for_anthropic_mcp_tool(mock_anthropic_client: MagicMock) -> None:
"""Test converting HostedMCPTool to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
tools = [HostedMCPTool(name="test-mcp", url="https://example.com/mcp")]
chat_options = ChatOptions(tools=[HostedMCPTool(name="test-mcp", url="https://example.com/mcp")])
result = chat_client._convert_tools_to_anthropic_format(tools)
result = chat_client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "mcp_servers" in result
@@ -328,18 +327,20 @@ def test_convert_tools_to_anthropic_format_mcp_tool(mock_anthropic_client: Magic
assert result["mcp_servers"][0]["url"] == "https://example.com/mcp"
def test_convert_tools_to_anthropic_format_mcp_with_auth(mock_anthropic_client: MagicMock) -> None:
def test_prepare_tools_for_anthropic_mcp_with_auth(mock_anthropic_client: MagicMock) -> None:
"""Test converting HostedMCPTool with authorization headers."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
tools = [
HostedMCPTool(
name="test-mcp",
url="https://example.com/mcp",
headers={"authorization": "Bearer token123"},
)
]
chat_options = ChatOptions(
tools=[
HostedMCPTool(
name="test-mcp",
url="https://example.com/mcp",
headers={"authorization": "Bearer token123"},
)
]
)
result = chat_client._convert_tools_to_anthropic_format(tools)
result = chat_client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "mcp_servers" in result
@@ -348,12 +349,12 @@ def test_convert_tools_to_anthropic_format_mcp_with_auth(mock_anthropic_client:
assert result["mcp_servers"][0]["authorization_token"] == "Bearer token123"
def test_convert_tools_to_anthropic_format_dict_tool(mock_anthropic_client: MagicMock) -> None:
def test_prepare_tools_for_anthropic_dict_tool(mock_anthropic_client: MagicMock) -> None:
"""Test converting dict tool to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
tools = [{"type": "custom", "name": "custom_tool", "description": "A custom tool"}]
chat_options = ChatOptions(tools=[{"type": "custom", "name": "custom_tool", "description": "A custom tool"}])
result = chat_client._convert_tools_to_anthropic_format(tools)
result = chat_client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "tools" in result
@@ -361,11 +362,12 @@ def test_convert_tools_to_anthropic_format_dict_tool(mock_anthropic_client: Magi
assert result["tools"][0]["name"] == "custom_tool"
def test_convert_tools_to_anthropic_format_none(mock_anthropic_client: MagicMock) -> None:
def test_prepare_tools_for_anthropic_none(mock_anthropic_client: MagicMock) -> None:
"""Test converting None tools."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
chat_options = ChatOptions()
result = chat_client._convert_tools_to_anthropic_format(None)
result = chat_client._prepare_tools_for_anthropic(chat_options)
assert result is None
@@ -373,14 +375,14 @@ def test_convert_tools_to_anthropic_format_none(mock_anthropic_client: MagicMock
# Run Options Tests
async def test_create_run_options_basic(mock_anthropic_client: MagicMock) -> None:
"""Test _create_run_options with basic ChatOptions."""
async def test_prepare_options_basic(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with basic ChatOptions."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role=Role.USER, text="Hello")]
chat_options = ChatOptions(max_tokens=100, temperature=0.7)
run_options = chat_client._create_run_options(messages, chat_options)
run_options = chat_client._prepare_options(messages, chat_options)
assert run_options["model"] == chat_client.model_id
assert run_options["max_tokens"] == 100
@@ -388,8 +390,8 @@ async def test_create_run_options_basic(mock_anthropic_client: MagicMock) -> Non
assert "messages" in run_options
async def test_create_run_options_with_system_message(mock_anthropic_client: MagicMock) -> None:
"""Test _create_run_options with system message."""
async def test_prepare_options_with_system_message(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with system message."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [
@@ -398,52 +400,52 @@ async def test_create_run_options_with_system_message(mock_anthropic_client: Mag
]
chat_options = ChatOptions()
run_options = chat_client._create_run_options(messages, chat_options)
run_options = chat_client._prepare_options(messages, chat_options)
assert run_options["system"] == "You are helpful."
assert len(run_options["messages"]) == 1 # System message not in messages list
async def test_create_run_options_with_tool_choice_auto(mock_anthropic_client: MagicMock) -> None:
"""Test _create_run_options with auto tool choice."""
async def test_prepare_options_with_tool_choice_auto(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with auto tool choice."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role=Role.USER, text="Hello")]
chat_options = ChatOptions(tool_choice="auto")
run_options = chat_client._create_run_options(messages, chat_options)
run_options = chat_client._prepare_options(messages, chat_options)
assert run_options["tool_choice"]["type"] == "auto"
async def test_create_run_options_with_tool_choice_required(mock_anthropic_client: MagicMock) -> None:
"""Test _create_run_options with required tool choice."""
async def test_prepare_options_with_tool_choice_required(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with required tool choice."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role=Role.USER, text="Hello")]
# For required with specific function, need to pass as dict
chat_options = ChatOptions(tool_choice={"mode": "required", "required_function_name": "get_weather"})
run_options = chat_client._create_run_options(messages, chat_options)
run_options = chat_client._prepare_options(messages, chat_options)
assert run_options["tool_choice"]["type"] == "tool"
assert run_options["tool_choice"]["name"] == "get_weather"
async def test_create_run_options_with_tool_choice_none(mock_anthropic_client: MagicMock) -> None:
"""Test _create_run_options with none tool choice."""
async def test_prepare_options_with_tool_choice_none(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with none tool choice."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role=Role.USER, text="Hello")]
chat_options = ChatOptions(tool_choice="none")
run_options = chat_client._create_run_options(messages, chat_options)
run_options = chat_client._prepare_options(messages, chat_options)
assert run_options["tool_choice"]["type"] == "none"
async def test_create_run_options_with_tools(mock_anthropic_client: MagicMock) -> None:
"""Test _create_run_options with tools."""
async def test_prepare_options_with_tools(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with tools."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
@ai_function
@@ -454,32 +456,32 @@ async def test_create_run_options_with_tools(mock_anthropic_client: MagicMock) -
messages = [ChatMessage(role=Role.USER, text="Hello")]
chat_options = ChatOptions(tools=[get_weather])
run_options = chat_client._create_run_options(messages, chat_options)
run_options = chat_client._prepare_options(messages, chat_options)
assert "tools" in run_options
assert len(run_options["tools"]) == 1
async def test_create_run_options_with_stop_sequences(mock_anthropic_client: MagicMock) -> None:
"""Test _create_run_options with stop sequences."""
async def test_prepare_options_with_stop_sequences(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with stop sequences."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role=Role.USER, text="Hello")]
chat_options = ChatOptions(stop=["STOP", "END"])
run_options = chat_client._create_run_options(messages, chat_options)
run_options = chat_client._prepare_options(messages, chat_options)
assert run_options["stop_sequences"] == ["STOP", "END"]
async def test_create_run_options_with_top_p(mock_anthropic_client: MagicMock) -> None:
"""Test _create_run_options with top_p."""
async def test_prepare_options_with_top_p(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with top_p."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role=Role.USER, text="Hello")]
chat_options = ChatOptions(top_p=0.9)
run_options = chat_client._create_run_options(messages, chat_options)
run_options = chat_client._prepare_options(messages, chat_options)
assert run_options["top_p"] == 0.9
@@ -540,41 +542,41 @@ def test_process_message_with_tool_use(mock_anthropic_client: MagicMock) -> None
assert response.finish_reason == FinishReason.TOOL_CALLS
def test_parse_message_usage_basic(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_message_usage with basic usage."""
def test_parse_usage_from_anthropic_basic(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_usage_from_anthropic with basic usage."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
usage = BetaUsage(input_tokens=10, output_tokens=5)
result = chat_client._parse_message_usage(usage)
result = chat_client._parse_usage_from_anthropic(usage)
assert result is not None
assert result.input_token_count == 10
assert result.output_token_count == 5
def test_parse_message_usage_none(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_message_usage with None usage."""
def test_parse_usage_from_anthropic_none(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_usage_from_anthropic with None usage."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
result = chat_client._parse_message_usage(None)
result = chat_client._parse_usage_from_anthropic(None)
assert result is None
def test_parse_message_contents_text(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_message_contents with text content."""
def test_parse_contents_from_anthropic_text(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_contents_from_anthropic with text content."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
content = [BetaTextBlock(type="text", text="Hello!")]
result = chat_client._parse_message_contents(content)
result = chat_client._parse_contents_from_anthropic(content)
assert len(result) == 1
assert isinstance(result[0], TextContent)
assert result[0].text == "Hello!"
def test_parse_message_contents_tool_use(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_message_contents with tool use."""
def test_parse_contents_from_anthropic_tool_use(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_contents_from_anthropic with tool use."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
content = [
@@ -585,7 +587,7 @@ def test_parse_message_contents_tool_use(mock_anthropic_client: MagicMock) -> No
input={"location": "SF"},
)
]
result = chat_client._parse_message_contents(content)
result = chat_client._parse_contents_from_anthropic(content)
assert len(result) == 1
assert isinstance(result[0], FunctionCallContent)