wehub-resource-sync
7a0da7932b
OSV-Scanner (Scheduled) / scan-scheduled (push) Failing after 0s
Create Release / test-gate (push) Has been cancelled
Create Release / release-gate (push) Has been cancelled
Create Release / ci-gate (push) Has been cancelled
Create Release / version-check (push) Has been cancelled
Create Release / e2e-test-gate (push) Has been cancelled
Create Release / responsive-test-gate (push) Has been cancelled
Create Release / compat-test-gate (push) Has been cancelled
Create Release / compose-integration-gate (push) Has been cancelled
Create Release / vulture-gate (push) Has been cancelled
Create Release / build (push) Has been cancelled
Create Release / provenance (push) Has been cancelled
Create Release / prerelease-docker (push) Has been cancelled
Create Release / publish-docker (push) Has been cancelled
Create Release / create-release (push) Has been cancelled
Create Release / cleanup-changelog (push) Has been cancelled
Create Release / trigger-pypi (push) Has been cancelled
Create Release / monitor-pypi (push) Has been cancelled
Create Release / Clean up orphan prerelease tags and signatures (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [research-form] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [research-metrics] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [research-workflow] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [settings-core] (push) Has been cancelled
CodeQL Advanced / Analyze (javascript-typescript) (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [history-news] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [library] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [link-analytics] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [chat-core] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [chat-lifecycle] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [error-benchmark] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [settings-pages] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) (push) Has been cancelled
Docker Tests (Consolidated) / Accessibility Tests (push) Has been cancelled
Docker Tests (Consolidated) / LLM Unit Tests (push) Has been cancelled
Docker Tests (Consolidated) / LLM Example Tests (push) Has been cancelled
Docker Tests (Consolidated) / Production Image Smoke Test (push) Has been cancelled
Docker Tests (Consolidated) / Infrastructure Tests (push) Has been cancelled
OSSF Scorecard / OSSF Security Scorecard Analysis (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [mobile] (push) Has been cancelled
Backwards Compatibility / Verify Encryption Constants (push) Has been cancelled
Backwards Compatibility / PyPI Version Compatibility (push) Has been cancelled
Backwards Compatibility / Database Migration Tests (push) Has been cancelled
CodeQL Advanced / Analyze (python) (push) Has been cancelled
Docker Tests (Consolidated) / detect-changes (push) Has been cancelled
Docker Tests (Consolidated) / Build Test Image (push) Has been cancelled
Docker Tests (Consolidated) / All Pytest Tests + Coverage (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [accessibility] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [api-crud] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [auth-login] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [auth-pages] (push) Has been cancelled
Docker Tests (Consolidated) / UI Tests (Puppeteer) [auth-register] (push) Has been cancelled
983 lines
35 KiB
Python
983 lines
35 KiB
Python
"""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"
|