DocsGPT/tests/test_token_management.py

"""
Tests for token management and compression features.

NOTE: These tests are for future planned features that are not yet implemented.
They are skipped until the following modules are created:
- docsgpt.compression (DocumentCompressor, HistoryCompressor, etc.)
- docsgpt.core.token_budget (TokenBudgetManager)
"""
# ruff: noqa: F821
import pytest

pytest.skip(
    "Token management features not yet implemented - planned for future release",
    allow_module_level=True,
)


class TestTokenBudgetManager:
    """Test TokenBudgetManager functionality"""

    def test_calculate_budget(self):
        """Test budget calculation"""
        manager = TokenBudgetManager(model_id="gpt-4o")
        budget = manager.calculate_budget()

        assert budget.total_budget > 0
        assert budget.system_prompt > 0
        assert budget.chat_history > 0
        assert budget.retrieved_docs > 0

    def test_measure_usage(self):
        """Test token usage measurement"""
        manager = TokenBudgetManager(model_id="gpt-4o")

        usage = manager.measure_usage(
            system_prompt="You are a helpful assistant.",
            current_query="What is Python?",
            chat_history=[
                {"prompt": "Hello", "response": "Hi there!"},
                {"prompt": "How are you?", "response": "I'm doing well, thanks!"},
            ],
        )

        assert usage.total > 0
        assert usage.system_prompt > 0
        assert usage.current_query > 0
        assert usage.chat_history > 0

    def test_compression_recommendation(self):
        """Test compression recommendation generation"""
        manager = TokenBudgetManager(model_id="gpt-4o")

        # Create scenario with excessive history
        large_history = [
            {"prompt": f"Question {i}" * 100, "response": f"Answer {i}" * 100}
            for i in range(100)
        ]

        budget, usage, recommendation = manager.check_and_recommend(
            system_prompt="You are a helpful assistant.",
            current_query="What is Python?",
            chat_history=large_history,
        )

        # Should recommend compression
        assert recommendation.needs_compression()
        assert recommendation.compress_history


class TestHistoryCompressor:
    """Test HistoryCompressor functionality"""

    def test_sliding_window_compression(self):
        """Test sliding window compression strategy"""
        compressor = HistoryCompressor()

        history = [
            {"prompt": f"Question {i}", "response": f"Answer {i}"} for i in range(20)
        ]

        compressed, metadata = compressor.compress(
            history, target_tokens=500, strategy="sliding_window"
        )

        assert len(compressed) < len(history)
        assert metadata["original_messages"] == 20
        assert metadata["compressed_messages"] < 20
        assert metadata["strategy"] == "sliding_window"

    def test_preserve_tool_calls(self):
        """Test that tool calls are preserved during compression"""
        compressor = HistoryCompressor()

        history = [
            {"prompt": "Question 1", "response": "Answer 1"},
            {
                "prompt": "Use a tool",
                "response": "Tool used",
                "tool_calls": [{"tool_name": "search", "result": "Found something"}],
            },
            {"prompt": "Question 3", "response": "Answer 3"},
        ]

        compressed, metadata = compressor.compress(
            history, target_tokens=200, strategy="sliding_window", preserve_tool_calls=True
        )

        # Tool call message should be preserved
        has_tool_calls = any("tool_calls" in msg for msg in compressed)
        assert has_tool_calls


class TestDocumentCompressor:
    """Test DocumentCompressor functionality"""

    def test_rerank_compression(self):
        """Test re-ranking compression strategy"""
        compressor = DocumentCompressor()

        docs = [
            {"text": f"Document {i} with some content here" * 20, "title": f"Doc {i}"}
            for i in range(10)
        ]

        compressed, metadata = compressor.compress(
            docs, target_tokens=500, query="Document 5", strategy="rerank"
        )

        assert len(compressed) < len(docs)
        assert metadata["original_docs"] == 10
        assert metadata["strategy"] == "rerank"

    def test_excerpt_extraction(self):
        """Test excerpt extraction strategy"""
        compressor = DocumentCompressor()

        docs = [
            {
                "text": "This is a long document. " * 100
                + "Python is great. "
                + "More text here. " * 100,
                "title": "Python Guide",
            }
        ]

        compressed, metadata = compressor.compress(
            docs, target_tokens=300, query="Python", strategy="excerpt"
        )

        assert metadata["excerpts_created"] > 0
        # Excerpt should contain the query term
        assert "python" in compressed[0]["text"].lower()


