Add unit tests for AnthropicTelegramInferenceBot
This commit is contained in:
cyclop-bot
@@ -0,0 +1,280 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch, AsyncMock, ANY
|
||||
import os
|
||||
|
||||
# Assuming anthropic_telegram_inference_bot.py is in the parent directory or PYTHONPATH is set
|
||||
from anthropic_telegram_inference_bot import AnthropicTelegramInferenceBot
|
||||
|
||||
# Mock response from Anthropic client's messages.create
|
||||
def create_mock_anthropic_response(content_text=None, stop_reason="end_turn", tool_use_parts=None):
|
||||
mock_response = MagicMock()
|
||||
mock_response.stop_reason = stop_reason
|
||||
|
||||
content_blocks = []
|
||||
if content_text:
|
||||
text_block = MagicMock()
|
||||
text_block.type = "text"
|
||||
text_block.text = content_text
|
||||
content_blocks.append(text_block)
|
||||
|
||||
if tool_use_parts:
|
||||
for tu_part in tool_use_parts: # tu_part = {"id": "toolu_123", "name": "get_weather", "input": {}}
|
||||
tool_block = MagicMock()
|
||||
tool_block.type = "tool_use"
|
||||
tool_block.id = tu_part["id"]
|
||||
tool_block.name = tu_part["name"]
|
||||
tool_block.input = tu_part["input"]
|
||||
content_blocks.append(tool_block)
|
||||
|
||||
mock_response.content = content_blocks
|
||||
return mock_response
|
||||
|
||||
class TestAnthropicTelegramInferenceBot(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.original_anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY")
|
||||
self.original_small_model = os.environ.get("ANTHROPIC_SMALL_MODEL")
|
||||
self.original_large_model = os.environ.get("ANTHROPIC_LARGE_MODEL")
|
||||
self.original_system_prompt_path = os.environ.get("SYSTEM_PROMPT_PATH")
|
||||
|
||||
for key in ["ANTHROPIC_API_KEY", "ANTHROPIC_SMALL_MODEL", "ANTHROPIC_LARGE_MODEL", "SYSTEM_PROMPT_PATH"]:
|
||||
if os.environ.get(key):
|
||||
del os.environ[key]
|
||||
|
||||
self.mock_anthropic_client_instance = MagicMock()
|
||||
self.mock_anthropic_client_instance.messages.create = MagicMock()
|
||||
|
||||
def tearDown(self):
|
||||
if self.original_anthropic_api_key: os.environ["ANTHROPIC_API_KEY"] = self.original_anthropic_api_key
|
||||
if self.original_small_model: os.environ["ANTHROPIC_SMALL_MODEL"] = self.original_small_model
|
||||
if self.original_large_model: os.environ["ANTHROPIC_LARGE_MODEL"] = self.original_large_model
|
||||
if self.original_system_prompt_path: os.environ["SYSTEM_PROMPT_PATH"] = self.original_system_prompt_path
|
||||
|
||||
@patch('anthropic.Anthropic')
|
||||
def test_init_with_anthropic_defaults_env_key(self, MockAnthropicConstructor):
|
||||
MockAnthropicConstructor.return_value = self.mock_anthropic_client_instance
|
||||
os.environ["ANTHROPIC_API_KEY"] = "test_anthropic_key"
|
||||
|
||||
bot = AnthropicTelegramInferenceBot()
|
||||
|
||||
MockAnthropicConstructor.assert_called_once_with(api_key="test_anthropic_key")
|
||||
self.assertEqual(bot.anthropic_client, self.mock_anthropic_client_instance)
|
||||
self.assertEqual(bot.model, os.environ.get("ANTHROPIC_SMALL_MODEL", "claude-3-haiku-20240307"))
|
||||
self.assertEqual(bot.max_tokens, int(os.environ.get("ANTHROPIC_SMALL_MODEL_MAX_TOKENS", 2000)))
|
||||
|
||||
@patch('anthropic.Anthropic')
|
||||
def test_init_with_provided_client_and_models(self, MockAnthropicConstructor):
|
||||
preconfigured_client = MagicMock()
|
||||
bot = AnthropicTelegramInferenceBot(
|
||||
anthropic_client=preconfigured_client,
|
||||
small_model_name="custom-small",
|
||||
small_model_max_tokens=100,
|
||||
large_model_name="custom-large",
|
||||
large_model_max_tokens=200
|
||||
)
|
||||
|
||||
MockAnthropicConstructor.assert_not_called()
|
||||
self.assertEqual(bot.anthropic_client, preconfigured_client)
|
||||
self.assertEqual(bot.model, "custom-small")
|
||||
self.assertEqual(bot.max_tokens, 100)
|
||||
self.assertEqual(bot.small_model_name, "custom-small")
|
||||
self.assertEqual(bot.large_model_name, "custom-large")
|
||||
|
||||
|
||||
def test_get_llm_description(self):
|
||||
bot = AnthropicTelegramInferenceBot(small_model_name="claude-test", small_model_max_tokens=500)
|
||||
self.assertEqual(bot.get_llm_description(), "LLM: claude-test, Max Tokens: 500")
|
||||
|
||||
async def test_switch_model(self):
|
||||
bot = AnthropicTelegramInferenceBot(
|
||||
small_model_name="claude-small", small_model_max_tokens=10,
|
||||
large_model_name="claude-large", large_model_max_tokens=20
|
||||
)
|
||||
self.assertEqual(bot.model, "claude-small")
|
||||
self.assertEqual(bot.max_tokens, 10)
|
||||
|
||||
status = await bot.switch_model()
|
||||
self.assertEqual(bot.model, "claude-large")
|
||||
self.assertEqual(bot.max_tokens, 20)
|
||||
self.assertEqual(status, "Switched to model: claude-large")
|
||||
|
||||
status = await bot.switch_model()
|
||||
self.assertEqual(bot.model, "claude-small")
|
||||
self.assertEqual(bot.max_tokens, 10)
|
||||
self.assertEqual(status, "Switched to model: claude-small")
|
||||
|
||||
def test_get_chat_response_success_text_only(self):
|
||||
bot = AnthropicTelegramInferenceBot(anthropic_client=self.mock_anthropic_client_instance)
|
||||
bot.model = "test-claude"
|
||||
bot.max_tokens = 150
|
||||
|
||||
mock_api_response = create_mock_anthropic_response(content_text="Hello from Anthropic API")
|
||||
self.mock_anthropic_client_instance.messages.create.return_value = mock_api_response
|
||||
|
||||
messages = [{"role": "user", "content": "Hi"}] # Anthropic format
|
||||
response = bot.get_chat_response(messages, []) # tools = empty list
|
||||
|
||||
self.mock_anthropic_client_instance.messages.create.assert_called_once_with(
|
||||
model="test-claude",
|
||||
max_tokens=150,
|
||||
messages=messages,
|
||||
system=bot.system_prompt, # Ensure system prompt is passed
|
||||
tools=None, # No tools passed to API if empty list or None
|
||||
tool_choice=None
|
||||
)
|
||||
self.assertEqual(response, mock_api_response)
|
||||
|
||||
def test_get_chat_response_with_tools(self):
|
||||
bot = AnthropicTelegramInferenceBot(anthropic_client=self.mock_anthropic_client_instance)
|
||||
bot.model = "claude-toolmaster"
|
||||
bot.max_tokens = 300
|
||||
|
||||
mock_tools_spec = [{"name": "get_weather", "description": "Gets weather", "input_schema": {"type": "object", "properties": {}}}]
|
||||
|
||||
mock_api_response = create_mock_anthropic_response(content_text="Thinking...", tool_use_parts=[
|
||||
{"id": "tool1", "name": "get_weather", "input": {"location": "here"}}
|
||||
])
|
||||
self.mock_anthropic_client_instance.messages.create.return_value = mock_api_response
|
||||
|
||||
messages = [{"role": "user", "content": "Weather?"}]
|
||||
response = bot.get_chat_response(messages, mock_tools_spec)
|
||||
|
||||
self.mock_anthropic_client_instance.messages.create.assert_called_once_with(
|
||||
model="claude-toolmaster",
|
||||
max_tokens=300,
|
||||
messages=messages,
|
||||
system=bot.system_prompt,
|
||||
tools=mock_tools_spec,
|
||||
tool_choice={"type": "auto"}
|
||||
)
|
||||
self.assertEqual(response.content[0].type, "text") # First part can be text
|
||||
self.assertEqual(response.content[1].type, "tool_use")
|
||||
|
||||
|
||||
def test_get_chat_response_api_error(self):
|
||||
bot = AnthropicTelegramInferenceBot(anthropic_client=self.mock_anthropic_client_instance)
|
||||
self.mock_anthropic_client_instance.messages.create.side_effect = Exception("Anthropic API Down")
|
||||
|
||||
with self.assertRaisesRegex(Exception, "Anthropic API Down"):
|
||||
bot.get_chat_response([{"role": "user", "content": "trigger"}], [])
|
||||
|
||||
|
||||
async def test_handle_message_simple_response_no_tools(self):
|
||||
# This test is more involved as it touches BaseTelegramInferenceBot's handle_message structure
|
||||
# which then calls the overridden get_chat_response.
|
||||
bot = AnthropicTelegramInferenceBot(anthropic_client=self.mock_anthropic_client_instance)
|
||||
bot.system_prompt = "System prompt for Anthropic"
|
||||
|
||||
# Mock get_chat_response directly to isolate its behavior from full handle_message logic of base
|
||||
# However, the point of this bot is its get_chat_response and subsequent processing.
|
||||
# So, let's mock the API call within get_chat_response.
|
||||
|
||||
api_response = create_mock_anthropic_response(content_text="Anthropic says hello.")
|
||||
self.mock_anthropic_client_instance.messages.create.return_value = api_response
|
||||
|
||||
# Ensure functions are empty for this test, so no tool logic is triggered
|
||||
bot.functions = []
|
||||
bot.tools = []
|
||||
|
||||
response_content = await bot.handle_message(user_id=101, user_message="Hello Anthropic")
|
||||
|
||||
self.assertEqual(response_content, "Anthropic says hello.")
|
||||
self.assertIn(101, bot.conversation_history)
|
||||
# Anthropic's handle_message structure:
|
||||
# 1. User message added to history.
|
||||
# 2. get_chat_response is called.
|
||||
# 3. Response content (text) is extracted.
|
||||
# 4. Assistant text response is added to history.
|
||||
# Expected history: [User, Assistant_Text_Response] (system prompt handled by get_chat_response)
|
||||
# The base class handle_message adds system prompt if not present.
|
||||
# Anthropic handle_message modifies history format before calling get_chat_response.
|
||||
|
||||
# Let's trace Base.handle_message -> Anthropic.handle_message -> Anthropic.get_chat_response
|
||||
# Base.handle_message:
|
||||
# - Adds system prompt to history if first turn: `self.conversation_history[user_id] = [{"role": "system", "content": self.system_prompt}]` (OpenAI style)
|
||||
# - Appends user message: `{"role": "user", "content": user_message}`
|
||||
# - Calls self.get_chat_response(messages, self.functions) -> This is Anthropic's get_chat_response
|
||||
# Anthropic.get_chat_response:
|
||||
# - Takes OpenAI style `messages` and `self.functions` (tool specs).
|
||||
# - Calls `anthropic_client.messages.create` with Anthropic style messages and system prompt.
|
||||
# Anthropic.handle_message (overridden):
|
||||
# - Prepares Anthropic-style messages from conversation_history (which is OpenAI style from Base)
|
||||
# - Calls get_chat_response with these Anthropic messages and self.functions (tool_specs)
|
||||
# - Processes response, extracts text, handles tool calls.
|
||||
# - Appends *user* message (original) and *assistant* text response to self.conversation_history (OpenAI style).
|
||||
|
||||
# For this test, we are calling AnthropicBot.handle_message directly.
|
||||
# 1. `user_id` not in `self.conversation_history`: `system_prompt` not added yet by Base logic.
|
||||
# Anthropic's `handle_message` will create `anthropic_messages` from this.
|
||||
# If `conversation_history` is empty, `anthropic_messages` = `[{"role": "user", "content": user_message}]`
|
||||
# 2. `get_chat_response` called with `anthropic_messages` and `bot.system_prompt` passed to API.
|
||||
# 3. Response "Anthropic says hello."
|
||||
# 4. Original `user_message` and "Anthropic says hello." (as assistant) added to `self.conversation_history`.
|
||||
|
||||
history = bot.conversation_history[101]
|
||||
self.assertEqual(len(history), 2) # User, Assistant
|
||||
self.assertEqual(history[0]["role"], "user")
|
||||
self.assertEqual(history[0]["content"], "Hello Anthropic")
|
||||
self.assertEqual(history[1]["role"], "assistant")
|
||||
self.assertEqual(history[1]["content"], "Anthropic says hello.")
|
||||
|
||||
# Check API call (made by the mocked get_chat_response indirectly)
|
||||
self.mock_anthropic_client_instance.messages.create.assert_called_once()
|
||||
call_args = self.mock_anthropic_client_instance.messages.create.call_args
|
||||
self.assertEqual(call_args.kwargs["system"], "System prompt for Anthropic")
|
||||
# Initial messages for API should just be the user message for first turn
|
||||
self.assertEqual(call_args.kwargs["messages"], [{"role": "user", "content": "Hello Anthropic"}])
|
||||
|
||||
|
||||
async def test_handle_message_with_tool_calls(self):
|
||||
bot = AnthropicTelegramInferenceBot(anthropic_client=self.mock_anthropic_client_instance)
|
||||
bot.system_prompt = "You are a helpful, tool-using assistant."
|
||||
|
||||
# Define a tool for the bot (OpenAI format, will be converted by Anthropic bot for API)
|
||||
mock_tool_oai_format = {"type": "function", "function": {"name": "get_weather", "description": "Get weather", "parameters": {}}}
|
||||
bot.functions = [mock_tool_oai_format] # This is used to generate anthropic_tools for API
|
||||
|
||||
# API Response 1: Request for tool call
|
||||
tool_use_part = {"id": "toolu_xyz", "name": "get_weather", "input": {"location": "paris"}}
|
||||
api_response_1 = create_mock_anthropic_response(tool_use_parts=[tool_use_part])
|
||||
|
||||
# API Response 2: Final text response after tool execution
|
||||
api_response_2 = create_mock_anthropic_response(content_text="The weather in Paris is nice.")
|
||||
|
||||
self.mock_anthropic_client_instance.messages.create.side_effect = [api_response_1, api_response_2]
|
||||
|
||||
# Mock the bot's call_tool method (from BaseTelegramInferenceBot)
|
||||
bot.call_tool = MagicMock(return_value='''{"weather": "sunny"}''') # Tool execution result
|
||||
|
||||
user_id = 102
|
||||
user_message = "What's the weather in Paris?"
|
||||
final_text_response = await bot.handle_message(user_id, user_message)
|
||||
|
||||
self.assertEqual(final_text_response, "The weather in Paris is nice.")
|
||||
self.assertEqual(self.mock_anthropic_client_instance.messages.create.call_count, 2)
|
||||
|
||||
bot.call_tool.assert_called_once_with("get_weather", {"location": "paris"}) # Anthropic passes input as dict
|
||||
|
||||
# Check conversation history (OpenAI style)
|
||||
history = bot.conversation_history[user_id]
|
||||
self.assertEqual(history[0]["role"], "user")
|
||||
self.assertEqual(history[0]["content"], user_message)
|
||||
|
||||
# Assistant message that requested tool call (Anthropic-specific format stored by its handle_message)
|
||||
# Anthropic's handle_message appends the raw tool_use block and then the tool_result
|
||||
self.assertEqual(history[1]["role"], "assistant")
|
||||
self.assertTrue(isinstance(history[1]["content"], list)) # Anthropic content is a list
|
||||
self.assertEqual(history[1]["content"][0]["type"], "tool_use")
|
||||
self.assertEqual(history[1]["content"][0]["id"], "toolu_xyz")
|
||||
|
||||
self.assertEqual(history[2]["role"], "tool")
|
||||
self.assertEqual(history[2]["tool_call_id"], "toolu_xyz")
|
||||
self.assertEqual(history[2]["name"], "get_weather")
|
||||
self.assertEqual(history[2]["content"], '''{"weather": "sunny"}''') # call_tool result
|
||||
|
||||
self.assertEqual(history[3]["role"], "assistant") # Final text response
|
||||
self.assertTrue(isinstance(history[3]["content"], str)) # simple text
|
||||
self.assertEqual(history[3]["content"], "The weather in Paris is nice.")
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user