cleveragents-core/tests/unit/agents/test_llm.py

"""
Unit tests for agents/llm.py

Tests the LLMAgent class.
"""

import pytest
from unittest.mock import Mock, patch, AsyncMock

from cleveragents.agents.llm import LLMAgent, DEFAULT_SYSTEM_MESSAGE
from cleveragents.core.exceptions import ConfigurationError, ExecutionError
from cleveragents.templates.renderer import TemplateRenderer


class TestLLMAgent:
    """Test suite for the LLMAgent class."""

    @pytest.fixture
    def template_renderer(self):
        """Create a basic template renderer."""
        return TemplateRenderer()

    @pytest.fixture
    def openai_config(self):
        """Create an OpenAI configuration."""
        return {
            "type": "llm",
            "provider": "openai",
            "model": "gpt-4",
            "temperature": 0.7,
            "max_tokens": 1000,
        }

    @pytest.fixture
    def anthropic_config(self):
        """Create an Anthropic configuration."""
        return {
            "type": "llm",
            "provider": "anthropic",
            "model": "claude-3-5-sonnet-20241022",
            "temperature": 0.8,
        }

    @pytest.fixture
    def google_config(self):
        """Create a Google configuration."""
        return {
            "type": "llm",
            "provider": "google",
            "model": "gemini-1.5-flash",
        }

    @patch('cleveragents.agents.llm.ChatOpenAI')
    def test_llm_agent_initialization_openai(self, mock_chat_openai, openai_config, template_renderer):
        """Test LLMAgent initialization with OpenAI."""
        mock_model = Mock()
        mock_chat_openai.return_value = mock_model

        agent = LLMAgent("test_llm", openai_config, template_renderer)

        assert agent.name == "test_llm"
        assert agent.provider == "openai"
        assert agent.model == "gpt-4"
        assert agent.temperature == 0.7
        assert agent.max_tokens == 1000
        assert agent.system_message == DEFAULT_SYSTEM_MESSAGE
        mock_chat_openai.assert_called_once()

    @patch('cleveragents.agents.llm.ChatAnthropic')
    def test_llm_agent_initialization_anthropic(self, mock_chat_anthropic, anthropic_config, template_renderer):
        """Test LLMAgent initialization with Anthropic."""
        mock_model = Mock()
        mock_chat_anthropic.return_value = mock_model

        agent = LLMAgent("test_llm", anthropic_config, template_renderer)

        assert agent.provider == "anthropic"
        assert agent.model == "claude-3-5-sonnet-20241022"
        mock_chat_anthropic.assert_called_once()

    @patch('cleveragents.agents.llm.ChatGoogleGenerativeAI')
    def test_llm_agent_initialization_google(self, mock_chat_google, google_config, template_renderer):
        """Test LLMAgent initialization with Google."""
        mock_model = Mock()
        mock_chat_google.return_value = mock_model

        agent = LLMAgent("test_llm", google_config, template_renderer)

        assert agent.provider == "google"
        assert agent.model == "gemini-1.5-flash"
        mock_chat_google.assert_called_once()

    @patch('cleveragents.agents.llm.ChatOpenAI')
    def test_llm_agent_default_values(self, mock_chat_openai, template_renderer):
        """Test LLMAgent with default values."""
        mock_model = Mock()
        mock_chat_openai.return_value = mock_model

        config = {"type": "llm"}
        agent = LLMAgent("test_llm", config, template_renderer)

        assert agent.provider == "openai"
        assert agent.model == "gpt-3.5-turbo"  # Default model
        assert agent.temperature == 0.7
        assert agent.max_tokens == 1000

    @patch('cleveragents.agents.llm.ChatOpenAI')
    def test_llm_agent_custom_system_message(self, mock_chat_openai, template_renderer):
        """Test LLMAgent with custom system message."""
        mock_model = Mock()
        mock_chat_openai.return_value = mock_model

        config = {
            "type": "llm",
            "system_prompt": "You are a code reviewer.",
        }
        agent = LLMAgent("test_llm", config, template_renderer)

        assert agent.system_message == "You are a code reviewer."

    @patch('cleveragents.agents.llm.ChatOpenAI')
    def test_llm_agent_get_capabilities(self, mock_chat_openai, openai_config, template_renderer):
        """Test getting agent capabilities."""
        mock_model = Mock()
        mock_chat_openai.return_value = mock_model

        agent = LLMAgent("test_llm", openai_config, template_renderer)
        capabilities = agent.get_capabilities()

        assert "text-generation" in capabilities
        assert "conversation" in capabilities

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_llm_agent_process_message(self, mock_chat_openai, openai_config, template_renderer):
        """Test processing a message."""
        # Mock the chat model
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "This is the AI response"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        agent = LLMAgent("test_llm", openai_config, template_renderer)
        result = await agent.process_message("Hello, AI!")

        assert result == "This is the AI response"
        mock_model.ainvoke.assert_called_once()

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_llm_agent_process_message_with_context(self, mock_chat_openai, openai_config, template_renderer):
        """Test processing a message with context."""
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "Response with context"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        agent = LLMAgent("test_llm", openai_config, template_renderer)
        context = {"user_name": "Alice"}
        result = await agent.process_message("Hello", context)

        assert result == "Response with context"

    @patch('cleveragents.agents.llm.ChatOpenAI')
    def test_llm_agent_get_default_model(self, mock_chat_openai, template_renderer):
        """Test getting default models for different providers."""
        mock_chat_openai.return_value = Mock()

        # Test OpenAI default
        config = {"type": "llm", "provider": "openai"}
        agent = LLMAgent("test", config, template_renderer)
        assert agent.model == "gpt-3.5-turbo"

    @patch('cleveragents.agents.llm.ChatOpenAI')
    def test_llm_agent_unsupported_provider(self, mock_chat_openai, template_renderer):
        """Test that unsupported provider raises error."""
        mock_chat_openai.return_value = Mock()

        config = {"type": "llm", "provider": "unsupported"}

        with pytest.raises(ConfigurationError) as exc_info:
            LLMAgent("test", config, template_renderer)

        assert "unsupported" in str(exc_info.value).lower()

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_llm_agent_execution_error(self, mock_chat_openai, openai_config, template_renderer):
        """Test that LangChain exceptions are wrapped in ExecutionError."""
        from langchain_core.exceptions import LangChainException

        mock_model = Mock()
        mock_model.ainvoke = AsyncMock(side_effect=LangChainException("API Error"))
        mock_chat_openai.return_value = mock_model

        agent = LLMAgent("test_llm", openai_config, template_renderer)

        with pytest.raises(ExecutionError) as exc_info:
            await agent.process_message("Hello")

        assert "LLM processing failed" in str(exc_info.value)

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_llm_agent_process_with_memory_context(self, mock_chat_openai, openai_config, template_renderer):
        """Test LLM processes with memory context."""
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "Response with memory"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        agent = LLMAgent("test_llm", openai_config, template_renderer)

        # Add some memory
        agent.memory["previous_message"] = "Previous context"

        result = await agent.process_message("New message")

        assert result == "Response with memory"

    @patch('cleveragents.agents.llm.ChatOpenAI')
    def test_llm_agent_get_metadata(self, mock_chat_openai, openai_config, template_renderer):
        """Test getting LLM agent metadata."""
        mock_model = Mock()
        mock_chat_openai.return_value = mock_model

        agent = LLMAgent("test_llm", openai_config, template_renderer)
        metadata = agent.get_metadata()

        assert metadata["name"] == "test_llm"
        assert metadata["type"] == "LLMAgent"
        assert "capabilities" in metadata

    @patch('cleveragents.agents.llm.ChatAnthropic')
    def test_llm_agent_default_anthropic_model(self, mock_anthropic, template_renderer):
        """Test default model for Anthropic provider."""
        config = {"type": "llm", "provider": "anthropic"}
        agent = LLMAgent("test_llm", config, template_renderer)

        assert agent.model == "claude-3-5-sonnet-20241022"

    @patch('cleveragents.agents.llm.ChatGoogleGenerativeAI')
    def test_llm_agent_default_google_model(self, mock_google, template_renderer):
        """Test default model for Google provider."""
        config = {"type": "llm", "provider": "google"}
        agent = LLMAgent("test_llm", config, template_renderer)

        assert agent.model == "gemini-1.5-flash"

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_llm_agent_empty_message(self, mock_chat_openai, openai_config, template_renderer):
        """Test processing an empty message."""
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "Response to empty"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        agent = LLMAgent("test_llm", openai_config, template_renderer)
        result = await agent.process_message("")

        assert result == "Response to empty"

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_llm_agent_long_message(self, mock_chat_openai, openai_config, template_renderer):
        """Test processing a long message."""
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "Response to long message"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        agent = LLMAgent("test_llm", openai_config, template_renderer)
        long_message = "x" * 10000
        result = await agent.process_message(long_message)

        assert result == "Response to long message"

    @patch('cleveragents.agents.llm.ChatOpenAI')
    def test_llm_agent_config_values(self, mock_chat_openai, template_renderer):
        """Test that config values are stored correctly."""
        config = {
            "type": "llm",
            "provider": "openai",
            "model": "gpt-4-turbo",
            "temperature": 0.9,
            "max_tokens": 2000,
            "system_prompt": "Custom system",
        }

        agent = LLMAgent("test_llm", config, template_renderer)

        assert agent.provider == "openai"
        assert agent.model == "gpt-4-turbo"
        assert agent.temperature == 0.9
        assert agent.max_tokens == 2000
        assert agent.system_message == "Custom system"