class TestToolResultCompressor:
    """Test ToolResultCompressor functionality"""

    def test_truncate_large_results(self):
        """Test truncation of large tool results"""
        compressor = ToolResultCompressor()

        tool_results = [
            {
                "tool_name": "search",
                "result": "Very long result " * 1000,
                "arguments": {},
            }
        ]

        compressed, metadata = compressor.compress(
            tool_results, target_tokens=100, strategy="truncate"
        )

        assert metadata["results_truncated"] > 0
        # Result should be shorter
        compressed_result_len = len(str(compressed[0]["result"]))
        original_result_len = len(tool_results[0]["result"])
        assert compressed_result_len < original_result_len

    def test_extract_json_fields(self):
        """Test extraction of key fields from JSON results"""
        compressor = ToolResultCompressor()

        tool_results = [
            {
                "tool_name": "api_call",
                "result": {
                    "data": {"important": "value"},
                    "metadata": {"verbose": "information" * 100},
                    "debug": {"lots": "of data" * 100},
                },
                "arguments": {},
            }
        ]

        compressed, metadata = compressor.compress(
            tool_results, target_tokens=100, strategy="extract"
        )

        # Should keep important fields, discard verbose ones
        assert "data" in compressed[0]["result"]


class TestPromptOptimizer:
    """Test PromptOptimizer functionality"""

    def test_compress_tool_descriptions(self):
        """Test compression of tool descriptions"""
        optimizer = PromptOptimizer()

        tools = [
            {
                "type": "function",
                "function": {
                    "name": f"tool_{i}",
                    "description": "This is a very long description " * 50,
                    "parameters": {},
                },
            }
            for i in range(10)
        ]

        optimized, metadata = optimizer.optimize_tools(
            tools, target_tokens=500, strategy="compress"
        )

        assert metadata["optimized_tokens"] < metadata["original_tokens"]
        assert metadata["descriptions_compressed"] > 0

    def test_lazy_load_tools(self):
        """Test lazy loading of tools based on query"""
        optimizer = PromptOptimizer()

        tools = [
            {
                "type": "function",
                "function": {
                    "name": "search_tool",
                    "description": "Search for information",
                    "parameters": {},
                },
            },
            {
                "type": "function",
                "function": {
                    "name": "calculate_tool",
                    "description": "Perform calculations",
                    "parameters": {},
                },
            },
            {
                "type": "function",
                "function": {
                    "name": "other_tool",
                    "description": "Do something else",
                    "parameters": {},
                },
            },
        ]

        optimized, metadata = optimizer.optimize_tools(
            tools, target_tokens=200, query="I want to search for something", strategy="lazy_load"
        )

        # Should prefer search tool
        assert len(optimized) < len(tools)
        tool_names = [t["function"]["name"] for t in optimized]
        # Search tool should be included due to query relevance
        assert any("search" in name for name in tool_names)


def test_integration_compression_workflow():
    """Test complete compression workflow"""
    # Simulate a scenario with large inputs
    manager = TokenBudgetManager(model_id="gpt-4o")
    history_compressor = HistoryCompressor()
    doc_compressor = DocumentCompressor()

    # Large chat history
    history = [
        {"prompt": f"Question {i}" * 50, "response": f"Answer {i}" * 50}
        for i in range(50)
    ]

    # Large documents
    docs = [
        {"text": f"Document {i} content" * 100, "title": f"Doc {i}"} for i in range(20)
    ]

    # Check budget
    budget, usage, recommendation = manager.check_and_recommend(
        system_prompt="You are a helpful assistant.",
        current_query="What is Python?",
        chat_history=history,
        retrieved_docs=docs,
    )

    # Should need compression
    assert recommendation.needs_compression()

    # Apply compression
    if recommendation.compress_history:
        compressed_history, hist_meta = history_compressor.compress(
            history, recommendation.target_history_tokens or budget.chat_history
        )
        assert len(compressed_history) < len(history)

    if recommendation.compress_docs:
        compressed_docs, doc_meta = doc_compressor.compress(
            docs,
            recommendation.target_docs_tokens or budget.retrieved_docs,
            query="Python",
        )
        assert len(compressed_docs) < len(docs)