133 lines
5.1 KiB
Python
133 lines
5.1 KiB
Python
|
|
from typing_extensions import TypedDict
|
||
|
|
from pydantic import BaseModel
|
||
|
|
from instructor.processing.response import handle_response_model
|
||
|
|
from instructor.providers.bedrock.utils import _prepare_bedrock_converse_kwargs_internal
|
||
|
|
|
||
|
|
|
||
|
|
def test_typed_dict_conversion() -> None:
|
||
|
|
class User(TypedDict): # type: ignore
|
||
|
|
name: str
|
||
|
|
age: int
|
||
|
|
|
||
|
|
_, user_tool_definition = handle_response_model(User)
|
||
|
|
|
||
|
|
class User(BaseModel):
|
||
|
|
name: str
|
||
|
|
age: int
|
||
|
|
|
||
|
|
_, pydantic_user_tool_definition = handle_response_model(User)
|
||
|
|
assert user_tool_definition == pydantic_user_tool_definition
|
||
|
|
|
||
|
|
|
||
|
|
def test_openai_to_bedrock_conversion() -> None:
|
||
|
|
"""OpenAI-style input should be fully converted to Bedrock format."""
|
||
|
|
call_kwargs = {
|
||
|
|
"model": "anthropic.claude-3-haiku-20240307-v1:0",
|
||
|
|
"messages": [
|
||
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
||
|
|
{"role": "user", "content": "Extract: Jason is 22 years old"},
|
||
|
|
{"role": "assistant", "content": "Sure! Jason is 22."},
|
||
|
|
],
|
||
|
|
}
|
||
|
|
result = _prepare_bedrock_converse_kwargs_internal(call_kwargs)
|
||
|
|
assert "model" not in result
|
||
|
|
assert result["modelId"] == "anthropic.claude-3-haiku-20240307-v1:0"
|
||
|
|
assert result["system"] == [{"text": "You are a helpful assistant."}]
|
||
|
|
assert len(result["messages"]) == 2
|
||
|
|
assert result["messages"][0]["role"] == "user"
|
||
|
|
assert result["messages"][0]["content"] == [
|
||
|
|
{"text": "Extract: Jason is 22 years old"}
|
||
|
|
]
|
||
|
|
assert result["messages"][1]["role"] == "assistant"
|
||
|
|
assert result["messages"][1]["content"] == [{"text": "Sure! Jason is 22."}]
|
||
|
|
|
||
|
|
|
||
|
|
def test_bedrock_native_preserved() -> None:
|
||
|
|
"""Bedrock-native input should be preserved as-is."""
|
||
|
|
call_kwargs = {
|
||
|
|
"modelId": "anthropic.claude-3-haiku-20240307-v1:0",
|
||
|
|
"system": [{"text": "You are a helpful assistant."}],
|
||
|
|
"messages": [
|
||
|
|
{"role": "user", "content": [{"text": "Extract: Jason is 22 years old"}]},
|
||
|
|
{"role": "assistant", "content": [{"text": "Sure! Jason is 22."}]},
|
||
|
|
],
|
||
|
|
}
|
||
|
|
result = _prepare_bedrock_converse_kwargs_internal(call_kwargs)
|
||
|
|
assert result["system"] == [{"text": "You are a helpful assistant."}]
|
||
|
|
assert len(result["messages"]) == 2
|
||
|
|
assert result["messages"][0]["content"] == [
|
||
|
|
{"text": "Extract: Jason is 22 years old"}
|
||
|
|
]
|
||
|
|
assert result["messages"][1]["content"] == [{"text": "Sure! Jason is 22."}]
|
||
|
|
|
||
|
|
|
||
|
|
def test_mixed_openai_and_bedrock() -> None:
|
||
|
|
"""Mixed input: OpenAI-style is converted, Bedrock-native is preserved."""
|
||
|
|
call_kwargs = {
|
||
|
|
"modelId": "anthropic.claude-3-haiku-20240307-v1:0",
|
||
|
|
"system": [{"text": "You are a helpful assistant."}],
|
||
|
|
"messages": [
|
||
|
|
{
|
||
|
|
"role": "user",
|
||
|
|
"content": "Extract: Jason is 22 years old",
|
||
|
|
}, # OpenAI style
|
||
|
|
{
|
||
|
|
"role": "assistant",
|
||
|
|
"content": [{"text": "Sure! Jason is 22."}],
|
||
|
|
}, # Bedrock style
|
||
|
|
],
|
||
|
|
}
|
||
|
|
result = _prepare_bedrock_converse_kwargs_internal(call_kwargs)
|
||
|
|
assert result["modelId"] == "anthropic.claude-3-haiku-20240307-v1:0"
|
||
|
|
assert result["system"] == [{"text": "You are a helpful assistant."}]
|
||
|
|
assert len(result["messages"]) == 2
|
||
|
|
# OpenAI-style user message converted
|
||
|
|
assert result["modelId"] == "anthropic.claude-3-haiku-20240307-v1:0"
|
||
|
|
assert result["messages"][0]["content"] == [
|
||
|
|
{"text": "Extract: Jason is 22 years old"}
|
||
|
|
]
|
||
|
|
# Bedrock-style assistant message preserved
|
||
|
|
assert result["messages"][1]["content"] == [{"text": "Sure! Jason is 22."}]
|
||
|
|
|
||
|
|
|
||
|
|
def test_bedrock_round_trip() -> None:
|
||
|
|
"""Bedrock input should be unchanged after round-trip through the function."""
|
||
|
|
call_kwargs = {
|
||
|
|
"modelId": "anthropic.claude-3-haiku-20240307-v1:0",
|
||
|
|
"system": [{"text": "Bedrock system."}],
|
||
|
|
"messages": [
|
||
|
|
{"role": "user", "content": [{"text": "Bedrock user message."}]},
|
||
|
|
],
|
||
|
|
}
|
||
|
|
import copy
|
||
|
|
|
||
|
|
original = copy.deepcopy(call_kwargs)
|
||
|
|
result = _prepare_bedrock_converse_kwargs_internal(call_kwargs)
|
||
|
|
assert result == original
|
||
|
|
|
||
|
|
|
||
|
|
def test_empty_and_missing_content() -> None:
|
||
|
|
"""Empty messages and missing content should be handled gracefully."""
|
||
|
|
# Empty messages
|
||
|
|
call_kwargs = {"messages": []}
|
||
|
|
result = _prepare_bedrock_converse_kwargs_internal(call_kwargs)
|
||
|
|
assert result["messages"] == []
|
||
|
|
# Message with no content
|
||
|
|
call_kwargs = {"messages": [{"role": "user"}]}
|
||
|
|
result = _prepare_bedrock_converse_kwargs_internal(call_kwargs)
|
||
|
|
assert result["messages"][0]["role"] == "user"
|
||
|
|
# Should not add a content key if not present
|
||
|
|
assert "content" not in result["messages"][0]
|
||
|
|
|
||
|
|
|
||
|
|
def test_bedrock_invalid_content_format() -> None:
|
||
|
|
"""Invalid content types should raise ValueError."""
|
||
|
|
call_kwargs = {
|
||
|
|
"messages": [{"role": "user", "content": 12345}] # Invalid content type
|
||
|
|
}
|
||
|
|
try:
|
||
|
|
_prepare_bedrock_converse_kwargs_internal(call_kwargs)
|
||
|
|
raise AssertionError("Should have raised ValueError")
|
||
|
|
except ValueError as e:
|
||
|
|
assert "Unsupported message content type for Bedrock" in str(e)
|