chopratejas--headroom/tests/test_integrations/langchain/test_memory.py

"""Tests for LangChain memory integration with automatic compression.

Tests cover:
1. HeadroomChatMessageHistory - Wrapper for chat message history with compression
2. Message conversion to/from OpenAI format
3. Rolling window compression behavior
4. Token counting and threshold detection
5. Compression statistics tracking
"""

from unittest.mock import MagicMock, patch

import pytest

# Check if LangChain is available
try:
    from langchain_core.messages import (
        AIMessage,
        BaseMessage,
        HumanMessage,
        SystemMessage,
        ToolMessage,
    )

    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False

# Skip all tests if LangChain not installed
pytestmark = pytest.mark.skipif(not LANGCHAIN_AVAILABLE, reason="LangChain not installed")


@pytest.fixture
def mock_base_history():
    """Create a mock BaseChatMessageHistory."""
    mock = MagicMock()
    mock.messages = []
    return mock


@pytest.fixture
def mock_provider():
    """Create a mock provider with token counter."""
    mock = MagicMock()
    mock_counter = MagicMock()
    mock_counter.count_text = MagicMock(side_effect=lambda text: len(text.split()))
    mock.get_token_counter = MagicMock(return_value=mock_counter)
    return mock


@pytest.fixture
def sample_langchain_messages():
    """Sample LangChain messages for testing."""
    return [
        SystemMessage(content="You are a helpful assistant."),
        HumanMessage(content="Hello, how are you?"),
        AIMessage(content="I am doing well, thank you!"),
        HumanMessage(content="What is the weather today?"),
        AIMessage(content="I don't have access to weather data."),
    ]


class TestHeadroomChatMessageHistoryInit:
    """Tests for HeadroomChatMessageHistory initialization."""

    def test_init_defaults(self, mock_base_history):
        """Initialize with default settings."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        with patch("headroom.integrations.langchain.memory.OpenAIProvider"):
            history = HeadroomChatMessageHistory(mock_base_history)

            assert history._base is mock_base_history
            assert history._threshold == 4000
            assert history._keep_recent_turns == 5
            assert history._model == "gpt-4o"
            assert history._compression_count == 0
            assert history._total_tokens_saved == 0

    def test_init_custom_threshold(self, mock_base_history, mock_provider):
        """Initialize with custom compression threshold."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(
            mock_base_history,
            compress_threshold_tokens=8000,
            keep_recent_turns=10,
            model="gpt-4-turbo",
            provider=mock_provider,
        )

        assert history._threshold == 8000
        assert history._keep_recent_turns == 10
        assert history._model == "gpt-4-turbo"
        assert history._provider is mock_provider


class TestHeadroomChatMessageHistoryMessages:
    """Tests for message access and compression."""

    def test_messages_returns_empty_when_no_messages(self, mock_base_history, mock_provider):
        """messages property returns empty list when no messages."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        mock_base_history.messages = []

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)
        messages = history.messages

        assert messages == []

    def test_messages_returns_uncompressed_when_below_threshold(
        self, mock_base_history, mock_provider, sample_langchain_messages
    ):
        """messages returns uncompressed when below token threshold."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        mock_base_history.messages = sample_langchain_messages

        history = HeadroomChatMessageHistory(
            mock_base_history,
            compress_threshold_tokens=10000,  # High threshold
            provider=mock_provider,
        )

        messages = history.messages

        # Should return all messages unchanged
        assert len(messages) == len(sample_langchain_messages)
        assert history._compression_count == 0

    def test_messages_compresses_when_over_threshold(self, mock_base_history, mock_provider):
        """messages applies compression when over token threshold."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        # Create messages that exceed threshold
        mock_base_history.messages = [
            SystemMessage(content="System " * 100),
            HumanMessage(content="User " * 100),
            AIMessage(content="Assistant " * 100),
        ]

        history = HeadroomChatMessageHistory(
            mock_base_history,
            compress_threshold_tokens=10,  # Very low threshold
            provider=mock_provider,
        )

        # Mock _apply_compression to return fewer messages
        with patch.object(history, "_apply_compression") as mock_apply:
            mock_apply.return_value = [
                SystemMessage(content="Compressed"),
            ]

            _ = history.messages

            mock_apply.assert_called_once()
            assert history._compression_count == 1

    def test_messages_tracks_tokens_saved(self, mock_base_history, mock_provider):
        """Compression tracks tokens saved."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        # Create messages that exceed threshold
        mock_base_history.messages = [
            SystemMessage(content="Word " * 50),
            HumanMessage(content="Word " * 50),
        ]

        history = HeadroomChatMessageHistory(
            mock_base_history,
            compress_threshold_tokens=10,  # Very low threshold
            provider=mock_provider,
        )

        # Mock _apply_compression to return fewer messages
        with patch.object(history, "_apply_compression") as mock_apply:
            mock_apply.return_value = [
                SystemMessage(content="Short"),
            ]

            _ = history.messages

            # tokens_saved should increase
            assert history._total_tokens_saved > 0