class TestLLMAgentProcessMessageExtended:
    """Extended tests for LLMAgent process_message method."""

    @pytest.fixture
    def template_renderer(self):
        """Provide template renderer for tests."""
        return TemplateRenderer()

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_process_message_with_template_in_config(self, mock_chat_openai, template_renderer):
        """Test process_message where template is in self.config."""
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "Template processed response"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        # Mock template rendering
        template_renderer.render = Mock(return_value="Rendered: Test message")

        config = {
            "provider": "openai",
            "model": "gpt-4",
            "template": "test_template",
            "template_vars": {"extra": "value"}
        }

        agent = LLMAgent("test_llm", config, template_renderer)
        result = await agent.process_message("Test message", {"key": "value"})

        # Verify template was called
        template_renderer.render.assert_called_once()
        assert "test_template" in str(template_renderer.render.call_args)
        assert result == "Template processed response"

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_process_message_with_conversation_history_in_context(self, mock_chat_openai, template_renderer):
        """Test process_message where conversation_history is in context."""
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "Response with context history"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        config = {"provider": "openai", "model": "gpt-4"}
        agent = LLMAgent("test_llm", config, template_renderer)

        context = {
            "conversation_history": [
                {"role": "user", "content": "Previous message"},
                {"role": "assistant", "content": "Previous response"}
            ]
        }

        result = await agent.process_message("New message", context)

        # Verify the model was called with history
        call_args = mock_model.ainvoke.call_args[0][0]
        assert len(call_args) == 4  # System + 2 history + current
        assert result == "Response with context history"

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_process_message_with_memory_enabled_no_context_history(self, mock_chat_openai, template_renderer):
        """Test where conversation_history is not in context but memory_enabled is true."""
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "Response from memory"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        config = {
            "provider": "openai",
            "model": "gpt-4",
            "memory_enabled": True
        }

        agent = LLMAgent("test_llm", config, template_renderer)

        # Pre-populate memory with conversation history
        agent.memory["conversation_history"] = [
            {"role": "user", "content": "Memory message"},
            {"role": "assistant", "content": "Memory response"}
        ]

        # Call without context history
        result = await agent.process_message("New message", {})

        assert result == "Response from memory"
        # Verify memory was used
        assert "conversation_history" in agent.memory

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_process_message_with_existing_history(self, mock_chat_openai, template_renderer):
        """Test where history exists."""
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "Response with history"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        config = {"provider": "openai", "model": "gpt-4"}
        agent = LLMAgent("test_llm", config, template_renderer)

        # Provide context with history
        context = {
            "conversation_history": [
                {"role": "user", "content": "First"},
                {"role": "assistant", "content": "Second"},
                {"role": "user", "content": "Third"}
            ]
        }

        result = await agent.process_message("Fourth message", context)

        # Verify history was included in the call
        call_args = mock_model.ainvoke.call_args[0][0]
        assert len(call_args) == 5  # System + 3 history + current
        assert result == "Response with history"

    @patch('cleveragents.agents.llm.ChatOpenAI')
    @pytest.mark.asyncio
    async def test_process_message_with_memory_enabled_updates_memory(self, mock_chat_openai, template_renderer):
        """Test where self.config.get('memory_enabled') is true and memory is updated."""
        mock_model = Mock()
        mock_response = Mock()
        mock_response.content = "Memory updated response"
        mock_model.ainvoke = AsyncMock(return_value=mock_response)
        mock_chat_openai.return_value = mock_model

        config = {
            "provider": "openai",
            "model": "gpt-4",
            "memory_enabled": True,
            "max_history": 10
        }

        agent = LLMAgent("test_llm", config, template_renderer)

        result = await agent.process_message("Test message")

        # Verify memory was updated
        assert "last_message" in agent.memory
        assert "last_response" in agent.memory
        assert "conversation_history" in agent.memory
        assert agent.memory["last_message"] == "Test message"
        assert agent.memory["last_response"] == "Memory updated response"
        assert len(agent.memory["conversation_history"]) == 2  # user + assistant
        assert result == "Memory updated response"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])