arc53--docsgpt/tests/api/answer/services/compression/test_orchestrator.py

"""Tests for application/api/answer/services/compression/orchestrator.py"""

from unittest.mock import MagicMock, patch

import pytest

from application.api.answer.services.compression.orchestrator import (
    CompressionOrchestrator,
)
from application.api.answer.services.compression.types import (
    CompressionMetadata,
    CompressionResult,
)


@pytest.fixture
def mock_conversation_service():
    svc = MagicMock()
    return svc


@pytest.fixture
def mock_threshold_checker():
    checker = MagicMock()
    return checker


@pytest.fixture
def orchestrator(mock_conversation_service, mock_threshold_checker):
    return CompressionOrchestrator(
        conversation_service=mock_conversation_service,
        threshold_checker=mock_threshold_checker,
    )


@pytest.fixture
def sample_conversation():
    return {
        "queries": [
            {"prompt": "q0", "response": "r0"},
            {"prompt": "q1", "response": "r1"},
            {"prompt": "q2", "response": "r2"},
        ],
        "compression_metadata": {},
        "agent_id": "agent-1",
    }


@pytest.fixture
def decoded_token():
    return {"sub": "user123"}


@pytest.mark.unit
class TestCompressIfNeeded:
    def test_conversation_not_found_returns_failure(
        self, orchestrator, mock_conversation_service
    ):
        mock_conversation_service.get_conversation.return_value = None

        result = orchestrator.compress_if_needed(
            conversation_id="conv1",
            user_id="user1",
            model_id="gpt-4",
            decoded_token={"sub": "user1"},
        )

        assert result.success is False
        assert "not found" in result.error

    def test_no_compression_needed(
        self,
        orchestrator,
        mock_conversation_service,
        mock_threshold_checker,
        sample_conversation,
    ):
        mock_conversation_service.get_conversation.return_value = sample_conversation
        mock_threshold_checker.should_compress.return_value = False

        result = orchestrator.compress_if_needed(
            conversation_id="conv1",
            user_id="user1",
            model_id="gpt-4",
            decoded_token={"sub": "user1"},
        )

        assert result.success is True
        assert result.compression_performed is False
        assert len(result.recent_queries) == 3

    def test_compression_performed_successfully(
        self,
        orchestrator,
        mock_conversation_service,
        mock_threshold_checker,
        sample_conversation,
        decoded_token,
    ):
        mock_conversation_service.get_conversation.return_value = sample_conversation
        mock_threshold_checker.should_compress.return_value = True

        mock_metadata = MagicMock(spec=CompressionMetadata)
        mock_metadata.compression_ratio = 5.0
        mock_metadata.original_token_count = 1000
        mock_metadata.compressed_token_count = 200
        mock_metadata.to_dict.return_value = {"query_index": 2}

        with patch.object(
            orchestrator, "_perform_compression"
        ) as mock_perform:
            mock_perform.return_value = CompressionResult.success_with_compression(
                "compressed summary",
                [{"prompt": "q2", "response": "r2"}],
                mock_metadata,
            )

            result = orchestrator.compress_if_needed(
                conversation_id="conv1",
                user_id="user1",
                model_id="gpt-4",
                decoded_token=decoded_token,
            )

            assert result.success is True
            assert result.compression_performed is True
            assert result.compressed_summary == "compressed summary"
            mock_perform.assert_called_once()

    def test_exception_returns_failure(
        self,
        orchestrator,
        mock_conversation_service,
    ):
        mock_conversation_service.get_conversation.side_effect = RuntimeError("DB down")

        result = orchestrator.compress_if_needed(
            conversation_id="conv1",
            user_id="user1",
            model_id="gpt-4",
            decoded_token={"sub": "user1"},
        )

        assert result.success is False
        assert "DB down" in result.error

    def test_custom_query_tokens(
        self,
        orchestrator,
        mock_conversation_service,
        mock_threshold_checker,
        sample_conversation,
    ):
        mock_conversation_service.get_conversation.return_value = sample_conversation
        mock_threshold_checker.should_compress.return_value = False

        orchestrator.compress_if_needed(
            conversation_id="conv1",
            user_id="user1",
            model_id="gpt-4",
            decoded_token={"sub": "user1"},
            current_query_tokens=1000,
        )

        # user_id flows through so BYOM custom-model UUIDs resolve to
        # the user's declared context window in the threshold check.
        mock_threshold_checker.should_compress.assert_called_once_with(
            sample_conversation, "gpt-4", 1000, user_id="user1"
        )


