"""Coverage tests for token_counter.py — LLM callback paths and save_to_db branches. Targets the specific uncovered lines: - on_llm_start: call stack capture, model extraction from invocation_params, Ollama _type detection - on_llm_end: usage_metadata from generations (Ollama-specific), context overflow detection (prompt_eval_count >= 80% of limit), raw Ollama response_metadata metrics extraction - on_llm_error: error status and type recorded - _save_to_db (background thread): no username -> warning + return early; no password -> warning + return early; success write path """ import threading import time from unittest.mock import MagicMock, patch from langchain_core.outputs import LLMResult from local_deep_research.metrics.token_counter import ( TokenCountingCallback, ) _MOD = "local_deep_research.metrics.token_counter" # --------------------------------------------------------------------------- # Helpers shared across test classes # --------------------------------------------------------------------------- def _make_callback(research_context=None, research_id="rid-1", **overrides): """Build a TokenCountingCallback with controllable state.""" ctx = ( research_context if research_context is not None else {"research_query": "q"} ) 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 _patch_worker_thread(): t = MagicMock() t.name = "WorkerThread-1" return patch.object(threading, "current_thread", return_value=t) def _patch_main_thread(): t = MagicMock() t.name = "MainThread" return patch.object(threading, "current_thread", return_value=t) 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.counts["by_model"][model_name] = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 1, "provider": provider, } # --------------------------------------------------------------------------- # 1. test_on_llm_start_call_stack_capture # --------------------------------------------------------------------------- def _fake_frame(filename, function, lineno=1): """Build a NamedTuple-like object that mimics inspect.FrameInfo. Using a plain object with real string attributes avoids the MagicMock attribute-chaining that causes pathlib to recurse into heavy imports. """ class _FI: pass fi = _FI() fi.filename = filename fi.function = function fi.lineno = lineno return fi class TestOnLlmStartCallStackCapture: """Verify that inspect.stack parsing populates call-stack tracking fields.""" def test_call_stack_populated_from_project_frame(self): """When a frame in local_deep_research (not site-packages/venv) is on the stack, calling_file and calling_function should be set.""" cb = _make_callback() # inspect.stack() returns list[FrameInfo]; element [0] is the method # itself (skipped), then [1:] are callers. project_frame = _fake_frame( "/home/user/local_deep_research/src/local_deep_research/runners/runner.py", "execute_research", 42, ) sentinel_frame = _fake_frame("/sentinel.py", "sentinel") with patch( f"{_MOD}.inspect.stack", return_value=[sentinel_frame, project_frame], ): cb.on_llm_start({}, ["hello"]) assert cb.calling_file is not None assert cb.calling_function == "execute_research" def test_call_stack_skips_site_packages_frames(self): """Frames from site-packages are skipped; if no project frame found, call_stack stays None.""" cb = _make_callback() site_frame = _fake_frame( "/usr/lib/python3/site-packages/langchain/llms/base.py", "_generate", 10, ) sentinel_frame = _fake_frame("/sentinel.py", "sentinel") with patch( f"{_MOD}.inspect.stack", return_value=[sentinel_frame, site_frame] ): cb.on_llm_start({}, ["hello"]) # No project frame → call stack fields remain None assert cb.calling_file is None assert cb.calling_function is None assert cb.call_stack is None def test_call_stack_uses_src_local_deep_research_split(self): """Paths containing src/local_deep_research are split on that segment.""" cb = _make_callback() project_frame = _fake_frame( "/home/user/repo/src/local_deep_research/metrics/token_counter.py", "on_llm_start", 70, ) sentinel_frame = _fake_frame("/sentinel.py", "sentinel") with patch( f"{_MOD}.inspect.stack", return_value=[sentinel_frame, project_frame], ): cb.on_llm_start({}, ["test prompt"]) assert cb.calling_file is not None assert "src/local_deep_research" not in cb.calling_file def test_call_stack_uses_local_deep_research_src_split(self): """Paths with local_deep_research/src are handled by the second branch.""" cb = _make_callback() project_frame = _fake_frame( "/home/user/local_deep_research/src/module/file.py", "some_function", 5, ) sentinel_frame = _fake_frame("/sentinel.py", "sentinel") with patch( f"{_MOD}.inspect.stack", return_value=[sentinel_frame, project_frame], ): cb.on_llm_start({}, ["p"]) assert cb.calling_function == "some_function" def test_call_stack_graceful_on_inspect_exception(self): """If inspect.stack raises, the callback continues without call stack info.""" cb = _make_callback() with patch( f"{_MOD}.inspect.stack", side_effect=RuntimeError("inspect fail") ): cb.on_llm_start({}, ["prompt"]) assert cb.calling_file is None assert cb.calling_function is None def test_call_stack_string_joined_with_arrow(self): """Multiple project frames produce a ' -> ' joined call_stack string.""" cb = _make_callback() frame_a = _fake_frame( "/home/user/local_deep_research/src/local_deep_research/a.py", "func_a", 10, ) frame_b = _fake_frame( "/home/user/local_deep_research/src/local_deep_research/b.py", "func_b", 20, ) sentinel_frame = _fake_frame("/sentinel.py", "sentinel") with patch( f"{_MOD}.inspect.stack", return_value=[sentinel_frame, frame_a, frame_b], ): cb.on_llm_start({}, ["prompt"]) if cb.call_stack: assert " -> " in cb.call_stack # --------------------------------------------------------------------------- # 2. test_on_llm_start_model_from_invocation_params # --------------------------------------------------------------------------- class TestOnLlmStartModelFromInvocationParams: """Various model-name extraction paths in on_llm_start.""" def test_model_key_in_invocation_params(self): cb = _make_callback() cb.on_llm_start({}, ["prompt"], invocation_params={"model": "gpt-4o"}) assert cb.current_model == "gpt-4o" def test_model_name_key_in_invocation_params(self): cb = _make_callback() cb.on_llm_start( {}, ["prompt"], invocation_params={"model_name": "claude-3-opus"} ) assert cb.current_model == "claude-3-opus" def test_model_from_kwargs_directly(self): cb = _make_callback() # model not in invocation_params, but passed as direct kwarg cb.on_llm_start({}, ["prompt"], model="gemini-pro") assert cb.current_model == "gemini-pro" def test_model_name_from_kwargs_directly(self): cb = _make_callback() cb.on_llm_start({}, ["prompt"], model_name="mistral-large") assert cb.current_model == "mistral-large" def test_model_from_serialized_kwargs(self): cb = _make_callback() cb.on_llm_start({"kwargs": {"model": "llama3"}}, ["prompt"]) assert cb.current_model == "llama3" def test_model_name_from_serialized_kwargs(self): cb = _make_callback() cb.on_llm_start({"kwargs": {"model_name": "phi-3"}}, ["prompt"]) assert cb.current_model == "phi-3" def test_model_from_serialized_name_field(self): cb = _make_callback() cb.on_llm_start({"name": "FancyModelName"}, ["prompt"]) assert cb.current_model == "FancyModelName" def test_fallback_to_type_field(self): cb = _make_callback() cb.on_llm_start({"_type": "SomeGenericLLM"}, ["prompt"]) assert cb.current_model == "SomeGenericLLM" def test_fallback_to_unknown_when_no_info(self): cb = _make_callback() cb.on_llm_start({}, ["prompt"]) assert cb.current_model == "unknown" def test_preset_model_overrides_all_extraction(self): """preset_model bypasses invocation_params extraction entirely.""" cb = _make_callback() cb.preset_model = "preset-gpt" cb.preset_provider = "openai" cb.on_llm_start( {"kwargs": {"model": "should-be-ignored"}}, ["prompt"], invocation_params={"model": "also-ignored"}, ) assert cb.current_model == "preset-gpt" assert cb.current_provider == "openai" def test_invocation_params_model_takes_priority_over_serialized(self): """invocation_params.model wins over serialized.kwargs.model.""" cb = _make_callback() cb.on_llm_start( {"kwargs": {"model": "from-serialized"}}, ["prompt"], invocation_params={"model": "from-invocation"}, ) assert cb.current_model == "from-invocation" # --------------------------------------------------------------------------- # 3. test_on_llm_start_ollama_type_detection # --------------------------------------------------------------------------- class TestOnLlmStartOllamaTypeDetection: """ChatOllama in _type field triggers Ollama-specific model and provider detection.""" def test_ollama_model_extracted_from_serialized_kwargs(self): cb = _make_callback() cb.on_llm_start( {"_type": "ChatOllama", "kwargs": {"model": "llama3:8b"}}, ["p"] ) assert cb.current_model == "llama3:8b" assert cb.current_provider == "ollama" def test_ollama_fallback_to_literal_ollama_when_no_kwargs_model(self): """When ChatOllama _type but no model in kwargs, model becomes 'ollama'.""" cb = _make_callback() cb.on_llm_start({"_type": "ChatOllama", "kwargs": {}}, ["p"]) assert cb.current_model == "ollama" assert cb.current_provider == "ollama" def test_ollama_no_kwargs_key_at_all(self): """ChatOllama with no kwargs key falls back to 'ollama'.""" cb = _make_callback() cb.on_llm_start({"_type": "ChatOllama"}, ["p"]) assert cb.current_model == "ollama" assert cb.current_provider == "ollama" def test_openai_type_sets_openai_provider(self): cb = _make_callback() cb.on_llm_start({"_type": "ChatOpenAI"}, ["p"]) assert cb.current_provider == "openai" def test_anthropic_type_sets_anthropic_provider(self): cb = _make_callback() cb.on_llm_start({"_type": "ChatAnthropic"}, ["p"]) assert cb.current_provider == "anthropic" def test_unknown_type_string_sets_provider_unknown(self): cb = _make_callback() cb.on_llm_start({"_type": "ChatMysteryProvider"}, ["p"]) assert cb.current_provider == "unknown" def test_no_type_field_sets_provider_unknown(self): cb = _make_callback() cb.on_llm_start({}, ["p"]) assert cb.current_provider == "unknown" def test_provider_kwarg_used_when_no_type(self): """When _type is absent, provider kwarg is used.""" cb = _make_callback() cb.on_llm_start({}, ["p"], provider="custom-provider") assert cb.current_provider == "custom-provider" # --------------------------------------------------------------------------- # 4. test_on_llm_end_usage_metadata_from_generations # --------------------------------------------------------------------------- class TestOnLlmEndUsageMetadataFromGenerations: """Ollama-specific path: usage_metadata on generation.message.""" def test_usage_metadata_extracted_and_counts_updated(self): cb = _make_callback() _setup_model_counts(cb, "mistral", "ollama") usage_meta = { "input_tokens": 30, "output_tokens": 15, "total_tokens": 45, } gen = _make_generation(usage_metadata=usage_meta, response_metadata={}) result = _make_llm_result(generations=[[gen]]) with patch.object(cb, "_save_to_db") as mock_save: cb.on_llm_end(result) assert cb.counts["total_prompt_tokens"] == 30 assert cb.counts["total_completion_tokens"] == 15 assert cb.counts["total_tokens"] == 45 mock_save.assert_called_once_with(30, 15) def test_usage_metadata_none_falls_through_to_response_metadata(self): """When usage_metadata is None, should continue to check response_metadata.""" cb = _make_callback() _setup_model_counts(cb, "llama3", "ollama") resp_meta = {"prompt_eval_count": 20, "eval_count": 10} gen = _make_generation(usage_metadata=None, response_metadata=resp_meta) result = _make_llm_result(generations=[[gen]]) with patch.object(cb, "_save_to_db") as mock_save: cb.on_llm_end(result) # Should pick up from response_metadata path assert cb.counts["total_tokens"] == 30 mock_save.assert_called_once_with(20, 10) def test_usage_metadata_zero_values_still_applied(self): """usage_metadata with zero token counts should still be applied.""" cb = _make_callback() _setup_model_counts(cb, "model-x", "openai") usage_meta = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} gen = _make_generation(usage_metadata=usage_meta, response_metadata={}) result = _make_llm_result(generations=[[gen]]) with patch.object(cb, "_save_to_db") as mock_save: cb.on_llm_end(result) # token_usage dict is built but all zeros → _save_to_db called mock_save.assert_called_once_with(0, 0) def test_no_message_attribute_on_generation_no_crash(self): """Generation without message attribute should not crash on_llm_end.""" cb = _make_callback() _setup_model_counts(cb, "model-y", "openai") gen = MagicMock(spec=[]) # no 'message' attribute result = _make_llm_result(generations=[[gen]]) # Must not raise with patch.object(cb, "_save_to_db"): cb.on_llm_end(result) def test_multiple_generation_lists_first_valid_used(self): """Only the first valid usage_metadata is consumed; loop breaks after.""" cb = _make_callback() _setup_model_counts(cb, "model-z", "openai") usage_meta_1 = { "input_tokens": 100, "output_tokens": 50, "total_tokens": 150, } usage_meta_2 = { "input_tokens": 999, "output_tokens": 999, "total_tokens": 1998, } gen1 = _make_generation( usage_metadata=usage_meta_1, response_metadata={} ) gen2 = _make_generation( usage_metadata=usage_meta_2, response_metadata={} ) result = _make_llm_result(generations=[[gen1], [gen2]]) with patch.object(cb, "_save_to_db"): cb.on_llm_end(result) # Only the first generation's tokens should be added assert cb.counts["total_tokens"] == 150 # --------------------------------------------------------------------------- # 5. test_on_llm_end_context_overflow_detection # --------------------------------------------------------------------------- class TestOnLlmEndContextOverflowDetection: """prompt_eval_count >= 80% of context_limit sets context_truncated = True.""" def _run_with_resp_meta(self, cb, resp_meta): gen = _make_generation(usage_metadata=None, response_metadata=resp_meta) result = _make_llm_result(generations=[[gen]]) with patch.object(cb, "_save_to_db"): cb.on_llm_end(result) def test_exactly_95_percent_triggers_overflow(self): cb = _make_callback() cb.context_limit = 1000 cb.original_prompt_estimate = 1100 _setup_model_counts(cb, "llama", "ollama") # 950 / 1000 = 95% — boundary triggers truncation self._run_with_resp_meta( cb, {"prompt_eval_count": 950, "eval_count": 50} ) assert cb.context_truncated is True assert cb.tokens_truncated > 0 assert 0.0 < cb.truncation_ratio <= 1.0 def test_above_95_percent_triggers_overflow(self): cb = _make_callback() cb.context_limit = 2000 cb.original_prompt_estimate = 2200 _setup_model_counts(cb, "phi3", "ollama") # 1950 / 2000 = 97.5% > 95% self._run_with_resp_meta( cb, {"prompt_eval_count": 1950, "eval_count": 50} ) assert cb.context_truncated is True def test_below_threshold_does_not_trigger_overflow(self): cb = _make_callback() cb.context_limit = 1000 cb.original_prompt_estimate = 700 _setup_model_counts(cb, "phi3", "ollama") # 700 / 1000 = 70% — below 80% threshold self._run_with_resp_meta( cb, {"prompt_eval_count": 700, "eval_count": 50} ) assert cb.context_truncated is False def test_no_context_limit_never_triggers_overflow(self): cb = _make_callback() cb.context_limit = None # not set _setup_model_counts(cb, "model", "ollama") self._run_with_resp_meta( cb, {"prompt_eval_count": 999999, "eval_count": 1} ) assert cb.context_truncated is False def test_tokens_truncated_zero_when_estimate_lte_actual(self): """When original estimate <= actual prompt_eval_count, tokens_truncated stays 0.""" cb = _make_callback() cb.context_limit = 1000 cb.original_prompt_estimate = ( 800 # less than actual → no truncation counted ) _setup_model_counts(cb, "llama", "ollama") self._run_with_resp_meta( cb, {"prompt_eval_count": 960, "eval_count": 40} ) # context_truncated may be True but tokens_truncated should be 0 # because original_prompt_estimate < prompt_eval_count assert cb.tokens_truncated == 0 def test_truncation_ratio_computed_correctly(self): cb = _make_callback() cb.context_limit = 1000 cb.original_prompt_estimate = 1200 _setup_model_counts(cb, "llama", "ollama") self._run_with_resp_meta( cb, {"prompt_eval_count": 960, "eval_count": 40} ) expected_truncated = 1200 - 960 # = 240 expected_ratio = expected_truncated / 1200 assert cb.tokens_truncated == expected_truncated assert abs(cb.truncation_ratio - expected_ratio) < 1e-9 # --------------------------------------------------------------------------- # 6. test_on_llm_end_ollama_response_metadata # --------------------------------------------------------------------------- class TestOnLlmEndOllamaResponseMetadata: """Raw Ollama metrics are captured from response_metadata.""" def test_all_ollama_fields_captured(self): cb = _make_callback() _setup_model_counts(cb, "mistral", "ollama") resp_meta = { "prompt_eval_count": 100, "eval_count": 40, "total_duration": 123456789, "load_duration": 111111, "prompt_eval_duration": 222222, "eval_duration": 333333, } gen = _make_generation(usage_metadata=None, response_metadata=resp_meta) result = _make_llm_result(generations=[[gen]]) with patch.object(cb, "_save_to_db"): cb.on_llm_end(result) assert cb.ollama_metrics["prompt_eval_count"] == 100 assert cb.ollama_metrics["eval_count"] == 40 assert cb.ollama_metrics["total_duration"] == 123456789 assert cb.ollama_metrics["load_duration"] == 111111 assert cb.ollama_metrics["prompt_eval_duration"] == 222222 assert cb.ollama_metrics["eval_duration"] == 333333 def test_missing_optional_fields_default_to_none(self): """Fields not present in response_metadata are stored as None in ollama_metrics.""" cb = _make_callback() _setup_model_counts(cb, "mistral", "ollama") # Only the mandatory trigger fields resp_meta = {"prompt_eval_count": 50, "eval_count": 20} gen = _make_generation(usage_metadata=None, response_metadata=resp_meta) result = _make_llm_result(generations=[[gen]]) with patch.object(cb, "_save_to_db"): cb.on_llm_end(result) assert cb.ollama_metrics["total_duration"] is None assert cb.ollama_metrics["load_duration"] is None assert cb.ollama_metrics["prompt_eval_duration"] is None assert cb.ollama_metrics["eval_duration"] is None def test_token_usage_built_from_response_metadata(self): """Token usage dict is constructed from prompt_eval_count + eval_count.""" cb = _make_callback() _setup_model_counts(cb, "qwen", "ollama") resp_meta = {"prompt_eval_count": 75, "eval_count": 25} gen = _make_generation(usage_metadata=None, response_metadata=resp_meta) result = _make_llm_result(generations=[[gen]]) with patch.object(cb, "_save_to_db") as mock_save: cb.on_llm_end(result) mock_save.assert_called_once_with(75, 25) assert cb.counts["total_tokens"] == 100 # 75 + 25 def test_only_eval_count_present_triggers_branch(self): """If only eval_count is set (not prompt_eval_count), the branch still fires because the condition is: prompt_eval_count OR eval_count.""" cb = _make_callback() _setup_model_counts(cb, "model", "ollama") resp_meta = {"eval_count": 30} gen = _make_generation(usage_metadata=None, response_metadata=resp_meta) result = _make_llm_result(generations=[[gen]]) with patch.object(cb, "_save_to_db") as mock_save: cb.on_llm_end(result) # ollama_metrics IS filled: eval_count captured, prompt_eval_count is None assert cb.ollama_metrics["eval_count"] == 30 assert cb.ollama_metrics["prompt_eval_count"] is None # token_usage is built: prompt=0, completion=30, total=30 mock_save.assert_called_once_with(0, 30) assert cb.counts["total_tokens"] == 30 def test_response_time_calculated_in_on_llm_end(self): """response_time_ms is set from start_time when on_llm_end is called.""" cb = _make_callback() cb.start_time = time.time() - 0.3 # 300 ms ago _setup_model_counts(cb, "model", "ollama") result = _make_llm_result(llm_output=None, generations=[]) cb.on_llm_end(result) assert cb.response_time_ms is not None assert cb.response_time_ms >= 200 # at least 200ms # --------------------------------------------------------------------------- # 7. test_on_llm_error_tracking # --------------------------------------------------------------------------- class TestOnLlmErrorTracking: """on_llm_error sets error status, error_type, response_time, and saves to db.""" def test_success_status_becomes_error(self): cb = _make_callback() assert cb.success_status == "success" with patch.object(cb, "_save_to_db"): cb.on_llm_error(ValueError("bad value")) assert cb.success_status == "error" def test_error_type_set_to_class_name(self): cb = _make_callback() with patch.object(cb, "_save_to_db"): cb.on_llm_error(TimeoutError("timed out")) assert cb.error_type == "TimeoutError" def test_error_type_for_custom_exception(self): class MyCustomError(Exception): pass cb = _make_callback() with patch.object(cb, "_save_to_db"): cb.on_llm_error(MyCustomError("oops")) assert cb.error_type == "MyCustomError" def test_save_to_db_called_with_zero_tokens(self): cb = _make_callback() with patch.object(cb, "_save_to_db") as mock_save: cb.on_llm_error(RuntimeError("crash")) mock_save.assert_called_once_with(0, 0) def test_response_time_calculated_when_start_time_set(self): cb = _make_callback() cb.start_time = time.time() - 1.0 # 1 second ago with patch.object(cb, "_save_to_db"): cb.on_llm_error(RuntimeError("late failure")) assert cb.response_time_ms is not None assert cb.response_time_ms >= 900 def test_no_db_save_when_no_research_id(self): """If research_id is None, _save_to_db should not be called.""" cb = TokenCountingCallback(research_id=None) with patch.object(cb, "_save_to_db") as mock_save: cb.on_llm_error(RuntimeError("boom")) mock_save.assert_not_called() def test_response_time_none_when_no_start_time(self): """If start_time was never set, response_time_ms stays None.""" cb = _make_callback() cb.start_time = None with patch.object(cb, "_save_to_db"): cb.on_llm_error(Exception("x")) assert cb.response_time_ms is None # --------------------------------------------------------------------------- # 8. test_save_to_db_no_username # --------------------------------------------------------------------------- class TestSaveToDbNoUsername: """In background thread: no username in research_context → warning, early return.""" def test_no_username_logs_warning_and_returns(self): cb = _make_callback(research_context={}) # empty context → no username with _patch_worker_thread(): with patch(f"{_MOD}.logger.warning") as mock_warn: with patch( "local_deep_research.database.thread_metrics.metrics_writer" ) as mock_writer: cb._save_to_db(50, 25) # Warning should have been logged mock_warn.assert_called_once() warning_msg = mock_warn.call_args[0][0] assert "username" in warning_msg.lower() or "no username" in warning_msg # No write attempted mock_writer.write_token_metrics.assert_not_called() def test_none_username_logs_warning(self): cb = _make_callback(research_context={"username": None}) with _patch_worker_thread(): with patch(f"{_MOD}.logger.warning") as mock_warn: cb._save_to_db(10, 5) mock_warn.assert_called_once() # --------------------------------------------------------------------------- # 9. test_save_to_db_no_password # --------------------------------------------------------------------------- class TestSaveToDbNoPassword: """In background thread: username present but no password → warning, early return.""" def test_no_password_logs_warning_and_returns(self): cb = _make_callback(research_context={"username": "alice"}) # no user_password key with _patch_worker_thread(): with patch(f"{_MOD}.logger.warning") as mock_warn: with patch( "local_deep_research.database.thread_metrics.metrics_writer" ) as mock_writer: cb._save_to_db(50, 25) mock_warn.assert_called_once() warning_msg = mock_warn.call_args[0][0] assert "password" in warning_msg.lower() mock_writer.write_token_metrics.assert_not_called() def test_none_password_logs_warning(self): cb = _make_callback( research_context={"username": "alice", "user_password": None} ) with _patch_worker_thread(): with patch(f"{_MOD}.logger.warning") as mock_warn: cb._save_to_db(10, 5) mock_warn.assert_called_once() # --------------------------------------------------------------------------- # 10. test_save_to_db_success # --------------------------------------------------------------------------- class TestSaveToDbSuccess: """Background thread success path: metrics_writer.write_token_metrics is called.""" def _make_ctx(self, **extra): ctx = { "username": "alice", "user_password": "secret", "research_query": "AI safety", "research_mode": "quick", "research_phase": "search", "search_iteration": 1, } ctx.update(extra) return ctx def test_write_token_metrics_called_with_correct_research_id(self): cb = _make_callback( research_context=self._make_ctx(), research_id="res-42" ) cb.current_model = "gpt-4" cb.current_provider = "openai" mock_writer = MagicMock() with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer", mock_writer, ): cb._save_to_db(100, 50) mock_writer.write_token_metrics.assert_called_once() args = mock_writer.write_token_metrics.call_args[0] assert args[0] == "alice" assert args[1] == "res-42" def test_token_data_contains_model_and_provider(self): cb = _make_callback(research_context=self._make_ctx()) cb.current_model = "claude-3" cb.current_provider = "anthropic" mock_writer = MagicMock() with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer", mock_writer, ): cb._save_to_db(200, 100) token_data = mock_writer.write_token_metrics.call_args[0][2] assert token_data["model_name"] == "claude-3" assert token_data["provider"] == "anthropic" assert token_data["prompt_tokens"] == 200 assert token_data["completion_tokens"] == 100 def test_set_user_password_called_before_write(self): """metrics_writer.set_user_password is invoked before write_token_metrics.""" cb = _make_callback(research_context=self._make_ctx()) cb.current_model = "gpt-4" cb.current_provider = "openai" mock_writer = MagicMock() call_order = [] mock_writer.set_user_password.side_effect = lambda *a: ( call_order.append("set_password") ) mock_writer.write_token_metrics.side_effect = lambda *a: ( call_order.append("write") ) with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer", mock_writer, ): cb._save_to_db(10, 5) assert call_order == ["set_password", "write"] def test_search_engines_planned_list_converted_to_json(self): """List values for search_engines_planned are JSON-serialised before write.""" import json ctx = self._make_ctx(search_engines_planned=["google", "brave"]) cb = _make_callback(research_context=ctx) cb.current_model = "gpt-4" cb.current_provider = "openai" mock_writer = MagicMock() with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer", mock_writer, ): cb._save_to_db(50, 25) token_data = mock_writer.write_token_metrics.call_args[0][2] assert isinstance(token_data["search_engines_planned"], str) assert json.loads(token_data["search_engines_planned"]) == [ "google", "brave", ] def test_context_overflow_fields_included_in_token_data(self): cb = _make_callback(research_context=self._make_ctx()) cb.current_model = "gpt-4" cb.current_provider = "openai" cb.context_limit = 4096 cb.context_truncated = True cb.tokens_truncated = 300 cb.truncation_ratio = 0.25 cb.ollama_metrics = {"prompt_eval_count": 3900} mock_writer = MagicMock() with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer", mock_writer, ): cb._save_to_db(100, 50) token_data = mock_writer.write_token_metrics.call_args[0][2] assert token_data["context_limit"] == 4096 assert token_data["context_truncated"] is True assert token_data["tokens_truncated"] == 300 assert token_data["ollama_prompt_eval_count"] == 3900 def test_exception_in_write_does_not_propagate(self): """Exception from metrics_writer.write_token_metrics is caught and logged.""" cb = _make_callback(research_context=self._make_ctx()) cb.current_model = "gpt-4" cb.current_provider = "openai" mock_writer = MagicMock() mock_writer.write_token_metrics.side_effect = RuntimeError("disk full") with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer", mock_writer, ): # Must not raise cb._save_to_db(10, 5) def test_calling_file_and_function_included(self): """call_stack tracking fields are forwarded to the token_data dict.""" cb = _make_callback(research_context=self._make_ctx()) cb.current_model = "gpt-4" cb.current_provider = "openai" cb.calling_file = "runner.py" cb.calling_function = "run_research" cb.call_stack = "runner.py:run_research:10" mock_writer = MagicMock() with _patch_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer", mock_writer, ): cb._save_to_db(30, 15) token_data = mock_writer.write_token_metrics.call_args[0][2] assert token_data["calling_file"] == "runner.py" assert token_data["calling_function"] == "run_research" assert token_data["call_stack"] == "runner.py:run_research:10"