cyclop-bot
159 lines
7.5 KiB
Python
159 lines
7.5 KiB
Python
import unittest
|
|
from unittest.mock import MagicMock, patch, ANY
|
|
import os
|
|
|
|
# Assuming chatgpt_telegram_inference_bot.py and its parent are accessible
|
|
from chatgpt_telegram_inference_bot import ChatGPTTelegramInferenceBot
|
|
from openai_compatible_inference_bot import OpenAICompatibleInferenceBot # For patching super
|
|
|
|
class TestChatGPTTelegramInferenceBot(unittest.IsolatedAsyncioTestCase):
|
|
|
|
def setUp(self):
|
|
# Store and clear relevant environment variables
|
|
self.original_openai_key = os.environ.get("OPENAI_API_KEY")
|
|
self.original_small_model = os.environ.get("OPENAI_SMALL_MODEL")
|
|
self.original_large_model = os.environ.get("OPENAI_LARGE_MODEL")
|
|
self.original_small_tokens = os.environ.get("OPENAI_SMALL_MODEL_MAX_TOKENS")
|
|
self.original_large_tokens = os.environ.get("OPENAI_LARGE_MODEL_MAX_TOKENS")
|
|
self.original_system_prompt_path = os.environ.get("SYSTEM_PROMPT_PATH")
|
|
|
|
for key in ["OPENAI_API_KEY", "OPENAI_SMALL_MODEL", "OPENAI_LARGE_MODEL",
|
|
"OPENAI_SMALL_MODEL_MAX_TOKENS", "OPENAI_LARGE_MODEL_MAX_TOKENS", "SYSTEM_PROMPT_PATH"]:
|
|
if os.environ.get(key):
|
|
del os.environ[key]
|
|
|
|
# Mock the OpenAI client that OpenAICompatibleInferenceBot's __init__ might create
|
|
self.mock_openai_client = MagicMock()
|
|
|
|
def tearDown(self):
|
|
# Restore environment variables
|
|
if self.original_openai_key: os.environ["OPENAI_API_KEY"] = self.original_openai_key
|
|
if self.original_small_model: os.environ["OPENAI_SMALL_MODEL"] = self.original_small_model
|
|
if self.original_large_model: os.environ["OPENAI_LARGE_MODEL"] = self.original_large_model
|
|
if self.original_small_tokens: os.environ["OPENAI_SMALL_MODEL_MAX_TOKENS"] = self.original_small_tokens
|
|
if self.original_large_tokens: os.environ["OPENAI_LARGE_MODEL_MAX_TOKENS"] = self.original_large_tokens
|
|
if self.original_system_prompt_path: os.environ["SYSTEM_PROMPT_PATH"] = self.original_system_prompt_path
|
|
|
|
|
|
@patch.object(OpenAICompatibleInferenceBot, '__init__') # Mock the superclass's __init__
|
|
def test_init_defaults_and_super_call(self, mock_super_init):
|
|
os.environ["OPENAI_API_KEY"] = "test_key_chatgpt"
|
|
os.environ["OPENAI_SMALL_MODEL"] = "gpt-3.5-turbo-env"
|
|
os.environ["OPENAI_SMALL_MODEL_MAX_TOKENS"] = "350"
|
|
|
|
bot = ChatGPTTelegramInferenceBot()
|
|
|
|
mock_super_init.assert_called_once_with(
|
|
client=None, # ChatGPT bot will let superclass create it
|
|
api_key="test_key_chatgpt", # Passed to super
|
|
base_url=None,
|
|
api_version=None,
|
|
azure_deployment=None,
|
|
model_name="gpt-3.5-turbo-env", # Default small model from env
|
|
max_tokens_str="350", # Default small model tokens from env
|
|
small_model_name="gpt-3.5-turbo-env",
|
|
small_model_max_tokens_str="350",
|
|
large_model_name=os.environ.get("OPENAI_LARGE_MODEL", "gpt-4-turbo-preview"), # Default large
|
|
large_model_max_tokens_str=os.environ.get("OPENAI_LARGE_MODEL_MAX_TOKENS"),
|
|
system_prompt_content=None,
|
|
system_prompt_path=None,
|
|
is_gemini=False,
|
|
max_history_length=20 # Default from OpenAICompatibleInferenceBot
|
|
)
|
|
|
|
@patch.object(OpenAICompatibleInferenceBot, '__init__')
|
|
def test_init_with_arguments(self, mock_super_init):
|
|
mock_client_arg = MagicMock()
|
|
bot = ChatGPTTelegramInferenceBot(
|
|
openai_client=mock_client_arg,
|
|
api_key="arg_key",
|
|
small_model_name="arg_small_model",
|
|
small_model_max_tokens="123",
|
|
large_model_name="arg_large_model",
|
|
large_model_max_tokens="456",
|
|
system_prompt_content="Arg prompt"
|
|
)
|
|
mock_super_init.assert_called_once_with(
|
|
client=mock_client_arg,
|
|
api_key="arg_key",
|
|
base_url=None,
|
|
api_version=None,
|
|
azure_deployment=None,
|
|
model_name="arg_small_model", # Initially configured with small model
|
|
max_tokens_str="123",
|
|
small_model_name="arg_small_model",
|
|
small_model_max_tokens_str="123",
|
|
large_model_name="arg_large_model",
|
|
large_model_max_tokens_str="456",
|
|
system_prompt_content="Arg prompt",
|
|
system_prompt_path=None,
|
|
is_gemini=False,
|
|
max_history_length=20
|
|
)
|
|
|
|
# Test switch_model - this method is part of ChatGPTTelegramInferenceBot
|
|
# It calls _configure_model_and_tokens which is in the superclass.
|
|
# We need a bot instance where _configure_model_and_tokens can be called.
|
|
@patch('openai.OpenAI') # To allow instantiation of the bot by mocking client creation
|
|
async def test_switch_model_logic(self, mock_openai_constructor):
|
|
mock_openai_constructor.return_value = self.mock_openai_client # Mock client creation in super
|
|
|
|
# Set env vars for model names that switch_model will use as fallback
|
|
os.environ["OPENAI_SMALL_MODEL"] = "env-small-gpt"
|
|
os.environ["OPENAI_SMALL_MODEL_MAX_TOKENS"] = "100"
|
|
os.environ["OPENAI_LARGE_MODEL"] = "env-large-gpt"
|
|
os.environ["OPENAI_LARGE_MODEL_MAX_TOKENS"] = "200"
|
|
|
|
# Instantiate with initial model (small)
|
|
bot = ChatGPTTelegramInferenceBot()
|
|
self.assertEqual(bot.model, "env-small-gpt")
|
|
self.assertEqual(bot.max_tokens, 100)
|
|
|
|
# Switch to large
|
|
status = await bot.switch_model()
|
|
self.assertEqual(bot.model, "env-large-gpt")
|
|
self.assertEqual(bot.max_tokens, 200)
|
|
self.assertEqual(status, "Switched to model: env-large-gpt")
|
|
|
|
# Switch back to small
|
|
status = await bot.switch_model()
|
|
self.assertEqual(bot.model, "env-small-gpt")
|
|
self.assertEqual(bot.max_tokens, 100)
|
|
self.assertEqual(status, "Switched to model: env-small-gpt")
|
|
|
|
@patch('openai.OpenAI')
|
|
async def test_switch_model_uses_instance_configs_if_provided(self, mock_openai_constructor):
|
|
mock_openai_constructor.return_value = self.mock_openai_client
|
|
|
|
# Instantiate with specific model names, overriding potential env vars
|
|
bot = ChatGPTTelegramInferenceBot(
|
|
small_model_name="init-small", small_model_max_tokens="50",
|
|
large_model_name="init-large", large_model_max_tokens="150"
|
|
)
|
|
self.assertEqual(bot.model, "init-small") # Starts with small
|
|
self.assertEqual(bot.max_tokens, 50)
|
|
|
|
# Switch to large
|
|
status = await bot.switch_model()
|
|
self.assertEqual(bot.model, "init-large")
|
|
self.assertEqual(bot.max_tokens, 150)
|
|
self.assertEqual(status, "Switched to model: init-large")
|
|
|
|
# Switch back to small
|
|
status = await bot.switch_model()
|
|
self.assertEqual(bot.model, "init-small")
|
|
self.assertEqual(bot.max_tokens, 50)
|
|
self.assertEqual(status, "Switched to model: init-small")
|
|
|
|
# get_llm_description is inherited from OpenAICompatibleInferenceBot.
|
|
# Test just to ensure it works in the context of a ChatGPTBot instance
|
|
@patch('openai.OpenAI')
|
|
def test_get_llm_description_for_chatgpt_bot(self, mock_openai_constructor):
|
|
mock_openai_constructor.return_value = self.mock_openai_client
|
|
bot = ChatGPTTelegramInferenceBot(small_model_name="gpt-3.5-desc", small_model_max_tokens="777")
|
|
# Initially configured with small model
|
|
self.assertEqual(bot.get_llm_description(), "LLM: gpt-3.5-desc, Max Tokens: 777, Azure: False")
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|