Ayaan Gazali
557 lines
22 KiB
Python
557 lines
22 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from mem0.configs.llms.aws_bedrock import AWSBedrockConfig
|
|
from mem0.llms.aws_bedrock import AWSBedrockLLM, extract_provider
|
|
from mem0.utils.factory import LlmFactory
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.fixture
|
|
def mock_boto3():
|
|
"""Patch boto3 so no real AWS calls are made during unit tests."""
|
|
with patch("mem0.llms.aws_bedrock.boto3") as mock_b3:
|
|
runtime_client = MagicMock()
|
|
bedrock_client = MagicMock()
|
|
bedrock_client.list_foundation_models.return_value = {"modelSummaries": []}
|
|
|
|
def _client(service, **kwargs):
|
|
if service == "bedrock-runtime":
|
|
return runtime_client
|
|
return bedrock_client
|
|
|
|
mock_b3.client.side_effect = _client
|
|
yield runtime_client
|
|
|
|
|
|
def _make_llm(model: str, mock_boto3, **kwargs) -> AWSBedrockLLM:
|
|
"""Instantiate AWSBedrockLLM with a given model, all AWS calls mocked."""
|
|
config = AWSBedrockConfig(model=model, **kwargs)
|
|
return AWSBedrockLLM(config)
|
|
|
|
|
|
def _converse_response(text: str = "ok") -> dict:
|
|
"""Minimal Converse API response dict."""
|
|
return {"output": {"message": {"content": [{"text": text}]}}}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# extract_provider
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestExtractProvider:
|
|
def test_standard_anthropic_model(self):
|
|
assert extract_provider("anthropic.claude-3-5-sonnet-20240620-v1:0") == "anthropic"
|
|
|
|
def test_inference_profile_us_prefix(self):
|
|
# Cross-region inference profile IDs look like us.anthropic.<model>
|
|
assert extract_provider("us.anthropic.claude-haiku-4-5-20251001-v1:0") == "anthropic"
|
|
|
|
def test_inference_profile_eu_prefix(self):
|
|
assert extract_provider("eu.anthropic.claude-sonnet-4-5-20250929-v1:0") == "anthropic"
|
|
|
|
def test_inference_profile_ap_prefix(self):
|
|
assert extract_provider("ap.anthropic.claude-3-opus-20240229-v1:0") == "anthropic"
|
|
|
|
def test_amazon_model(self):
|
|
assert extract_provider("amazon.nova-3-mini-20241119-v1:0") == "amazon"
|
|
|
|
def test_meta_model(self):
|
|
assert extract_provider("meta.llama3-8b-instruct-v1:0") == "meta"
|
|
|
|
def test_mistral_model(self):
|
|
assert extract_provider("mistral.mistral-7b-instruct-v0:2") == "mistral"
|
|
|
|
def test_unknown_model_raises(self):
|
|
with pytest.raises(ValueError, match="Unknown provider"):
|
|
extract_provider("unknown-vendor.some-model-v1:0")
|
|
|
|
def test_application_inference_profile_arn_without_override_raises(self):
|
|
arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"
|
|
with pytest.raises(ValueError, match="Unknown provider"):
|
|
extract_provider(arn)
|
|
|
|
def test_application_inference_profile_arn_with_explicit_provider(self):
|
|
arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"
|
|
assert extract_provider(arn, "anthropic") == "anthropic"
|
|
|
|
def test_explicit_provider_takes_precedence_over_regex(self):
|
|
assert extract_provider("anthropic.claude-3-5-sonnet-20240620-v1:0", "amazon") == "amazon"
|
|
|
|
def test_explicit_provider_typo_raises(self):
|
|
arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"
|
|
with pytest.raises(ValueError, match="Unknown provider_override 'anthorpic'"):
|
|
extract_provider(arn, "anthorpic")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AWSBedrockConfig
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestAWSBedrockConfig:
|
|
def test_top_p_defaults_to_none(self):
|
|
config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0")
|
|
assert config.top_p is None
|
|
|
|
def test_top_p_explicit_value_stored(self):
|
|
config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0", top_p=0.8)
|
|
assert config.top_p == 0.8
|
|
|
|
def test_get_model_config_excludes_top_p_by_default(self):
|
|
config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0", temperature=0.5)
|
|
model_cfg = config.get_model_config()
|
|
assert "top_p" not in model_cfg
|
|
|
|
def test_get_model_config_includes_top_p_when_set(self):
|
|
config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0", top_p=0.7)
|
|
model_cfg = config.get_model_config()
|
|
assert model_cfg["top_p"] == 0.7
|
|
|
|
def test_get_model_config_top_p_via_model_kwargs(self):
|
|
"""model_kwargs can supply top_p after merge; same semantics as explicit top_p."""
|
|
config = AWSBedrockConfig(
|
|
model="anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
model_kwargs={"top_p": 0.88},
|
|
)
|
|
assert config.get_model_config()["top_p"] == 0.88
|
|
|
|
def test_get_model_config_always_includes_temperature(self):
|
|
config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0", temperature=0.3)
|
|
model_cfg = config.get_model_config()
|
|
assert model_cfg["temperature"] == 0.3
|
|
|
|
def test_aws_region_stored(self):
|
|
config = AWSBedrockConfig(
|
|
model="anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
aws_region="us-east-2",
|
|
)
|
|
assert config.aws_region == "us-east-2"
|
|
|
|
def test_provider_override_defaults_to_none(self):
|
|
config = AWSBedrockConfig(model="anthropic.claude-3-5-sonnet-20240620-v1:0")
|
|
assert config.provider_override is None
|
|
|
|
def test_provider_override_stored(self):
|
|
arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"
|
|
config = AWSBedrockConfig(model=arn, provider_override="anthropic")
|
|
assert config.provider_override == "anthropic"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AWSBedrockLLM with application inference profile ARNs
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestApplicationInferenceProfileArn:
|
|
ARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz"
|
|
|
|
def test_arn_without_provider_override_raises(self, mock_boto3):
|
|
with pytest.raises(ValueError, match="Unknown provider"):
|
|
_make_llm(self.ARN, mock_boto3)
|
|
|
|
def test_arn_with_provider_override_resolves(self, mock_boto3):
|
|
llm = _make_llm(self.ARN, mock_boto3, provider_override="anthropic")
|
|
assert llm.provider == "anthropic"
|
|
assert llm.supports_tools is True
|
|
|
|
def test_plain_model_id_unaffected(self, mock_boto3):
|
|
llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)
|
|
assert llm.provider == "anthropic"
|
|
|
|
def test_cross_region_inference_profile_unaffected(self, mock_boto3):
|
|
llm = _make_llm("us.anthropic.claude-haiku-4-5-20251001-v1:0", mock_boto3)
|
|
assert llm.provider == "anthropic"
|
|
|
|
def test_arn_with_misspelled_provider_override_raises(self, mock_boto3):
|
|
with pytest.raises(ValueError, match="Unknown provider_override 'anthorpic'"):
|
|
_make_llm(self.ARN, mock_boto3, provider_override="anthorpic")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LlmFactory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestLlmFactory:
|
|
def test_aws_bedrock_uses_aws_bedrock_config(self):
|
|
_, config_class = LlmFactory.provider_to_class["aws_bedrock"]
|
|
assert config_class is AWSBedrockConfig
|
|
|
|
def test_factory_create_accepts_aws_region(self, mock_boto3):
|
|
"""LlmFactory.create must not crash when aws_region is in the config dict.
|
|
|
|
Before the fix, the factory mapped aws_bedrock to BaseLlmConfig which has
|
|
no aws_region parameter, causing: TypeError: __init__() got an unexpected
|
|
keyword argument 'aws_region'.
|
|
"""
|
|
user_dict = {
|
|
"model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
"aws_region": "us-east-2",
|
|
"temperature": 0.1,
|
|
"max_tokens": 2000,
|
|
}
|
|
mock_boto3.converse.return_value = _converse_response()
|
|
|
|
# This must not raise TypeError
|
|
llm = LlmFactory.create("aws_bedrock", user_dict)
|
|
|
|
assert isinstance(llm.config, AWSBedrockConfig)
|
|
assert llm.config.aws_region == "us-east-2"
|
|
assert llm.config.top_p is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _build_inference_config
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestBuildInferenceConfig:
|
|
"""
|
|
Unit tests for the _build_inference_config helper.
|
|
Validates the exact keys present in the returned dict.
|
|
"""
|
|
|
|
def test_anthropic_only_temperature_by_default(self, mock_boto3):
|
|
llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3, temperature=0.5)
|
|
cfg = llm._build_inference_config()
|
|
assert "temperature" in cfg
|
|
assert cfg["temperature"] == 0.5
|
|
assert "topP" not in cfg, "topP must be absent when top_p not configured"
|
|
|
|
def test_anthropic_top_p_explicitly_set_still_omits_top_p(self, mock_boto3):
|
|
# Anthropic rejects both; topP must still be omitted even when set
|
|
llm = _make_llm(
|
|
"anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
mock_boto3,
|
|
temperature=0.5,
|
|
top_p=0.9,
|
|
)
|
|
cfg = llm._build_inference_config()
|
|
assert "temperature" in cfg
|
|
assert "topP" not in cfg
|
|
|
|
def test_anthropic_inference_profile_omits_top_p(self, mock_boto3):
|
|
# Cross-region inference profiles (us.anthropic.*) follow the same rule
|
|
llm = _make_llm("us.anthropic.claude-haiku-4-5-20251001-v1:0", mock_boto3, temperature=0.1)
|
|
cfg = llm._build_inference_config()
|
|
assert "topP" not in cfg
|
|
|
|
def test_amazon_includes_top_p_when_set(self, mock_boto3):
|
|
llm = _make_llm(
|
|
"amazon.nova-3-mini-20241119-v1:0",
|
|
mock_boto3,
|
|
temperature=0.5,
|
|
top_p=0.85,
|
|
)
|
|
cfg = llm._build_inference_config()
|
|
assert cfg["temperature"] == 0.5
|
|
assert cfg["topP"] == 0.85
|
|
|
|
def test_amazon_omits_top_p_when_not_set(self, mock_boto3):
|
|
llm = _make_llm("amazon.nova-3-mini-20241119-v1:0", mock_boto3, temperature=0.5)
|
|
cfg = llm._build_inference_config()
|
|
assert "topP" not in cfg
|
|
|
|
def test_max_tokens_present(self, mock_boto3):
|
|
llm = _make_llm(
|
|
"anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
mock_boto3,
|
|
max_tokens=1024,
|
|
)
|
|
cfg = llm._build_inference_config()
|
|
assert cfg["maxTokens"] == 1024
|
|
|
|
def test_nova_fallback_max_tokens_when_absent(self, mock_boto3):
|
|
"""Legacy Nova Converse used 5000 when max_tokens was missing from the dict."""
|
|
llm = _make_llm("amazon.nova-3-mini-20241119-v1:0", mock_boto3)
|
|
llm.model_config.pop("max_tokens", None)
|
|
cfg = llm._build_inference_config()
|
|
assert cfg["maxTokens"] == 5000
|
|
|
|
def test_anthropic_fallback_max_tokens_when_absent(self, mock_boto3):
|
|
llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)
|
|
llm.model_config.pop("max_tokens", None)
|
|
cfg = llm._build_inference_config()
|
|
assert cfg["maxTokens"] == 2000
|
|
|
|
def test_minimax_omits_top_p_when_explicitly_set(self, mock_boto3):
|
|
# MiniMax M2.x (reasoning model) rejects both temperature and topP simultaneously.
|
|
# Even when the user explicitly configures top_p, it must be omitted.
|
|
llm = _make_llm(
|
|
"minimax.minimax-m2.5",
|
|
mock_boto3,
|
|
temperature=0.1,
|
|
top_p=0.9,
|
|
)
|
|
cfg = llm._build_inference_config()
|
|
assert "temperature" in cfg
|
|
assert "topP" not in cfg, "topP must be absent for MiniMax reasoning models"
|
|
|
|
def test_minimax_only_temperature_by_default(self, mock_boto3):
|
|
llm = _make_llm("minimax.minimax-m2.5", mock_boto3, temperature=0.1)
|
|
cfg = llm._build_inference_config()
|
|
assert cfg["temperature"] == 0.1
|
|
assert "topP" not in cfg
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# generate_response — Converse API call assertions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
MESSAGES = [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "Hello"},
|
|
]
|
|
|
|
TOOLS = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "add_memory",
|
|
"description": "Store a memory",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"data": {"type": "string"}},
|
|
"required": ["data"],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
|
|
|
|
class TestGenerateResponseConverse:
|
|
def test_standard_anthropic_no_top_p_in_converse_call(self, mock_boto3):
|
|
mock_boto3.converse.return_value = _converse_response()
|
|
llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3, temperature=0.2)
|
|
|
|
llm.generate_response(MESSAGES)
|
|
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
inference_cfg = kwargs["inferenceConfig"]
|
|
assert "topP" not in inference_cfg
|
|
assert inference_cfg["temperature"] == 0.2
|
|
|
|
def test_standard_anthropic_inference_profile_no_top_p(self, mock_boto3):
|
|
mock_boto3.converse.return_value = _converse_response()
|
|
llm = _make_llm("us.anthropic.claude-haiku-4-5-20251001-v1:0", mock_boto3, temperature=0.1)
|
|
|
|
llm.generate_response(MESSAGES)
|
|
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
assert "topP" not in kwargs["inferenceConfig"]
|
|
|
|
def test_with_tools_anthropic_no_top_p(self, mock_boto3):
|
|
mock_boto3.converse.return_value = _converse_response()
|
|
llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3, temperature=0.2)
|
|
|
|
llm.generate_response(MESSAGES, tools=TOOLS)
|
|
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
assert "topP" not in kwargs["inferenceConfig"]
|
|
|
|
def test_with_tools_anthropic_top_p_set_still_omitted(self, mock_boto3):
|
|
# Even if user explicitly sets top_p, Anthropic must not receive topP
|
|
mock_boto3.converse.return_value = _converse_response()
|
|
llm = _make_llm(
|
|
"anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
mock_boto3,
|
|
temperature=0.2,
|
|
top_p=0.9,
|
|
)
|
|
|
|
llm.generate_response(MESSAGES, tools=TOOLS)
|
|
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
assert "topP" not in kwargs["inferenceConfig"]
|
|
|
|
def test_nova_includes_top_p_when_explicitly_set(self, mock_boto3):
|
|
mock_boto3.converse.return_value = _converse_response()
|
|
llm = _make_llm(
|
|
"amazon.nova-3-mini-20241119-v1:0",
|
|
mock_boto3,
|
|
temperature=0.5,
|
|
top_p=0.85,
|
|
)
|
|
|
|
llm.generate_response(MESSAGES)
|
|
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
assert kwargs["inferenceConfig"]["topP"] == 0.85
|
|
|
|
def test_nova_omits_top_p_when_not_set(self, mock_boto3):
|
|
mock_boto3.converse.return_value = _converse_response()
|
|
llm = _make_llm("amazon.nova-3-mini-20241119-v1:0", mock_boto3, temperature=0.5)
|
|
|
|
llm.generate_response(MESSAGES)
|
|
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
assert "topP" not in kwargs["inferenceConfig"]
|
|
|
|
def test_anthropic_model_kwargs_top_p_still_omits_top_p_in_converse(self, mock_boto3):
|
|
"""top_p injected via model_kwargs must not add topP for Anthropic Converse."""
|
|
mock_boto3.converse.return_value = _converse_response()
|
|
llm = _make_llm(
|
|
"anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
mock_boto3,
|
|
temperature=0.2,
|
|
model_kwargs={"top_p": 0.88},
|
|
)
|
|
|
|
llm.generate_response(MESSAGES)
|
|
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
assert "topP" not in kwargs["inferenceConfig"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MiniMax provider
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMiniMaxProvider:
|
|
"""Tests for MiniMax models via Bedrock Converse API."""
|
|
|
|
def test_extract_provider(self):
|
|
assert extract_provider("minimax.minimax-m2.5") == "minimax"
|
|
assert extract_provider("minimax.minimax-m2") == "minimax"
|
|
|
|
def test_generate_response_text_only(self, mock_boto3):
|
|
"""Standard response: single text block."""
|
|
mock_boto3.converse.return_value = _converse_response("Hello!")
|
|
llm = _make_llm("minimax.minimax-m2.5", mock_boto3)
|
|
|
|
result = llm.generate_response([{"role": "user", "content": "say hi"}])
|
|
|
|
assert result == "Hello!"
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
assert kwargs["modelId"] == "minimax.minimax-m2.5"
|
|
assert kwargs["messages"][0]["role"] == "user"
|
|
assert kwargs["messages"][0]["content"][0]["text"] == "say hi"
|
|
|
|
def test_generate_response_reasoning_model(self, mock_boto3):
|
|
"""MiniMax M2.5 is a reasoning model: reasoningContent block comes before text."""
|
|
reasoning_response = {
|
|
"output": {
|
|
"message": {
|
|
"content": [
|
|
{"reasoningContent": {"reasoningText": {"text": "Let me think..."}}},
|
|
{"text": "Hello!"},
|
|
]
|
|
}
|
|
}
|
|
}
|
|
mock_boto3.converse.return_value = reasoning_response
|
|
llm = _make_llm("minimax.minimax-m2.5", mock_boto3)
|
|
|
|
result = llm.generate_response([{"role": "user", "content": "say hi"}])
|
|
|
|
# Must skip reasoningContent and return the actual text block
|
|
assert result == "Hello!"
|
|
|
|
def test_inference_config(self, mock_boto3):
|
|
"""inferenceConfig should include maxTokens and temperature; no topP."""
|
|
mock_boto3.converse.return_value = _converse_response()
|
|
llm = _make_llm("minimax.minimax-m2.5", mock_boto3, temperature=0.2, max_tokens=512)
|
|
|
|
llm.generate_response([{"role": "user", "content": "hi"}])
|
|
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
assert kwargs["inferenceConfig"]["maxTokens"] == 512
|
|
assert kwargs["inferenceConfig"]["temperature"] == 0.2
|
|
assert "topP" not in kwargs["inferenceConfig"]
|
|
|
|
def test_system_prompt_passed_correctly(self, mock_boto3):
|
|
"""System messages must be sent via top-level `system` param, not as a message role."""
|
|
mock_boto3.converse.return_value = _converse_response('{"facts": ["test"]}')
|
|
llm = _make_llm("minimax.minimax-m2.5", mock_boto3)
|
|
|
|
llm.generate_response([
|
|
{"role": "system", "content": "Return JSON only."},
|
|
{"role": "user", "content": "Extract facts from: test"},
|
|
])
|
|
|
|
_, kwargs = mock_boto3.converse.call_args
|
|
# system prompt must be in top-level "system" key
|
|
assert "system" in kwargs
|
|
assert kwargs["system"][0]["text"] == "Return JSON only."
|
|
# messages list must NOT contain a system role entry
|
|
for msg in kwargs["messages"]:
|
|
assert msg["role"] != "system"
|
|
# user message must be present
|
|
assert kwargs["messages"][0]["role"] == "user"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _parse_response — legacy InvokeModel provider-specific parsing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestParseResponseLegacy:
|
|
def test_ai21_missing_completions_returns_empty(self, mock_boto3):
|
|
"""When AI21 response lacks 'completions', the fallback default must
|
|
be a valid dict (not a set literal), returning empty string."""
|
|
llm = _make_llm("ai21.j2-mid-v1", mock_boto3)
|
|
import io
|
|
import json
|
|
body = io.BytesIO(json.dumps({"not_completions": True}).encode())
|
|
response = {"body": body}
|
|
result = llm._parse_response(response, tools=None)
|
|
assert result == ""
|
|
|
|
def test_ai21_normal_response(self, mock_boto3):
|
|
llm = _make_llm("ai21.j2-mid-v1", mock_boto3)
|
|
import io
|
|
import json
|
|
body = io.BytesIO(json.dumps({
|
|
"completions": [{"data": {"text": "hello from ai21"}}]
|
|
}).encode())
|
|
response = {"body": body}
|
|
result = llm._parse_response(response, tools=None)
|
|
assert result == "hello from ai21"
|
|
|
|
|
|
class TestAnthropicConverseContentParsing:
|
|
"""The Anthropic Converse branch must not assume content[0] is the text
|
|
block: Claude reasoning models emit a reasoningContent block before the
|
|
text block, and some stop conditions produce an empty content array. The
|
|
parser iterates for the first block carrying text, like the MiniMax branch.
|
|
"""
|
|
|
|
def test_text_after_reasoning_content_block(self, mock_boto3):
|
|
mock_boto3.converse.return_value = {
|
|
"output": {
|
|
"message": {
|
|
"content": [
|
|
{"reasoningContent": {"reasoningText": {"text": "step by step..."}}},
|
|
{"text": "final answer"},
|
|
]
|
|
}
|
|
}
|
|
}
|
|
llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)
|
|
|
|
assert llm.generate_response(MESSAGES) == "final answer"
|
|
|
|
def test_empty_content_returns_empty_string(self, mock_boto3):
|
|
mock_boto3.converse.return_value = {"output": {"message": {"content": []}}}
|
|
llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)
|
|
|
|
assert llm.generate_response(MESSAGES) == ""
|
|
|
|
def test_plain_text_content_still_returned(self, mock_boto3):
|
|
mock_boto3.converse.return_value = _converse_response("plain answer")
|
|
llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)
|
|
|
|
assert llm.generate_response(MESSAGES) == "plain answer"
|
|
|
|
def test_object_style_response_iterates_blocks(self, mock_boto3):
|
|
# Defensive attr-style branch: object wrapper with a reasoning block first.
|
|
from types import SimpleNamespace
|
|
|
|
reasoning_block = SimpleNamespace(reasoningContent={"reasoningText": {"text": "hmm"}})
|
|
text_block = SimpleNamespace(text="object answer")
|
|
mock_boto3.converse.return_value = SimpleNamespace(
|
|
output=SimpleNamespace(message=SimpleNamespace(content=[reasoning_block, text_block]))
|
|
)
|
|
llm = _make_llm("anthropic.claude-3-5-sonnet-20240620-v1:0", mock_boto3)
|
|
|
|
assert llm.generate_response(MESSAGES) == "object answer"
|