@pytest.mark.unit
class TestPerformCompression:
    @patch(
        "application.api.answer.services.compression.orchestrator.get_provider_from_model_id"
    )
    @patch(
        "application.api.answer.services.compression.orchestrator.get_api_key_for_provider"
    )
    @patch("application.api.answer.services.compression.orchestrator.LLMCreator")
    @patch("application.api.answer.services.compression.orchestrator.CompressionService")
    @patch("application.api.answer.services.compression.orchestrator.settings")
    def test_successful_compression(
        self,
        mock_settings,
        MockCompressionService,
        MockLLMCreator,
        mock_get_api_key,
        mock_get_provider,
        mock_conversation_service,
        mock_threshold_checker,
        sample_conversation,
        decoded_token,
    ):
        mock_settings.COMPRESSION_MODEL_OVERRIDE = None
        mock_get_provider.return_value = "openai"
        mock_get_api_key.return_value = "sk-test"
        MockLLMCreator.create_llm.return_value = MagicMock()

        mock_metadata = MagicMock(spec=CompressionMetadata)
        mock_metadata.compression_ratio = 5.0
        mock_metadata.original_token_count = 500
        mock_metadata.compressed_token_count = 100

        mock_svc_instance = MagicMock()
        mock_svc_instance.compress_and_save.return_value = mock_metadata
        mock_svc_instance.get_compressed_context.return_value = (
            "compressed text",
            [{"prompt": "q2", "response": "r2"}],
        )
        MockCompressionService.return_value = mock_svc_instance

        # After compression, reload conversation
        mock_conversation_service.get_conversation.return_value = sample_conversation

        orch = CompressionOrchestrator(
            conversation_service=mock_conversation_service,
            threshold_checker=mock_threshold_checker,
        )

        result = orch._perform_compression(
            "conv1", sample_conversation, "gpt-4", decoded_token
        )

        assert result.success is True
        assert result.compression_performed is True
        assert result.compressed_summary == "compressed text"
        mock_svc_instance.compress_and_save.assert_called_once()

    @patch(
        "application.api.answer.services.compression.orchestrator.get_provider_from_model_id"
    )
    @patch(
        "application.api.answer.services.compression.orchestrator.get_api_key_for_provider"
    )
    @patch("application.api.answer.services.compression.orchestrator.LLMCreator")
    @patch("application.api.answer.services.compression.orchestrator.settings")
    def test_uses_compression_model_override(
        self,
        mock_settings,
        MockLLMCreator,
        mock_get_api_key,
        mock_get_provider,
        mock_conversation_service,
        mock_threshold_checker,
        decoded_token,
    ):
        mock_settings.COMPRESSION_MODEL_OVERRIDE = "gpt-3.5-turbo"
        mock_get_provider.return_value = "openai"
        mock_get_api_key.return_value = "sk-test"
        MockLLMCreator.create_llm.return_value = MagicMock()

        conversation = {"queries": [{"prompt": "q", "response": "r"}], "agent_id": "a"}

        with patch(
            "application.api.answer.services.compression.orchestrator.CompressionService"
        ) as MockCS:
            mock_svc = MagicMock()
            mock_svc.compress_and_save.return_value = MagicMock(
                compression_ratio=3.0,
                original_token_count=300,
                compressed_token_count=100,
            )
            mock_svc.get_compressed_context.return_value = ("s", [])
            MockCS.return_value = mock_svc

            mock_conversation_service.get_conversation.return_value = conversation

            orch = CompressionOrchestrator(
                conversation_service=mock_conversation_service,
                threshold_checker=mock_threshold_checker,
            )
            orch._perform_compression("c1", conversation, "gpt-4", decoded_token)

            # Verify the override model was used. user_id flows from
            # decoded_token['sub'] so per-user BYOM custom-model UUIDs
            # resolve.
            mock_get_provider.assert_called_with(
                "gpt-3.5-turbo", user_id=decoded_token["sub"]
            )

    @patch(
        "application.api.answer.services.compression.orchestrator.get_provider_from_model_id"
    )
    @patch(
        "application.api.answer.services.compression.orchestrator.get_api_key_for_provider"
    )
    @patch("application.api.answer.services.compression.orchestrator.LLMCreator")
    @patch("application.api.answer.services.compression.orchestrator.CompressionService")
    @patch("application.api.answer.services.compression.orchestrator.settings")
    def test_no_queries_returns_no_compression(
        self,
        mock_settings,
        MockCompressionService,
        MockLLMCreator,
        mock_get_api_key,
        mock_get_provider,
        mock_conversation_service,
        mock_threshold_checker,
        decoded_token,
    ):
        mock_settings.COMPRESSION_MODEL_OVERRIDE = None
        mock_get_provider.return_value = "openai"
        mock_get_api_key.return_value = "sk-test"
        MockLLMCreator.create_llm.return_value = MagicMock()

        conversation = {"queries": [], "agent_id": "a"}

        orch = CompressionOrchestrator(
            conversation_service=mock_conversation_service,
            threshold_checker=mock_threshold_checker,
        )
        result = orch._perform_compression("c1", conversation, "gpt-4", decoded_token)

        assert result.success is True
        assert result.compression_performed is False

    def test_exception_returns_failure(
        self,
        mock_conversation_service,
        mock_threshold_checker,
        decoded_token,
    ):
        conversation = {
            "queries": [{"prompt": "q", "response": "r"}],
            "agent_id": "a",
        }

        with patch(
            "application.api.answer.services.compression.orchestrator.settings"
        ) as mock_settings, patch(
            "application.api.answer.services.compression.orchestrator.get_provider_from_model_id",
            side_effect=RuntimeError("provider error"),
        ):
            mock_settings.COMPRESSION_MODEL_OVERRIDE = None

            orch = CompressionOrchestrator(
                conversation_service=mock_conversation_service,
                threshold_checker=mock_threshold_checker,
            )
            result = orch._perform_compression(
                "c1", conversation, "gpt-4", decoded_token
            )

            assert result.success is False
            assert "provider error" in result.error


