"""Tests for TokenCountingCallback real behavior - covering token extraction, model detection, provider detection, context overflow, and error handling.""" import time from unittest.mock import Mock, patch from langchain_core.outputs import LLMResult from local_deep_research.metrics.token_counter import ( TokenCounter, TokenCountingCallback, ) class TestTokenCountingCallbackModelDetection: """Tests for model name detection from various sources.""" def test_preset_model_takes_priority(self): """Preset model name should override all other detection.""" callback = TokenCountingCallback() callback.preset_model = "my-preset-model" callback.preset_provider = "openai" serialized = {"name": "different-model", "_type": "ChatOpenAI"} callback.on_llm_start(serialized, ["test prompt"]) assert callback.current_model == "my-preset-model" def test_model_from_invocation_params(self): """Model should be extracted from invocation_params.""" callback = TokenCountingCallback() serialized = {} kwargs = {"invocation_params": {"model": "gpt-4-turbo"}} callback.on_llm_start(serialized, ["test"], **kwargs) assert callback.current_model == "gpt-4-turbo" def test_model_from_serialized_kwargs(self): """Model should be extracted from serialized kwargs.""" callback = TokenCountingCallback() serialized = {"kwargs": {"model": "claude-3-opus"}} callback.on_llm_start(serialized, ["test"]) assert callback.current_model == "claude-3-opus" def test_model_from_serialized_name(self): """Model should be extracted from serialized name.""" callback = TokenCountingCallback() serialized = {"name": "my-model"} callback.on_llm_start(serialized, ["test"]) assert callback.current_model == "my-model" def test_ollama_model_extraction(self): """Ollama model should be extracted from serialized type and kwargs.""" callback = TokenCountingCallback() serialized = { "_type": "ChatOllama", "kwargs": {"model": "llama3:8b"}, } callback.on_llm_start(serialized, ["test"]) assert callback.current_model == "llama3:8b" def test_ollama_type_without_model_name(self): """Ollama type without model name should default to 'ollama'.""" callback = TokenCountingCallback() serialized = {"_type": "ChatOllama"} callback.on_llm_start(serialized, ["test"]) assert callback.current_model == "ollama" def test_unknown_model_fallback(self): """Unknown model should fall back to 'unknown'.""" callback = TokenCountingCallback() serialized = {} callback.on_llm_start(serialized, ["test"]) assert callback.current_model == "unknown" def test_type_as_model_fallback(self): """_type should be used as model name if no other source found.""" callback = TokenCountingCallback() serialized = {"_type": "CustomLLM"} callback.on_llm_start(serialized, ["test"]) assert callback.current_model == "CustomLLM" class TestTokenCountingCallbackProviderDetection: """Tests for provider detection from various sources.""" def test_preset_provider_takes_priority(self): """Preset provider should override detection.""" callback = TokenCountingCallback() callback.preset_provider = "anthropic" serialized = {"_type": "ChatOpenAI"} callback.on_llm_start(serialized, ["test"]) assert callback.current_provider == "anthropic" def test_openai_provider_detection(self): """ChatOpenAI type should detect 'openai' provider.""" callback = TokenCountingCallback() serialized = {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}} callback.on_llm_start(serialized, ["test"]) assert callback.current_provider == "openai" def test_anthropic_provider_detection(self): """ChatAnthropic type should detect 'anthropic' provider.""" callback = TokenCountingCallback() serialized = { "_type": "ChatAnthropic", "kwargs": {"model": "claude-3-opus"}, } callback.on_llm_start(serialized, ["test"]) assert callback.current_provider == "anthropic" def test_ollama_provider_detection(self): """ChatOllama type should detect 'ollama' provider.""" callback = TokenCountingCallback() serialized = { "_type": "ChatOllama", "kwargs": {"model": "llama3"}, } callback.on_llm_start(serialized, ["test"]) assert callback.current_provider == "ollama" def test_unknown_provider_fallback(self): """Unknown type should fall back to 'unknown' provider.""" callback = TokenCountingCallback() serialized = {} callback.on_llm_start(serialized, ["test"]) assert callback.current_provider == "unknown" class TestTokenCountingCallbackTokenExtraction: """Tests for token extraction from LLM responses.""" def test_token_extraction_from_llm_output(self): """Tokens should be extracted from response.llm_output.""" callback = TokenCountingCallback() callback.current_model = "gpt-4" callback.counts["by_model"]["gpt-4"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "openai", } response = Mock(spec=LLMResult) response.llm_output = { "token_usage": { "prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150, } } response.generations = [] callback.on_llm_end(response) assert callback.counts["total_prompt_tokens"] == 100 assert callback.counts["total_completion_tokens"] == 50 assert callback.counts["total_tokens"] == 150 def test_token_extraction_from_usage_metadata(self): """Tokens should be extracted from generation.message.usage_metadata.""" callback = TokenCountingCallback() callback.current_model = "llama3" callback.counts["by_model"]["llama3"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "ollama", } # Create mock generation with usage_metadata mock_message = Mock() mock_message.usage_metadata = { "input_tokens": 200, "output_tokens": 80, "total_tokens": 280, } mock_message.response_metadata = {} mock_generation = Mock() mock_generation.message = mock_message response = Mock(spec=LLMResult) response.llm_output = None response.generations = [[mock_generation]] callback.on_llm_end(response) assert callback.counts["total_prompt_tokens"] == 200 assert callback.counts["total_completion_tokens"] == 80 assert callback.counts["total_tokens"] == 280 def test_token_extraction_from_response_metadata_ollama(self): """Tokens should be extracted from Ollama response_metadata.""" callback = TokenCountingCallback() callback.current_model = "llama3" callback.counts["by_model"]["llama3"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "ollama", } mock_message = Mock() mock_message.usage_metadata = None mock_message.response_metadata = { "prompt_eval_count": 300, "eval_count": 120, "total_duration": 5000000000, "load_duration": 100000000, "prompt_eval_duration": 2000000000, "eval_duration": 2900000000, } mock_generation = Mock() mock_generation.message = mock_message response = Mock(spec=LLMResult) response.llm_output = None response.generations = [[mock_generation]] callback.on_llm_end(response) assert callback.counts["total_prompt_tokens"] == 300 assert callback.counts["total_completion_tokens"] == 120 assert callback.counts["total_tokens"] == 420 def test_no_token_usage_available(self): """No token usage data should not crash and not update counts.""" callback = TokenCountingCallback() callback.current_model = "test" callback.counts["by_model"]["test"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "unknown", } response = Mock(spec=LLMResult) response.llm_output = None response.generations = [] callback.on_llm_end(response) assert callback.counts["total_tokens"] == 0 def test_cumulative_token_counting(self): """Multiple on_llm_end calls should accumulate token counts.""" callback = TokenCountingCallback() callback.current_model = "gpt-4" callback.counts["by_model"]["gpt-4"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "openai", } for i in range(3): response = Mock(spec=LLMResult) response.llm_output = { "token_usage": { "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, } } response.generations = [] callback.on_llm_end(response) assert callback.counts["total_tokens"] == 45 assert callback.counts["total_prompt_tokens"] == 30 assert callback.counts["total_completion_tokens"] == 15 class TestTokenCountingCallbackContextOverflow: """Tests for context overflow detection.""" def test_context_overflow_detected(self): """Context overflow should be detected when near context limit.""" callback = TokenCountingCallback( research_context={"context_limit": 4096} ) callback.current_model = "llama3" callback.counts["by_model"]["llama3"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "ollama", } # Estimate original prompt is much larger callback.original_prompt_estimate = 5000 # Simulate start to set context limit callback.on_llm_start({"_type": "ChatOllama"}, ["x" * 20000]) mock_message = Mock() mock_message.usage_metadata = None mock_message.response_metadata = { "prompt_eval_count": 3900, # 95% of 4096 "eval_count": 100, } mock_generation = Mock() mock_generation.message = mock_message response = Mock(spec=LLMResult) response.llm_output = None response.generations = [[mock_generation]] callback.on_llm_end(response) assert callback.context_truncated is True assert callback.tokens_truncated > 0 def test_no_overflow_when_below_threshold(self): """No overflow when prompt tokens are below 95% of context limit.""" callback = TokenCountingCallback( research_context={"context_limit": 4096} ) callback.current_model = "llama3" callback.counts["by_model"]["llama3"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "ollama", } callback.on_llm_start({"_type": "ChatOllama"}, ["x" * 4000]) mock_message = Mock() mock_message.usage_metadata = None mock_message.response_metadata = { "prompt_eval_count": 2000, # Well below 95% of 4096 "eval_count": 100, } mock_generation = Mock() mock_generation.message = mock_message response = Mock(spec=LLMResult) response.llm_output = None response.generations = [[mock_generation]] callback.on_llm_end(response) assert callback.context_truncated is False def test_no_overflow_without_context_limit(self): """No overflow detection when context_limit is not set.""" callback = TokenCountingCallback() callback.current_model = "gpt-4" callback.counts["by_model"]["gpt-4"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "openai", } mock_message = Mock() mock_message.usage_metadata = None mock_message.response_metadata = { "prompt_eval_count": 100000, "eval_count": 100, } mock_generation = Mock() mock_generation.message = mock_message response = Mock(spec=LLMResult) response.llm_output = None response.generations = [[mock_generation]] callback.on_llm_end(response) assert callback.context_truncated is False class TestTokenCountingCallbackMissingUsageData: """Tests for recording calls when the provider reports no usage (#4457).""" @staticmethod def _response_without_usage(): response = Mock(spec=LLMResult) response.llm_output = None response.generations = [] return response def test_no_usage_with_research_id_still_saves_to_db(self): """A call without usage data must still be recorded (zero counts).""" callback = TokenCountingCallback(research_id="research-123") callback.current_model = "test-model" callback._save_to_db = Mock() callback.on_llm_end(self._response_without_usage()) callback._save_to_db.assert_called_once_with(0, 0) def test_no_usage_without_research_id_does_not_save(self): """Without a research_id there is nothing to persist.""" callback = TokenCountingCallback() callback.current_model = "test-model" callback._save_to_db = Mock() callback.on_llm_end(self._response_without_usage()) callback._save_to_db.assert_not_called() def test_no_usage_does_not_update_counts(self): """Zero-count recording must not inflate in-memory token counts.""" callback = TokenCountingCallback(research_id="research-123") callback.current_model = "test-model" callback._save_to_db = Mock() callback.on_llm_end(self._response_without_usage()) assert callback.counts["total_tokens"] == 0 assert callback.counts["total_prompt_tokens"] == 0 assert callback.counts["total_completion_tokens"] == 0 def test_repeated_no_usage_records_every_call_warns_once(self): """Every no-usage call is recorded, but the warning fires once.""" import local_deep_research.metrics.token_counter as tc_mod callback = TokenCountingCallback(research_id="research-123") callback.current_model = "test-model" callback._save_to_db = Mock() with patch.object(tc_mod.logger, "warning") as mock_warn: for _ in range(5): callback.on_llm_end(self._response_without_usage()) assert callback._save_to_db.call_count == 5 mock_warn.assert_called_once() def test_usage_present_saves_actual_counts(self): """Sanity check: usage data still saves the real counts.""" callback = TokenCountingCallback(research_id="research-123") callback.current_model = "gpt-4" callback.counts["by_model"]["gpt-4"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "openai", } callback._save_to_db = Mock() response = Mock(spec=LLMResult) response.llm_output = { "token_usage": { "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, } } response.generations = [] callback.on_llm_end(response) callback._save_to_db.assert_called_once_with(10, 5) class TestTokenCountingCallbackErrorHandling: """Tests for error tracking in callback.""" def test_error_sets_status_and_type(self): """on_llm_error should set error status and type.""" callback = TokenCountingCallback() callback.start_time = time.time() error = ValueError("test error") callback.on_llm_error(error) assert callback.success_status == "error" assert callback.error_type == "ValueError" def test_error_calculates_response_time(self): """on_llm_error should calculate response time.""" callback = TokenCountingCallback() callback.start_time = time.time() - 0.5 # 500ms ago callback.on_llm_error(RuntimeError("test")) assert callback.response_time_ms is not None assert callback.response_time_ms >= 400 # At least 400ms def test_error_without_start_time(self): """on_llm_error without start_time should not crash.""" callback = TokenCountingCallback() callback.on_llm_error(RuntimeError("test")) assert callback.success_status == "error" assert callback.response_time_ms is None class TestTokenCountingCallbackTimingAndCallStack: """Tests for timing and call stack tracking.""" def test_start_time_set_on_llm_start(self): """on_llm_start should record start time.""" callback = TokenCountingCallback() before = time.time() callback.on_llm_start({}, ["test"]) after = time.time() assert before <= callback.start_time <= after def test_response_time_calculated_on_end(self): """on_llm_end should calculate response time in ms.""" callback = TokenCountingCallback() callback.current_model = "test" callback.counts["by_model"]["test"] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": "unknown", } callback.start_time = time.time() - 0.1 # 100ms ago response = Mock(spec=LLMResult) response.llm_output = None response.generations = [] callback.on_llm_end(response) assert callback.response_time_ms is not None assert callback.response_time_ms >= 50 # At least 50ms def test_prompt_estimate_from_prompts(self): """original_prompt_estimate should be set from prompt length.""" callback = TokenCountingCallback() callback.on_llm_start({}, ["Hello world!"]) # 12 chars ~ 3 tokens assert callback.original_prompt_estimate == 3 # 12 // 4 def test_prompt_estimate_multiple_prompts(self): """Multiple prompts should sum their character counts for estimate.""" callback = TokenCountingCallback() callback.on_llm_start({}, ["aaaa", "bbbb", "cccc"]) # 12 chars total assert callback.original_prompt_estimate == 3 # 12 // 4 def test_call_count_incremented(self): """Call count should increment on each on_llm_start.""" callback = TokenCountingCallback() callback.on_llm_start( {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}}, ["test"] ) callback.on_llm_start( {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}}, ["test"] ) assert callback.counts["by_model"]["gpt-4"]["calls"] == 2 class TestTokenCountingCallbackGetContextOverflowFields: """Tests for _get_context_overflow_fields helper.""" def test_fields_when_no_overflow(self): """Fields should indicate no overflow when not truncated.""" callback = TokenCountingCallback() fields = callback._get_context_overflow_fields() assert fields["context_truncated"] is False assert fields["tokens_truncated"] is None assert fields["truncation_ratio"] is None def test_fields_when_overflow(self): """Fields should contain overflow data when truncated.""" callback = TokenCountingCallback() callback.context_limit = 4096 callback.context_truncated = True callback.tokens_truncated = 500 callback.truncation_ratio = 0.12 callback.ollama_metrics = { "prompt_eval_count": 3900, "eval_count": 100, } fields = callback._get_context_overflow_fields() assert fields["context_truncated"] is True assert fields["tokens_truncated"] == 500 assert fields["truncation_ratio"] == 0.12 assert fields["context_limit"] == 4096 assert fields["ollama_prompt_eval_count"] == 3900 class TestTokenCounterManager: """Tests for TokenCounter manager class.""" def test_create_callback_returns_callback(self): """create_callback should return a TokenCountingCallback instance.""" counter = TokenCounter() callback = counter.create_callback( research_id="test-123", research_context={"key": "value"}, ) assert isinstance(callback, TokenCountingCallback) assert callback.research_id == "test-123" assert callback.research_context == {"key": "value"} def test_create_callback_without_args(self): """create_callback without args should work.""" counter = TokenCounter() callback = counter.create_callback() assert isinstance(callback, TokenCountingCallback) assert callback.research_id is None def test_empty_metrics_structure(self): """_get_empty_metrics should return proper structure.""" counter = TokenCounter() metrics = counter._get_empty_metrics() assert metrics["total_tokens"] == 0 assert metrics["total_researches"] == 0 assert metrics["by_model"] == [] assert metrics["recent_researches"] == [] assert "token_breakdown" in metrics