class TestHeadroomChatMessageHistoryAddMessage:
    """Tests for add_message methods."""

    def test_add_message(self, mock_base_history, mock_provider):
        """add_message delegates to base history."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        msg = HumanMessage(content="Hello")
        history.add_message(msg)

        mock_base_history.add_message.assert_called_once_with(msg)

    def test_add_user_message(self, mock_base_history, mock_provider):
        """add_user_message delegates to base history."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        history.add_user_message("Hello")

        mock_base_history.add_user_message.assert_called_once_with("Hello")

    def test_add_ai_message(self, mock_base_history, mock_provider):
        """add_ai_message delegates to base history."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        history.add_ai_message("Response")

        mock_base_history.add_ai_message.assert_called_once_with("Response")

    def test_clear(self, mock_base_history, mock_provider):
        """clear delegates to base history."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        history.clear()

        mock_base_history.clear.assert_called_once()


class TestHeadroomChatMessageHistoryConversion:
    """Tests for message format conversion."""

    def test_convert_to_openai_system_message(self, mock_base_history, mock_provider):
        """Convert SystemMessage to OpenAI format."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        messages = [SystemMessage(content="You are helpful.")]
        result = history._convert_to_openai(messages)

        assert len(result) == 1
        assert result[0]["role"] == "system"
        assert result[0]["content"] == "You are helpful."

    def test_convert_to_openai_human_message(self, mock_base_history, mock_provider):
        """Convert HumanMessage to OpenAI format."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        messages = [HumanMessage(content="Hello")]
        result = history._convert_to_openai(messages)

        assert result[0]["role"] == "user"
        assert result[0]["content"] == "Hello"

    def test_convert_to_openai_ai_message(self, mock_base_history, mock_provider):
        """Convert AIMessage to OpenAI format."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        messages = [AIMessage(content="I can help.")]
        result = history._convert_to_openai(messages)

        assert result[0]["role"] == "assistant"
        assert result[0]["content"] == "I can help."

    def test_convert_to_openai_ai_message_with_tool_calls(self, mock_base_history, mock_provider):
        """Convert AIMessage with tool_calls to OpenAI format."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        messages = [
            AIMessage(
                content="Calling tool...",
                tool_calls=[{"id": "call_1", "name": "search", "args": {"q": "test"}}],
            )
        ]
        result = history._convert_to_openai(messages)

        assert result[0]["role"] == "assistant"
        assert "tool_calls" in result[0]
        assert result[0]["tool_calls"][0]["id"] == "call_1"

    def test_convert_to_openai_tool_message(self, mock_base_history, mock_provider):
        """Convert ToolMessage to OpenAI format."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        messages = [ToolMessage(content='{"result": "data"}', tool_call_id="call_1")]
        result = history._convert_to_openai(messages)

        assert result[0]["role"] == "tool"
        assert result[0]["tool_call_id"] == "call_1"
        assert result[0]["content"] == '{"result": "data"}'

    def test_convert_from_openai_system(self, mock_base_history, mock_provider):
        """Convert OpenAI system message back to LangChain."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        openai_msgs = [{"role": "system", "content": "System prompt"}]
        result = history._convert_from_openai(openai_msgs)

        assert len(result) == 1
        assert isinstance(result[0], SystemMessage)
        assert result[0].content == "System prompt"

    def test_convert_from_openai_user(self, mock_base_history, mock_provider):
        """Convert OpenAI user message back to LangChain."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        openai_msgs = [{"role": "user", "content": "Hello"}]
        result = history._convert_from_openai(openai_msgs)

        assert isinstance(result[0], HumanMessage)
        assert result[0].content == "Hello"

    def test_convert_from_openai_assistant(self, mock_base_history, mock_provider):
        """Convert OpenAI assistant message back to LangChain."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        openai_msgs = [{"role": "assistant", "content": "Response"}]
        result = history._convert_from_openai(openai_msgs)

        assert isinstance(result[0], AIMessage)
        assert result[0].content == "Response"

    def test_convert_from_openai_assistant_with_tool_calls(self, mock_base_history, mock_provider):
        """Convert OpenAI assistant message with tool_calls back to LangChain."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        openai_msgs = [
            {
                "role": "assistant",
                "content": "",
                "tool_calls": [{"id": "call_1", "name": "search", "args": {}}],
            }
        ]
        result = history._convert_from_openai(openai_msgs)

        assert isinstance(result[0], AIMessage)
        # LangChain may add a 'type' field to tool_calls, so just check key fields
        assert len(result[0].tool_calls) == 1
        assert result[0].tool_calls[0]["id"] == "call_1"
        assert result[0].tool_calls[0]["name"] == "search"
        assert result[0].tool_calls[0]["args"] == {}

    def test_convert_from_openai_tool(self, mock_base_history, mock_provider):
        """Convert OpenAI tool message back to LangChain."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider)

        openai_msgs = [{"role": "tool", "tool_call_id": "call_1", "content": '{"data": 1}'}]
        result = history._convert_from_openai(openai_msgs)

        assert isinstance(result[0], ToolMessage)
        assert result[0].tool_call_id == "call_1"
        assert result[0].content == '{"data": 1}'


