cyclop-bot
281 lines
15 KiB
Python
281 lines
15 KiB
Python
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()
|