tests/test_token_optimizer.py

import pytest
from unittest.mock import patch, MagicMock

from src.utils.token_optimizer import (
    TokenOptimizer,
    _ApproximateEncoder,
    get_token_optimizer,
    reset_token_optimizer,
)


class TestTokenOptimizerInit:

    def test_default_max_tokens(self):
        optimizer = TokenOptimizer()
        assert optimizer._max_tokens == 4000

    def test_custom_max_tokens(self):
        optimizer = TokenOptimizer(max_tokens=2000)
        assert optimizer._max_tokens == 2000

    def test_encoder_lazy_loaded(self):
        optimizer = TokenOptimizer()
        assert optimizer._encoder is None

    def test_encoder_loaded_on_access(self):
        optimizer = TokenOptimizer()
        enc = optimizer.encoder
        assert enc is not None


class TestCountTokens:

    def test_empty_string_returns_zero(self):
        optimizer = TokenOptimizer()
        assert optimizer.count_tokens("") == 0

    def test_none_returns_zero(self):
        optimizer = TokenOptimizer()
        assert optimizer.count_tokens(None) == 0

    def test_simple_text_returns_positive(self):
        optimizer = TokenOptimizer()
        count = optimizer.count_tokens("Hello world this is a test")
        assert count > 0

    def test_longer_text_more_tokens(self):
        optimizer = TokenOptimizer()
        short = optimizer.count_tokens("hello")
        long = optimizer.count_tokens("hello world this is a much longer text with many words")
        assert long > short

    def test_consistent_results(self):
        optimizer = TokenOptimizer()
        text = "consistent token counting"
        assert optimizer.count_tokens(text) == optimizer.count_tokens(text)

    def test_turkish_text(self):
        optimizer = TokenOptimizer()
        count = optimizer.count_tokens("Satışları göster ve toplam tutarı hesapla")
        assert count > 0


class TestTruncateChatHistory:

    def test_empty_messages_returns_empty(self):
        optimizer = TokenOptimizer()
        result, tokens = optimizer.truncate_chat_history([])
        assert result == []
        assert tokens == 0

    def test_short_history_unchanged(self):
        optimizer = TokenOptimizer()
        messages = [
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "hi there"},
        ]
        result, tokens = optimizer.truncate_chat_history(messages, max_tokens=500)
        assert len(result) == 2
        assert tokens > 0

    def test_truncation_keeps_recent_messages(self):
        optimizer = TokenOptimizer()
        messages = [
            {"role": "user", "content": "old message " * 200},
            {"role": "assistant", "content": "old reply " * 200},
            {"role": "user", "content": "recent question"},
            {"role": "assistant", "content": "recent answer"},
        ]
        result, tokens = optimizer.truncate_chat_history(messages, max_tokens=50)
        assert any("recent" in m.get("content", "") for m in result)

    def test_never_drops_last_message(self):
        optimizer = TokenOptimizer()
        messages = [
            {"role": "user", "content": "very " * 500},
        ]
        result, tokens = optimizer.truncate_chat_history(messages, max_tokens=10)
        assert len(result) >= 1

    def test_preserves_message_pairs(self):
        optimizer = TokenOptimizer()
        messages = [
            {"role": "user", "content": "q1"},
            {"role": "assistant", "content": "a1"},
            {"role": "user", "content": "q2"},
            {"role": "assistant", "content": "a2"},
        ]
        result, _ = optimizer.truncate_chat_history(messages, max_tokens=500)
        roles = [m["role"] for m in result]
        if len(roles) >= 2:
            assert roles[0] == "user" or (roles[0] == "assistant" and len(result) == len(messages))

    def test_respects_token_budget(self):
        optimizer = TokenOptimizer()
        messages = [{"role": "user", "content": f"message number {i} " * 20} for i in range(20)]
        result, tokens = optimizer.truncate_chat_history(messages, max_tokens=100)
        assert tokens <= 200


