130 lines
4.9 KiB
Python
130 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from behave import then, when
|
|
|
|
from cleveragents.providers.llm.anthropic_provider import AnthropicChatProvider
|
|
|
|
|
|
def _register_cleanup(context, cleanup):
|
|
if hasattr(context, "add_cleanup"):
|
|
context.add_cleanup(cleanup)
|
|
else:
|
|
cleanup_handlers = getattr(context, "_cleanup_handlers", [])
|
|
cleanup_handlers.append(cleanup)
|
|
context._cleanup_handlers = cleanup_handlers
|
|
|
|
|
|
@when("I attempt to create an Anthropic chat provider without an API key")
|
|
def step_attempt_create_without_key(context):
|
|
context.anthropic_creation_error = None
|
|
try:
|
|
AnthropicChatProvider(api_key="", model="claude-3-opus-20240229")
|
|
except Exception as exc: # pragma: no cover - defensive test guard
|
|
context.anthropic_creation_error = exc
|
|
|
|
|
|
@then('the Anthropic provider creation should fail with "{error_message}"')
|
|
def step_assert_creation_failure(context, error_message):
|
|
error = getattr(context, "anthropic_creation_error", None)
|
|
assert error is not None, "Anthropic provider creation should have failed"
|
|
assert isinstance(error, ValueError)
|
|
assert str(error) == error_message
|
|
|
|
|
|
@when(
|
|
'I create an Anthropic chat provider with API key "{api_key}" and model "{model}" with custom overrides'
|
|
)
|
|
def step_create_anthropic_provider_with_overrides(context, api_key, model):
|
|
patcher = patch("cleveragents.providers.llm.anthropic_provider.ChatAnthropic")
|
|
context.chat_anthropic_patcher = patcher
|
|
mock_chat_class = patcher.start()
|
|
_register_cleanup(context, patcher.stop)
|
|
mock_chat_instance = MagicMock(name="ChatAnthropicInstance")
|
|
mock_chat_instance.get_num_tokens = MagicMock(return_value=0)
|
|
mock_chat_class.return_value = mock_chat_instance
|
|
|
|
context.anthropic_overrides = {
|
|
"temperature": 0.15,
|
|
"max_tokens": 2048,
|
|
"default_request_timeout": 12,
|
|
}
|
|
|
|
context.provider = AnthropicChatProvider(
|
|
api_key=api_key,
|
|
model=model,
|
|
max_retries=2,
|
|
**context.anthropic_overrides,
|
|
)
|
|
context.chat_anthropic_class = mock_chat_class
|
|
context.chat_anthropic_instance = mock_chat_instance
|
|
|
|
|
|
@when("I request plan generation from the Anthropic provider")
|
|
def step_request_plan_generation_anthropic(context):
|
|
patcher = patch(
|
|
"cleveragents.providers.llm.langchain_chat_provider.PlanGenerationGraph"
|
|
)
|
|
context.plan_generation_patcher = patcher
|
|
mock_graph_class = patcher.start()
|
|
_register_cleanup(context, patcher.stop)
|
|
mock_graph_instance = MagicMock(name="PlanGenerationGraphInstance")
|
|
mock_graph_instance.invoke.return_value = {
|
|
"generated_changes": [],
|
|
"validation_result": {"status": "PASS"},
|
|
"error": None,
|
|
}
|
|
mock_graph_instance.stream.return_value = iter(())
|
|
mock_graph_class.return_value = mock_graph_instance
|
|
|
|
context.plan_generation_graph = mock_graph_instance
|
|
context.plan_generation_graph_class = mock_graph_class
|
|
|
|
context.response = context.provider.generate_changes(
|
|
context.project,
|
|
context.plan,
|
|
context.contexts,
|
|
)
|
|
context.chat_anthropic_call = context.chat_anthropic_class.call_args
|
|
context.plan_generation_graph_call = mock_graph_class.call_args
|
|
|
|
|
|
@then(
|
|
'the Anthropic provider should construct ChatAnthropic with api key "{api_key}" and model "{model}"'
|
|
)
|
|
def step_assert_chat_anthropic_construction(context, api_key, model):
|
|
call = getattr(context, "chat_anthropic_call", None)
|
|
assert call is not None, "ChatAnthropic should have been constructed"
|
|
call_args, call_kwargs = call
|
|
assert call_args == (), "ChatAnthropic should be called with keyword arguments"
|
|
assert call_kwargs["api_key"] == api_key
|
|
assert call_kwargs["model"] == model
|
|
|
|
graph_call = getattr(context, "plan_generation_graph_call", None)
|
|
assert graph_call is not None, "PlanGenerationGraph should be constructed"
|
|
_, graph_kwargs = graph_call
|
|
assert graph_kwargs["llm"] is context.chat_anthropic_instance
|
|
|
|
|
|
@then("the Anthropic provider should apply the custom overrides to ChatAnthropic")
|
|
def step_assert_anthropic_overrides(context):
|
|
call = getattr(context, "chat_anthropic_call", None)
|
|
assert call is not None, "ChatAnthropic should have been constructed"
|
|
_, call_kwargs = call
|
|
overrides = getattr(context, "anthropic_overrides", None)
|
|
assert overrides, "Custom overrides should have been captured"
|
|
|
|
for key, value in overrides.items():
|
|
assert call_kwargs.get(key) == value, f"Override {key} was not forwarded"
|
|
|
|
|
|
@then(
|
|
'the Anthropic provider metadata should report name "{expected_name}" and model "{expected_model}"'
|
|
)
|
|
def step_assert_anthropic_metadata(context, expected_name, expected_model):
|
|
provider = getattr(context, "provider", None)
|
|
assert provider is not None, "Anthropic provider should exist in context"
|
|
assert provider.name == expected_name
|
|
assert provider.model_id == expected_model
|