@pytest.mark.unit
class TestCompressMidExecution:
    def test_with_provided_conversation(
        self,
        orchestrator,
        sample_conversation,
        decoded_token,
    ):
        with patch.object(
            orchestrator, "_perform_compression"
        ) as mock_perform:
            mock_perform.return_value = CompressionResult.success_no_compression([])

            orchestrator.compress_mid_execution(
                conversation_id="conv1",
                user_id="user1",
                model_id="gpt-4",
                decoded_token=decoded_token,
                current_conversation=sample_conversation,
            )

            mock_perform.assert_called_once_with(
                "conv1",
                sample_conversation,
                "gpt-4",
                decoded_token,
                user_id="user1",
                model_user_id=None,
            )

    def test_loads_conversation_when_not_provided(
        self,
        orchestrator,
        mock_conversation_service,
        sample_conversation,
        decoded_token,
    ):
        mock_conversation_service.get_conversation.return_value = sample_conversation

        with patch.object(
            orchestrator, "_perform_compression"
        ) as mock_perform:
            mock_perform.return_value = CompressionResult.success_no_compression([])

            orchestrator.compress_mid_execution(
                conversation_id="conv1",
                user_id="user1",
                model_id="gpt-4",
                decoded_token=decoded_token,
            )

            mock_conversation_service.get_conversation.assert_called_once_with(
                "conv1", "user1"
            )
            mock_perform.assert_called_once()

    def test_conversation_not_found_returns_failure(
        self,
        orchestrator,
        mock_conversation_service,
        decoded_token,
    ):
        mock_conversation_service.get_conversation.return_value = None

        result = orchestrator.compress_mid_execution(
            conversation_id="conv1",
            user_id="user1",
            model_id="gpt-4",
            decoded_token=decoded_token,
        )

        assert result.success is False
        assert "not found" in result.error

    def test_exception_returns_failure(
        self,
        orchestrator,
        mock_conversation_service,
        decoded_token,
    ):
        mock_conversation_service.get_conversation.side_effect = RuntimeError("fail")

        result = orchestrator.compress_mid_execution(
            conversation_id="conv1",
            user_id="user1",
            model_id="gpt-4",
            decoded_token=decoded_token,
        )

        assert result.success is False
        assert "fail" in result.error


@pytest.mark.unit
class TestOrchestratorInit:
    def test_default_threshold_checker(self, mock_conversation_service):
        orch = CompressionOrchestrator(
            conversation_service=mock_conversation_service
        )
        assert orch.threshold_checker is not None
        assert orch.conversation_service is mock_conversation_service

    def test_custom_threshold_checker(
        self, mock_conversation_service, mock_threshold_checker
    ):
        orch = CompressionOrchestrator(
            conversation_service=mock_conversation_service,
            threshold_checker=mock_threshold_checker,
        )
        assert orch.threshold_checker is mock_threshold_checker