class TestTruncateContext:

    def test_empty_context(self):
        optimizer = TokenOptimizer()
        result, tokens = optimizer.truncate_context("")
        assert result == ""
        assert tokens == 0

    def test_short_context_unchanged(self):
        optimizer = TokenOptimizer()
        text = "Short context text."
        result, tokens = optimizer.truncate_context(text, max_tokens=500)
        assert result == text

    def test_long_context_truncated(self):
        optimizer = TokenOptimizer()
        text = "This is a sentence. " * 500
        result, tokens = optimizer.truncate_context(text, max_tokens=50)
        assert tokens <= 60
        assert result.endswith("...")

    def test_preserve_start_true(self):
        optimizer = TokenOptimizer()
        text = "First sentence. Middle sentence. Last sentence."
        result, _ = optimizer.truncate_context(text, max_tokens=10, preserve_start=True)
        assert "First" in result

    def test_preserve_start_false(self):
        optimizer = TokenOptimizer()
        text = "First sentence here. " * 10 + "Last sentence here."
        result, _ = optimizer.truncate_context(text, max_tokens=15, preserve_start=False)
        assert result.startswith("...")

    def test_single_long_sentence_hard_truncated(self):
        optimizer = TokenOptimizer()
        text = "word " * 1000
        result, tokens = optimizer.truncate_context(text, max_tokens=20)
        assert tokens <= 30


class TestOptimizeSchemaContext:

    def test_empty_schema(self):
        optimizer = TokenOptimizer()
        assert optimizer.optimize_schema_context("") == ""

    def test_short_schema_unchanged(self):
        optimizer = TokenOptimizer()
        schema = "Table: users\nColumns: id, name"
        result = optimizer.optimize_schema_context(schema, max_tokens=500)
        assert result == schema

    def test_prioritizes_relevant_tables(self):
        optimizer = TokenOptimizer()
        schema = """Table: customers
Columns: id, name, city

Table: products
Columns: id, name, price

Table: sales
Columns: id, product_id, amount"""
        result = optimizer.optimize_schema_context(schema, relevant_tables=["sales"], max_tokens=50)
        assert "sales" in result

    def test_truncates_when_over_budget(self):
        optimizer = TokenOptimizer()
        schema = "\n\n".join([f"Table: table_{i}\nColumns: " + ", ".join([f"col_{j}" for j in range(20)]) for i in range(50)])
        result = optimizer.optimize_schema_context(schema, max_tokens=100)
        assert optimizer.count_tokens(result) <= 120

    def test_no_relevant_tables_keeps_order(self):
        optimizer = TokenOptimizer()
        schema = "Table: alpha\nColumns: a\n\nTable: beta\nColumns: b"
        result = optimizer.optimize_schema_context(schema, max_tokens=500)
        assert "alpha" in result
        assert "beta" in result


class TestParseTableBlocks:

    def test_parses_multiple_tables(self):
        optimizer = TokenOptimizer()
        schema = "Table: users\nColumns: id, name\n\nTable: orders\nColumns: id, user_id"
        blocks = optimizer._parse_table_blocks(schema)
        assert len(blocks) == 2
        assert blocks[0][0] == "users"
        assert blocks[1][0] == "orders"

    def test_empty_schema_returns_empty(self):
        optimizer = TokenOptimizer()
        assert optimizer._parse_table_blocks("") == []

    def test_single_table(self):
        optimizer = TokenOptimizer()
        schema = "Table: users\nColumns: id, name"
        blocks = optimizer._parse_table_blocks(schema)
        assert len(blocks) == 1

    def test_no_table_prefix_returns_empty(self):
        optimizer = TokenOptimizer()
        assert optimizer._parse_table_blocks("just some text") == []


class TestSelectFewShotExamples:

    def test_empty_examples(self):
        optimizer = TokenOptimizer()
        assert optimizer.select_few_shot_examples([], "query") == []

    def test_selects_relevant_examples(self):
        optimizer = TokenOptimizer()
        examples = [
            {"query": "show total sales amount", "sql": "SELECT SUM(amount) FROM sales"},
            {"query": "list all customers", "sql": "SELECT * FROM customers"},
            {"query": "calculate average sales price", "sql": "SELECT AVG(price) FROM sales"},
        ]
        result = optimizer.select_few_shot_examples(examples, "total sales", max_tokens=500)
        assert len(result) > 0
        assert any("sales" in ex.get("query", "") and "total" in ex.get("query", "") for ex in result)

    def test_respects_token_budget(self):
        optimizer = TokenOptimizer()
        examples = [{"query": f"example {i} " * 50, "sql": f"SELECT {i}"} for i in range(20)]
        result = optimizer.select_few_shot_examples(examples, "example 1", max_tokens=100)
        total = sum(optimizer.count_tokens(optimizer._example_text(ex)) for ex in result)
        assert total <= 200

    def test_ranks_by_keyword_similarity(self):
        optimizer = TokenOptimizer()
        examples = [
            {"query": "customer list", "sql": "SELECT * FROM customers"},
            {"query": "sales report total", "sql": "SELECT SUM(total) FROM sales"},
        ]
        result = optimizer.select_few_shot_examples(examples, "sales total report", max_tokens=500)
        assert result[0]["query"] == "sales report total"


