"""Provider-level coverage tests for token_counter.py. Targets specific uncovered lines: - on_llm_start: model extraction fallbacks (kwargs direct, serialized kwargs model_name) - on_llm_start: provider extraction for unrecognized _type with provider kwarg - on_llm_end: token extraction from llm_output 'usage' key (alternative to 'token_usage') - on_llm_end: Ollama context overflow with zero original_prompt_estimate guard - _save_to_db: background thread success path with full token_data dict verification """ import threading from unittest.mock import MagicMock, patch from langchain_core.outputs import LLMResult from local_deep_research.metrics.token_counter import TokenCountingCallback # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_callback(research_context=None, research_id="rid-test", **overrides): """Build a TokenCountingCallback with controllable state.""" ctx = research_context if research_context is not None else {} cb = TokenCountingCallback(research_id=research_id, research_context=ctx) for k, v in overrides.items(): setattr(cb, k, v) return cb def _make_llm_result(llm_output=None, generations=None): """Build a minimal mock LLMResult.""" result = MagicMock(spec=LLMResult) result.llm_output = llm_output result.generations = generations or [] return result def _make_generation(usage_metadata=None, response_metadata=None): """Build a mock generation with a message carrying metadata.""" gen = MagicMock() msg = MagicMock() msg.usage_metadata = usage_metadata msg.response_metadata = ( response_metadata if response_metadata is not None else {} ) gen.message = msg return gen def _setup_model_counts(cb, model_name="test-model", provider="openai"): """Register model in cb.counts so on_llm_end can update them.""" cb.current_model = model_name cb.current_provider = provider cb.counts["by_model"][model_name] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": provider, } def _patch_worker_thread(): t = MagicMock() t.name = "WorkerThread-1" return patch.object(threading, "current_thread", return_value=t) # --------------------------------------------------------------------------- # 1. on_llm_start – model extraction fallback: kwargs["model"] when # invocation_params has no model key (line 163) # --------------------------------------------------------------------------- class TestModelFromKwargsModel: """Model extracted from kwargs['model'] when invocation_params is empty.""" def test_model_from_kwargs_model_key(self): """kwargs['model'] is used when invocation_params has no model.""" cb = _make_callback() cb.on_llm_start( {}, ["hello"], invocation_params={}, model="gemma-2-27b", ) assert cb.current_model == "gemma-2-27b" def test_model_from_kwargs_model_name_key(self): """kwargs['model_name'] is used when invocation_params and kwargs['model'] are absent.""" cb = _make_callback() cb.on_llm_start({}, ["hello"], model_name="codellama-34b") assert cb.current_model == "codellama-34b" # --------------------------------------------------------------------------- # 2. on_llm_start – model from serialized["kwargs"]["model_name"] # (line 167-169, specifically the model_name branch) # --------------------------------------------------------------------------- class TestModelFromSerializedKwargsModelName: """Model extracted from serialized['kwargs']['model_name'] when 'model' absent.""" def test_serialized_kwargs_model_name(self): """serialized['kwargs']['model_name'] provides the model when 'model' key missing.""" cb = _make_callback() cb.on_llm_start( {"kwargs": {"model_name": "deepseek-coder"}}, ["prompt"], ) assert cb.current_model == "deepseek-coder" # --------------------------------------------------------------------------- # 3. on_llm_start – provider from kwargs when _type is unrecognized # (line 210-211: else branch inside the if "_type" in serialized block) # --------------------------------------------------------------------------- class TestProviderFromKwargsWithUnrecognizedType: """When _type exists but is not ChatOllama/ChatOpenAI/ChatAnthropic, provider falls back to kwargs.get('provider', 'unknown').""" def test_unrecognized_type_with_provider_kwarg(self): """Unrecognized _type uses provider kwarg.""" cb = _make_callback() cb.on_llm_start( {"_type": "ChatGoogleGenAI", "kwargs": {"model": "gemini-pro"}}, ["prompt"], provider="google", ) assert cb.current_provider == "google" def test_unrecognized_type_without_provider_kwarg(self): """Unrecognized _type without provider kwarg defaults to 'unknown'.""" cb = _make_callback() cb.on_llm_start( {"_type": "ChatGoogleGenAI", "kwargs": {"model": "gemini-pro"}}, ["prompt"], ) assert cb.current_provider == "unknown" # --------------------------------------------------------------------------- # 4. on_llm_end – token extraction from llm_output 'usage' key # (line 240-241: the or branch: response.llm_output.get("usage", {})) # --------------------------------------------------------------------------- class TestTokenExtractionFromUsageKey: """Tokens extracted from llm_output['usage'] when 'token_usage' absent.""" def test_usage_key_provides_tokens(self): """llm_output['usage'] provides prompt_tokens/completion_tokens.""" cb = _make_callback() _setup_model_counts(cb, "gpt-4o", "openai") result = _make_llm_result( llm_output={ "usage": { "prompt_tokens": 45, "completion_tokens": 30, "total_tokens": 75, } }, ) cb.on_llm_end(result) assert cb.counts["total_prompt_tokens"] == 45 assert cb.counts["total_completion_tokens"] == 30 assert cb.counts["total_tokens"] == 75 # --------------------------------------------------------------------------- # 5. on_llm_end – Ollama context overflow: original_prompt_estimate == 0 # guards against division-by-zero (line 312-314) # --------------------------------------------------------------------------- class TestContextOverflowZeroPromptEstimate: """When original_prompt_estimate is 0 the truncation_ratio is 0, not a ZeroDivisionError.""" def test_zero_prompt_estimate_no_division_error(self): """truncation_ratio stays 0 when original_prompt_estimate is 0.""" cb = _make_callback(research_context={"context_limit": 100}) _setup_model_counts(cb, "llama3", "ollama") # context_limit is set during on_llm_start; set it directly here cb.context_limit = 100 cb.original_prompt_estimate = 0 gen = _make_generation( usage_metadata=None, response_metadata={ "prompt_eval_count": 96, # >= 95 (95% of 100) "eval_count": 20, }, ) result = _make_llm_result(generations=[[gen]]) cb.on_llm_end(result) assert cb.context_truncated is True # original_prompt_estimate (0) <= prompt_eval_count (96), so # the inner if branch is not entered; tokens_truncated stays 0 assert cb.tokens_truncated == 0 assert cb.truncation_ratio == 0.0 def test_context_overflow_calculates_truncation(self): """When prompt estimate > actual, truncation fields are populated.""" cb = _make_callback(research_context={"context_limit": 1000}) _setup_model_counts(cb, "llama3", "ollama") # context_limit is set during on_llm_start; set it directly here cb.context_limit = 1000 cb.original_prompt_estimate = 1200 gen = _make_generation( usage_metadata=None, response_metadata={ "prompt_eval_count": 960, # >= 950 (95% of 1000) "eval_count": 50, }, ) result = _make_llm_result(generations=[[gen]]) cb.on_llm_end(result) assert cb.context_truncated is True assert cb.tokens_truncated == 240 # 1200 - 960 expected_ratio = 240 / 1200 assert abs(cb.truncation_ratio - expected_ratio) < 1e-9 # --------------------------------------------------------------------------- # 6. _save_to_db – background thread success path # (lines 408-482): verifies metrics_writer is called with correct data # --------------------------------------------------------------------------- class TestSaveToDbBackgroundThread: """_save_to_db in a background thread writes via metrics_writer.""" def test_metrics_writer_called_with_token_data(self): """metrics_writer.write_token_metrics receives correct token_data dict.""" cb = _make_callback( research_context={ "username": "alice", "user_password": "secret", "research_query": "quantum computing", "research_mode": "detailed", "research_phase": "search", "search_iteration": 2, "search_engines_planned": ["google", "bing"], "search_engine_selected": "google", }, research_id="rid-42", ) cb.current_model = "gpt-4o" cb.current_provider = "openai" cb.response_time_ms = 350 cb.success_status = "success" cb.error_type = None cb.calling_file = "runner.py" cb.calling_function = "execute" cb.call_stack = "runner.py:execute:10" with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer" ) as mock_writer: cb._save_to_db(prompt_tokens=100, completion_tokens=50) mock_writer.set_user_password.assert_called_once_with( "alice", "secret" ) mock_writer.write_token_metrics.assert_called_once() call_args = mock_writer.write_token_metrics.call_args assert call_args[0][0] == "alice" assert call_args[0][1] == "rid-42" token_data = call_args[0][2] assert token_data["model_name"] == "gpt-4o" assert token_data["provider"] == "openai" assert token_data["prompt_tokens"] == 100 assert token_data["completion_tokens"] == 50 assert token_data["research_query"] == "quantum computing" assert token_data["research_mode"] == "detailed" assert token_data["research_phase"] == "search" assert token_data["search_iteration"] == 2 # search_engines_planned list should be JSON-serialized assert ( token_data["search_engines_planned"] == '["google", "bing"]' ) assert token_data["search_engine_selected"] == "google" assert token_data["response_time_ms"] == 350 assert token_data["calling_file"] == "runner.py" def test_no_username_logs_warning_and_returns(self): """Missing username in research_context causes early return.""" cb = _make_callback( research_context={"user_password": "secret"}, research_id="rid-99", ) cb.current_model = "gpt-4" cb.current_provider = "openai" with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer" ) as mock_writer: cb._save_to_db(prompt_tokens=10, completion_tokens=5) mock_writer.write_token_metrics.assert_not_called() def test_no_password_logs_warning_and_returns(self): """Missing password in research_context causes early return after username check.""" cb = _make_callback( research_context={"username": "bob"}, research_id="rid-88", ) cb.current_model = "gpt-4" cb.current_provider = "openai" with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer" ) as mock_writer: cb._save_to_db(prompt_tokens=10, completion_tokens=5) mock_writer.write_token_metrics.assert_not_called()