class TestHeadroomChatMessageHistoryTokenCounting:
    """Tests for token counting."""

    def test_count_tokens(self, mock_base_history, mock_provider):
        """Count tokens using provider's tokenizer."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(
            mock_base_history,
            provider=mock_provider,
            model="gpt-4o",
        )

        messages = [
            HumanMessage(content="Hello world"),
            AIMessage(content="Hi there"),
        ]

        count = history._count_tokens(messages)

        # Mock counts words, so "Hello world" = 2, "Hi there" = 2
        assert count == 4
        mock_provider.get_token_counter.assert_called_with("gpt-4o")


class TestHeadroomChatMessageHistoryStats:
    """Tests for compression statistics."""

    def test_get_compression_stats_initial(self, mock_base_history, mock_provider):
        """Get initial compression stats."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(
            mock_base_history,
            compress_threshold_tokens=4000,
            keep_recent_turns=5,
            provider=mock_provider,
        )

        stats = history.get_compression_stats()

        assert stats["compression_count"] == 0
        assert stats["total_tokens_saved"] == 0
        assert stats["threshold_tokens"] == 4000
        assert stats["keep_recent_turns"] == 5

    def test_get_compression_stats_after_compression(self, mock_base_history, mock_provider):
        """Get compression stats after compression."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        mock_base_history.messages = [
            SystemMessage(content="Word " * 100),
            HumanMessage(content="Word " * 100),
        ]

        history = HeadroomChatMessageHistory(
            mock_base_history,
            compress_threshold_tokens=10,
            provider=mock_provider,
        )

        # Mock _apply_compression
        with patch.object(history, "_apply_compression") as mock_apply:
            mock_apply.return_value = [SystemMessage(content="Short")]

            _ = history.messages

        stats = history.get_compression_stats()

        assert stats["compression_count"] == 1
        assert stats["total_tokens_saved"] > 0


class TestHeadroomChatMessageHistoryCompression:
    """Tests for rolling window compression."""

    def test_apply_compression_calls_pipeline(self, mock_base_history, mock_provider):
        """_apply_compression uses TransformPipeline."""
        from headroom.integrations.langchain.memory import HeadroomChatMessageHistory

        history = HeadroomChatMessageHistory(
            mock_base_history,
            compress_threshold_tokens=1000,
            keep_recent_turns=5,
            provider=mock_provider,
        )

        messages = [
            HumanMessage(content="Hello"),
            AIMessage(content="Hi there"),
        ]

        with patch("headroom.integrations.langchain.memory.TransformPipeline") as MockPipeline:
            mock_instance = MagicMock()
            mock_result = MagicMock()
            mock_result.messages = [
                {"role": "user", "content": "Hello"},
                {"role": "assistant", "content": "Hi there"},
            ]
            mock_instance.apply.return_value = mock_result
            MockPipeline.return_value = mock_instance

            result = history._apply_compression(messages)

            MockPipeline.assert_called_once()
            mock_instance.apply.assert_called_once()

            # Result should be converted back to LangChain messages
            assert all(isinstance(m, BaseMessage) for m in result)


class TestLangChainNotAvailable:
    """Tests for behavior when LangChain is not available."""

    def test_check_raises_import_error(self):
        """_check_langchain_available raises ImportError when not available."""
        from headroom.integrations.langchain.memory import _check_langchain_available

        # When LangChain IS available, should not raise
        try:
            _check_langchain_available()
        except ImportError:
            pytest.fail("Should not raise when LangChain is available")