class TestKeywordSimilarity:

    def test_identical_text(self):
        optimizer = TokenOptimizer()
        assert optimizer._keyword_similarity("hello world", "hello world") == 1.0

    def test_no_overlap(self):
        optimizer = TokenOptimizer()
        assert optimizer._keyword_similarity("hello", "goodbye") == 0.0

    def test_partial_overlap(self):
        optimizer = TokenOptimizer()
        score = optimizer._keyword_similarity("hello world", "hello there")
        assert 0 < score < 1

    def test_empty_query(self):
        optimizer = TokenOptimizer()
        assert optimizer._keyword_similarity("", "hello") == 0.0

    def test_case_insensitive(self):
        optimizer = TokenOptimizer()
        assert optimizer._keyword_similarity("Hello", "hello") == 1.0


class TestOptimizeForLlmCall:

    def test_all_components_present(self):
        optimizer = TokenOptimizer()
        components = {
            "query": "what are total sales?",
            "system_prompt": "you are a sql assistant",
            "schema": "Table: sales\nColumns: id, amount",
            "few_shots": "Example: SELECT COUNT(*) FROM sales",
            "chat_history": "user: show sales\nassistant: here are sales",
        }
        result = optimizer.optimize_for_llm_call(components, total_budget=3000)
        assert "query" in result
        assert "system_prompt" in result
        assert "schema" in result
        assert "few_shots" in result
        assert "chat_history" in result

    def test_query_preserved(self):
        optimizer = TokenOptimizer()
        components = {"query": "full query text", "system_prompt": "sp", "schema": "s" * 5000}
        result = optimizer.optimize_for_llm_call(components, total_budget=100)
        assert result["query"] == "full query text"

    def test_system_prompt_preserved(self):
        optimizer = TokenOptimizer()
        components = {"query": "q", "system_prompt": "system prompt text", "schema": "s" * 5000}
        result = optimizer.optimize_for_llm_call(components, total_budget=100)
        assert result["system_prompt"] == "system prompt text"

    def test_empty_components(self):
        optimizer = TokenOptimizer()
        result = optimizer.optimize_for_llm_call({}, total_budget=1000)
        assert result["query"] == ""
        assert result["system_prompt"] == ""

    def test_large_schema_truncated(self):
        optimizer = TokenOptimizer()
        components = {
            "query": "q",
            "system_prompt": "sp",
            "schema": "Table: big\nColumns: " + ", ".join([f"col_{i}" for i in range(500)]),
        }
        result = optimizer.optimize_for_llm_call(components, total_budget=100)
        assert optimizer.count_tokens(result["schema"]) < optimizer.count_tokens(components["schema"])


class TestGetOptimizationStats:

    def test_stats_structure(self):
        optimizer = TokenOptimizer()
        original = {"query": "hello world", "schema": "big " * 100}
        optimized = {"query": "hello world", "schema": "big"}
        stats = optimizer.get_optimization_stats(original, optimized)
        assert "original_tokens" in stats
        assert "optimized_tokens" in stats
        assert "tokens_saved" in stats
        assert "reduction_ratio" in stats

    def test_reduction_calculated(self):
        optimizer = TokenOptimizer()
        original = {"text": "word " * 100}
        optimized = {"text": "word " * 10}
        stats = optimizer.get_optimization_stats(original, optimized)
        assert stats["tokens_saved"] > 0
        assert stats["reduction_ratio"] > 0

    def test_no_reduction(self):
        optimizer = TokenOptimizer()
        data = {"text": "same"}
        stats = optimizer.get_optimization_stats(data, data)
        assert stats["tokens_saved"] == 0
        assert stats["reduction_ratio"] == 0.0

    def test_empty_input(self):
        optimizer = TokenOptimizer()
        stats = optimizer.get_optimization_stats({}, {})
        assert stats["original_tokens"] == 0


class TestApproximateEncoder:

    def test_encode_returns_list(self):
        enc = _ApproximateEncoder()
        result = enc.encode("hello world")
        assert isinstance(result, list)

    def test_decode_returns_string(self):
        enc = _ApproximateEncoder()
        result = enc.decode([1, 2, 3])
        assert isinstance(result, str)

    def test_approximate_proportional(self):
        enc = _ApproximateEncoder()
        short = len(enc.encode("hi"))
        long = len(enc.encode("hello world this is longer"))
        assert long > short


class TestSingleton:

    def setup_method(self):
        reset_token_optimizer()

    def teardown_method(self):
        reset_token_optimizer()

    def test_get_returns_instance(self):
        optimizer = get_token_optimizer()
        assert isinstance(optimizer, TokenOptimizer)

    def test_get_returns_same_instance(self):
        first = get_token_optimizer()
        second = get_token_optimizer()
        assert first is second

    def test_reset_clears_instance(self):
        first = get_token_optimizer()
        reset_token_optimizer()
        second = get_token_optimizer()
        assert first is not second

    def test_get_with_custom_max_tokens(self):
        optimizer = get_token_optimizer(max_tokens=2000)
        assert optimizer._max_tokens == 2000


class TestExampleText:

    def test_extracts_query_and_sql(self):
        optimizer = TokenOptimizer()
        ex = {"query": "show sales", "sql": "SELECT * FROM sales"}
        text = optimizer._example_text(ex)
        assert "show sales" in text
        assert "SELECT * FROM sales" in text

    def test_handles_missing_keys(self):
        optimizer = TokenOptimizer()
        ex = {"other": "data"}
        text = optimizer._example_text(ex)
        assert len(text) > 0

    def test_empty_example(self):
        optimizer = TokenOptimizer()
        ex = {}
        text = optimizer._example_text(ex)
        assert isinstance(text, str)


class TestEdgeCases:

    def test_count_tokens_with_special_characters(self):
        optimizer = TokenOptimizer()
        count = optimizer.count_tokens("!@#$%^&*()_+{}|:<>?")
        assert count > 0

    def test_truncate_context_single_word(self):
        optimizer = TokenOptimizer()
        result, _ = optimizer.truncate_context("word", max_tokens=500)
        assert result == "word"

    def test_schema_with_no_table_prefix(self):
        optimizer = TokenOptimizer()
        result = optimizer.optimize_schema_context("just plain text", max_tokens=500)
        assert result == "just plain text"

    def test_few_shot_with_single_example(self):
        optimizer = TokenOptimizer()
        examples = [{"query": "test", "sql": "SELECT 1"}]
        result = optimizer.select_few_shot_examples(examples, "test", max_tokens=500)
        assert len(result) == 1

    def test_optimize_llm_call_missing_optional_components(self):
        optimizer = TokenOptimizer()
        components = {"query": "hello"}
        result = optimizer.optimize_for_llm_call(components, total_budget=1000)
        assert result["schema"] == ""
        assert result["few_shots"] == ""
        assert result["chat_history"] == ""

    def test_truncate_history_with_empty_content(self):
        optimizer = TokenOptimizer()
        messages = [{"role": "user", "content": ""}, {"role": "assistant", "content": "reply"}]
        result, _ = optimizer.truncate_chat_history(messages, max_tokens=500)
        assert len(result) == 2

    def test_preserve_message_pairs_no_orphan_assistant(self):
        optimizer = TokenOptimizer()
        messages = [
            {"role": "user", "content": "q1 " * 200},
            {"role": "assistant", "content": "a1 " * 200},
            {"role": "user", "content": "q2"},
            {"role": "assistant", "content": "a2"},
        ]
        result, _ = optimizer.truncate_chat_history(messages, max_tokens=20)
        if result and result[0].get("role") == "assistant":
            idx = messages.index(result[0])
            assert idx > 0

    @patch("src.utils.token_optimizer.TokenOptimizer.encoder", new_callable=lambda: property(lambda self: _ApproximateEncoder()))
    def test_fallback_encoder_works(self, _):
        optimizer = TokenOptimizer()
        optimizer._encoder = _ApproximateEncoder()
        count = optimizer.count_tokens("hello world test")
        assert count >= 0


class TestConfigIntegration:

    def test_config_constants_available(self):
        from utils.token_optimizer import NL2SQL_MAX_CONTEXT_TOKENS, CHAT_HISTORY_MAX_TOKENS, SCHEMA_CONTEXT_MAX_TOKENS, FEW_SHOT_MAX_TOKENS
        assert NL2SQL_MAX_CONTEXT_TOKENS > 0
        assert CHAT_HISTORY_MAX_TOKENS > 0
        assert SCHEMA_CONTEXT_MAX_TOKENS > 0
        assert FEW_SHOT_MAX_TOKENS > 0

    def test_optimizer_uses_config_defaults(self):
        optimizer = TokenOptimizer()
        messages = [{"role": "user", "content": "test"}]
        result, _ = optimizer.truncate_chat_history(messages)
        assert len(result) == 1