diff --git a/src/skillspector/llm_analyzer_base.py b/src/skillspector/llm_analyzer_base.py index 28a9029b3..fcabe9c4b 100644 --- a/src/skillspector/llm_analyzer_base.py +++ b/src/skillspector/llm_analyzer_base.py @@ -40,10 +40,12 @@ from random import uniform from typing import Any, Literal, cast, get_args +import tiktoken from langchain_anthropic import ChatAnthropic from langchain_core.messages import BaseMessage from langchain_openai.chat_models.base import BaseChatOpenAI from pydantic import BaseModel, Field, ValidationError, field_validator +from tiktoken import registry as tiktoken_registry from skillspector.inference_usage import ( InferenceUsageRecord, @@ -177,6 +179,76 @@ def _require_structured_response(response: object) -> None: ) +REQUEST_FRAMING_TOKENS = 256 + + +def _input_token_counter_key(chat_model: object) -> tuple[object, ...]: + """Retain the selected counter for equivalent client configurations.""" + if not isinstance(chat_model, BaseChatOpenAI): + return ("conservative_utf8_bytes",) + custom = getattr(chat_model, "custom_get_token_ids", None) + if callable(custom): + # The selected counter retains this callable, so its identity cannot be + # recycled while this key is cached. Do not retain whole HTTP clients. + return ("configured_chat_model_tokenizer", id(custom)) + label = getattr(chat_model, "tiktoken_model_name", None) or getattr( + chat_model, "model_name", None + ) + return ("openai_codec", label if isinstance(label, str) else None) + + +def _input_token_counter(chat_model: object) -> tuple[Callable[[str], int], str]: + """Use a configured/recognized OpenAI codec; unknown codecs use UTF-8 bytes. + + A model label must not silently inherit another model's tokenizer. Byte + length is a conservative bound for text when no verified codec is known. + """ + + def byte_bound(text: str) -> int: + return len(text.encode("utf-8", errors="replace")) + + if not isinstance(chat_model, BaseChatOpenAI): + return byte_bound, "conservative_utf8_bytes" + custom = getattr(chat_model, "custom_get_token_ids", None) + if callable(custom): + + def configured(text: str) -> int: + try: + count = len(custom(text)) + return count if count > 0 or not text else byte_bound(text) + except Exception: + return byte_bound(text) + + return configured, "configured_chat_model_tokenizer" + label = getattr(chat_model, "tiktoken_model_name", None) or getattr( + chat_model, "model_name", None + ) + if not isinstance(label, str): + return byte_bound, "conservative_utf8_bytes" + try: + # Name lookup is a pure table lookup. Never initialize a codec here: + # get_encoding/encoding_for_model can download its vocabulary without a + # timeout, outside the shared scan deadline. Use only an existing codec. + encoding_name = tiktoken.encoding_name_for_model(label) + encoding = tiktoken_registry.ENCODINGS.get(encoding_name) + except Exception: + return byte_bound, "conservative_utf8_bytes" + if encoding is None: + return byte_bound, "conservative_utf8_bytes" + + def recognized(text: str) -> int: + try: + return len(encoding.encode(text, disallowed_special=())) + except Exception: + return byte_bound(text) + + return recognized, "recognized_tiktoken_codec" + + +class LLMInputBudgetError(RuntimeError): + """Retain an unsubmitted batch that cannot fit the serialized input budget.""" + + class LLMRuntimeLimitError(RuntimeError): """Signal that no shared scan time remains for an LLM operation.""" @@ -1074,6 +1146,7 @@ def chunk_file_by_lines( overlap_lines: int = CHUNK_OVERLAP_LINES, *, max_bytes: int | None = None, + fits_chunk: Callable[[str, int, int], bool] | None = None, ) -> list[tuple[str, int, int]]: """Split *content* into line-range chunks that each fit within *max_tokens*. @@ -1081,10 +1154,14 @@ def chunk_file_by_lines( are 1-indexed. Consecutive chunks share *overlap_lines* lines of context so findings near chunk boundaries still have surrounding code. + The token estimate includes numbered line prefixes and rounds cumulative + serialized size, rather than rounding each line independently. + When *max_bytes* is set, each chunk must also fit within that many UTF-8 bytes *after* :func:`number_lines` has prefixed it, which is what a CLI provider's input cap measures. A single line that alone exceeds the budget - still becomes its own chunk: content is never dropped or truncated. + still becomes its own chunk: content is never dropped or truncated. The + invocation guard records an explicit failure instead of submitting it. """ lines = content.splitlines(keepends=True) if not lines: @@ -1098,25 +1175,37 @@ def chunk_file_by_lines( start_idx = 0 while start_idx < len(lines): - token_count = 0 byte_count = 0 end_idx = start_idx while end_idx < len(lines): - line_tokens = estimate_tokens(lines[end_idx]) - line_bytes = ( - len(lines[end_idx].encode("utf-8", errors="replace")) + prefix_bytes - if max_bytes is not None - else 0 - ) - over_tokens = token_count + line_tokens > max_tokens + line_bytes = len(lines[end_idx].encode("utf-8", errors="replace")) + prefix_bytes + # Round the cumulative serialized size, not each short line. Prefixes + # and UTF-8 bytes conservatively bound the chars/4 token estimate. + over_tokens = ( + byte_count + line_bytes + CHARS_PER_TOKEN - 1 + ) // CHARS_PER_TOKEN > max_tokens over_bytes = max_bytes is not None and byte_count + line_bytes > max_bytes if (over_tokens or over_bytes) and end_idx > start_idx: break - token_count += line_tokens byte_count += line_bytes end_idx += 1 + if fits_chunk is not None and not fits_chunk( + "".join(lines[start_idx:end_idx]), start_idx + 1, end_idx + ): + low, high = start_idx + 1, end_idx + best = None + while low <= high: + middle = (low + high) // 2 + if fits_chunk("".join(lines[start_idx:middle]), start_idx + 1, middle): + best = middle + low = middle + 1 + else: + high = middle - 1 + # An unfit single line stays represented; the call guard refuses it. + end_idx = best if best is not None else start_idx + 1 + chunk_text = "".join(lines[start_idx:end_idx]) chunks.append((chunk_text, start_idx + 1, end_idx)) # 1-indexed @@ -1241,6 +1330,8 @@ def __init__( self._require_time_remaining() raise self._require_time_remaining() + self._token_counter, self._token_counter_source = _input_token_counter(self._llm) + self._call_token_counters = {_input_token_counter_key(self._llm): self._token_counter} # Native SDK retries cannot re-read a workflow-wide deadline between # attempts. A dynamic deadline therefore uses our explicit retry loop, # which checks and caps every retry/backoff against remaining time. @@ -1360,13 +1451,13 @@ def get_batches( ) -> list[Batch]: """Create one :class:`Batch` per file, splitting oversized files into chunks. - For CLI providers every batch must also keep the built prompt under the + Every batch accounts for the real prompt wrapper, findings and schema + as well as line prefixes. For CLI providers it must also stay under the provider's byte cap (see :func:`_cli_max_input_bytes`). The fixed part of the prompt is measured with :meth:`build_prompt` itself, called with *prompt_kwargs* (the same keyword arguments later passed to the run), so subclass prompt layouts are accounted for rather than estimated. """ - base_overhead = estimate_tokens(self.base_prompt) max_input_bytes = _cli_max_input_bytes() findings_by_file: dict[str, list[Finding]] = defaultdict(list) @@ -1376,27 +1467,46 @@ def get_batches( batches: list[Batch] = [] for path in file_paths: - content = file_cache.get(path) or "No content available for this file." + content = file_cache.get(path) + if content is None: + content = "No content available for this file." file_findings = findings_by_file.get(path, []) - extra = self._estimate_extra_overhead(file_findings) - content_budget = max(self._input_budget - base_overhead - extra, 1024) - last_line = max(len(content.splitlines()), 1) - fixed_bytes = 0 + # Only the irreducible wrapper belongs in the coarse budget. The + # range predicate measures findings that actually occur in a chunk. + fixed_bytes = self._fixed_prompt_bytes(path, [], last_line, **prompt_kwargs) + minimum = Batch(file_path=path, content="", findings=[]) + if not self._prompt_fits_budget(self.build_prompt(minimum, **prompt_kwargs)): + # No line split can shrink a fixed wrapper. Keep all content as + # one explicit unsubmitted failure instead of refining every line. + retained = Batch(file_path=path, content=content, findings=file_findings) + batches.append(retained) + logger.warning("Fixed LLM prompt cannot fit input budget; batch content retained") + if max_input_bytes is not None and fixed_bytes > max_input_bytes: + logger.warning( + "LLM prompt exceeds the provider input limit; batch content retained" + ) + continue + fixed_tokens = (fixed_bytes + CHARS_PER_TOKEN - 1) // CHARS_PER_TOKEN + content_budget = max(self._input_budget - fixed_tokens, 1) byte_budget: int | None = None if max_input_bytes is not None: - fixed_bytes = self._fixed_prompt_bytes( - path, file_findings, last_line, **prompt_kwargs - ) byte_budget = max_input_bytes - fixed_bytes - PROMPT_BYTE_SAFETY_MARGIN first_batch = len(batches) - content_tokens = estimate_tokens(content) + content_tokens = ( + _numbered_size_bound(content, last_line) + CHARS_PER_TOKEN - 1 + ) // CHARS_PER_TOKEN fits_bytes = byte_budget is None or ( _numbered_size_bound(content, last_line) <= byte_budget ) - if content_tokens <= content_budget and fits_bytes: + whole = Batch(file_path=path, content=content, findings=file_findings) + if ( + content_tokens <= content_budget + and fits_bytes + and self._prompt_fits_budget(self.build_prompt(whole, **prompt_kwargs)) + ): batches.append( Batch( file_path=path, @@ -1405,13 +1515,38 @@ def get_batches( ) ) else: - chunk_budget = max(int(content_budget), 1024) + chunk_budget = content_budget # A non-positive byte budget means the fixed prompt alone is over the # cap; splitting the content further cannot help, and every batch is # reported below and fails closed at the provider. chunk_bytes = byte_budget if byte_budget is not None and byte_budget > 0 else None + + def fits_chunk( + text: str, + start: int, + end: int, + *, + path: str = path, + file_findings: list[Finding] = file_findings, + ) -> bool: + return self._prompt_fits_budget( + self.build_prompt( + Batch( + file_path=path, + content=text, + start_line=start, + end_line=end, + findings=findings_in_range(file_findings, start, end), + ), + **prompt_kwargs, + ) + ) + for chunk_text, s_line, e_line in chunk_file_by_lines( - content, chunk_budget, max_bytes=chunk_bytes + content, + chunk_budget, + max_bytes=chunk_bytes, + fits_chunk=fits_chunk, ): chunk_findings = findings_in_range(file_findings, s_line, e_line) batches.append( @@ -1497,6 +1632,45 @@ def parse_response(self, response: object, batch: Batch) -> list[Finding]: # -- Run loop ----------------------------------------------------------- + def _serialized_input(self, prompt: str) -> str: + if self.response_schema is not None: + return prompt + structured_output_instructions(self.response_schema) + return prompt + + def _count_prompt_tokens(self, prompt: str, *, chat_model: object | None = None) -> int: + if chat_model is None: + counter = self._token_counter + else: + # Reuse a chosen fallback even if another request later populates + # the process-wide codec registry. Equivalent refreshed HTTP clients + # share a counter without keeping their connection pools alive. + counters = getattr(self, "_call_token_counters", None) + if counters is None: + counters = self._call_token_counters = {} + key = _input_token_counter_key(chat_model) + if key not in counters: + counters[key] = _input_token_counter(chat_model)[0] + counter = counters[key] + return counter(self._serialized_input(prompt)) + REQUEST_FRAMING_TOKENS + + def _prompt_fits_budget(self, prompt: str) -> bool: + try: + self._check_prompt_budget(prompt) + except LLMInputBudgetError: + return False + return True + + def _check_prompt_budget(self, prompt: str, *, chat_model: object | None = None) -> None: + """Reject an unfit serialized batch before it reaches a provider.""" + size = len(self._serialized_input(prompt).encode("utf-8", errors="replace")) + count = self._count_prompt_tokens(prompt, chat_model=chat_model) + byte_cap = _cli_max_input_bytes() + if count > self._input_budget or (byte_cap is not None and size > byte_cap): + raise LLMInputBudgetError( + "serialized LLM input exceeds its token or byte budget; " + "batch content retained but not submitted" + ) + def _record_failed_request(self, exc: BaseException, prompt: str, ordinal: int) -> None: # Diagnostics must never replace the original provider exception or report usage. try: @@ -1536,6 +1710,7 @@ def _invoke_batch(self, batch: Batch, prompt: str) -> tuple[Batch, list]: len(batch.findings), ) llm, structured_llm = self._model_for_call() + self._check_prompt_budget(prompt, chat_model=llm) if structured_llm: try: response = self._invoke_with_failure_diagnostic(structured_llm, prompt) @@ -1566,6 +1741,7 @@ def _formatting_call(self, source: str) -> tuple[object, str]: "formatting input exceeds model budget", issue=ResponseIssue.LIMIT ) llm, _ = self._model_for_call() + self._check_prompt_budget(prompt, chat_model=llm) return bind_structured_output( llm, schema, @@ -1874,6 +2050,9 @@ def _recover_batch(self, batch: Batch, prompt: str, source: str) -> _RecoveredBa return self._recovered_result(batch, items, False, exc) try: llm, _ = self._model_for_call() + # Recovery repeats source plus quoted prior answers; measure the + # complete prompt against the actual refreshed request client too. + self._check_prompt_budget(question.prompt, chat_model=llm) response = self._invoke_with_failure_diagnostic(llm, question.prompt) if (getattr(response, "response_metadata", {}) or {}).get( "completion_source" @@ -1901,6 +2080,9 @@ async def _arecover_batch(self, batch: Batch, prompt: str, source: str) -> _Reco return self._recovered_result(batch, items, False, exc) try: llm, _ = self._model_for_call() + # Keep the async raw-recovery path under the same serialized + # token/byte limits without dropping already recovered findings. + self._check_prompt_budget(question.prompt, chat_model=llm) response = await self._ainvoke_with_failure_diagnostic(llm, question.prompt) if (getattr(response, "response_metadata", {}) or {}).get( "completion_source" @@ -2066,6 +2248,7 @@ async def _ainvoke_batch(self, batch: Batch, prompt: str) -> tuple[Batch, list]: len(batch.findings), ) llm, structured_llm = self._model_for_call() + self._check_prompt_budget(prompt, chat_model=llm) if structured_llm: try: response = await self._ainvoke_with_failure_diagnostic(structured_llm, prompt) diff --git a/tests/nodes/analyzers/test_dependency_sources.py b/tests/nodes/analyzers/test_dependency_sources.py index 7d67440b0..46999ec3a 100644 --- a/tests/nodes/analyzers/test_dependency_sources.py +++ b/tests/nodes/analyzers/test_dependency_sources.py @@ -14,7 +14,6 @@ from skillspector.llm_analyzer_base import Batch from skillspector.models import Finding from skillspector.nodes.meta_analyzer import ( - PER_FILE_ANALYSIS_PROMPT, LLMMetaAnalyzer, _fallback_filtered, _passthrough_with_defaults, @@ -1234,14 +1233,13 @@ def test_cargo_url_credentials_are_redacted_for_supported_schemes(destination: s assert "packages.example.invalid" in serialized -def test_sc10_credentials_are_redacted_before_provider_prompt_construction() -> None: +def test_sc10_credentials_are_redacted_before_provider_prompt_construction(mock_llm) -> None: username = "provider-user-sentinel" password = "provider-password-sentinel" token = "provider-token-sentinel" content = f"registry=ssh://{username}:{password}@packages.example.invalid/index?token={token}\n" finding = _analyze({".npmrc": content})[0] - analyzer = LLMMetaAnalyzer.__new__(LLMMetaAnalyzer) - analyzer.base_prompt = PER_FILE_ANALYSIS_PROMPT + analyzer = LLMMetaAnalyzer(model="fixture") analyzer._input_budget = 100_000 batch = analyzer.get_batches([".npmrc"], {".npmrc": content}, [finding])[0] @@ -1251,6 +1249,8 @@ def test_sc10_credentials_are_redacted_before_provider_prompt_construction() -> assert secret not in batch.content assert secret not in prompt assert "packages.example.invalid" in prompt + mock_llm.invoke.assert_not_called() + mock_llm.ainvoke.assert_not_called() def test_hidden_source_finding_is_marked_local_only() -> None: diff --git a/tests/nodes/analyzers/test_semantic_security_discovery.py b/tests/nodes/analyzers/test_semantic_security_discovery.py index ad74002f8..ba8a6f0c3 100644 --- a/tests/nodes/analyzers/test_semantic_security_discovery.py +++ b/tests/nodes/analyzers/test_semantic_security_discovery.py @@ -314,7 +314,8 @@ def test_oversized_file_multiple_batches(self, base_state) -> None: with patch( "skillspector.llm_analyzer_base.get_max_input_tokens", - return_value=50, + # Leave room for the fixed prompt/schema; only source needs splitting. + return_value=8192, ): with patch.object( LLMAnalyzerBase, "arun_batches", new_callable=AsyncMock, return_value=[] diff --git a/tests/nodes/test_llm_analyzer_base.py b/tests/nodes/test_llm_analyzer_base.py index c7e761cc8..9ac6d8c78 100644 --- a/tests/nodes/test_llm_analyzer_base.py +++ b/tests/nodes/test_llm_analyzer_base.py @@ -1142,13 +1142,22 @@ def test_cli_provider_splits_multibyte_file_under_token_budget( assert len(analyzer.build_prompt(batch).encode("utf-8")) <= MAX_INPUT_BYTES @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) - def test_http_provider_batching_unchanged(self, http_provider: None) -> None: + def test_http_provider_serialized_prompts_fit_and_keep_all_lines( + self, http_provider: None + ) -> None: analyzer = self._analyzer() content = "".join("z" * 10 + "\n" for _ in range(150_000)) - budget = max(analyzer._input_budget - estimate_tokens("test"), 1024) - expected = [(s, e) for _t, s, e in chunk_file_by_lines(content, budget)] - batches = analyzer.get_batches(["big.js"], {"big.js": content}) - assert [(b.start_line, b.end_line) for b in batches] == expected + covered: set[int] = set() + for batch in analyzer.get_batches(["big.js"], {"big.js": content}): + prompt = analyzer.build_prompt(batch) + if analyzer.response_schema is not None: + from skillspector.llm_utils import structured_output_instructions + + prompt += structured_output_instructions(analyzer.response_schema) + assert (len(prompt.encode("utf-8")) + 3) // 4 <= analyzer._input_budget + assert batch.end_line is not None + covered.update(range(batch.start_line, batch.end_line + 1)) + assert covered == set(range(1, 150_001)) @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) def test_cli_payload_with_schema_instructions_fits_without_margin( @@ -3209,6 +3218,12 @@ def test_cli_batches_fit_with_large_metadata_and_findings( content = "".join("m" * 40 + "\n" for _ in range(30_000)) findings = [self._make_finding("big.js", line=i) for i in range(1, 30_000, 100)] metadata = "d" * 400_000 + # Exercise the CLI byte cap with an explicit known local counting fixture; + # unknown codecs instead use the conservative byte-bound refusal test. + import tiktoken + + codec = tiktoken.get_encoding("cl100k_base") + analyzer._token_counter = lambda text: len(codec.encode(text, disallowed_special=())) batches = analyzer.get_batches( ["big.js"], {"big.js": content}, findings, metadata_text=metadata ) diff --git a/tests/nodes/test_llm_prompt_budget.py b/tests/nodes/test_llm_prompt_budget.py new file mode 100644 index 000000000..31b186680 --- /dev/null +++ b/tests/nodes/test_llm_prompt_budget.py @@ -0,0 +1,439 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Serialized prompt budgets must include every line and the real prompt wrapper.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import skillspector.llm_analyzer_base as lab +from skillspector.models import Finding + + +@pytest.fixture +def analyzer(monkeypatch: pytest.MonkeyPatch) -> lab.LLMAnalyzerBase: + """Construct the real planner with an inert provider; no model credentials or calls.""" + monkeypatch.setattr(lab, "get_max_input_tokens", lambda model: 8192) + monkeypatch.setattr(lab, "get_chat_model", lambda **kwargs: MagicMock()) + monkeypatch.setattr(lab, "_uses_native_connection_retries", lambda *args, **kwargs: False) + monkeypatch.setattr(lab, "new_inference_usage_collector", lambda **kwargs: MagicMock()) + monkeypatch.setattr(lab.LLMAnalyzerBase, "_bind_structured_output", lambda *args: MagicMock()) + monkeypatch.setattr(lab, "_cli_max_input_bytes", lambda: None) + return lab.LLMAnalyzerBase("Review every supplied line.", "fixture") + + +def test_short_line_rounding_does_not_exceed_chunk_budget() -> None: + chunks = lab.chunk_file_by_lines("ab\nab\nab", 1) + assert all(lab.estimate_tokens(text) <= 1 for text, _, _ in chunks) + assert {line for _, start, end in chunks for line in range(start, end + 1)} == {1, 2, 3} + + +@pytest.mark.parametrize("line", ["ab\n", "ab\r\n", "漢字🙂\n"]) +@pytest.mark.parametrize("compact", [False, True]) +def test_real_serialized_prompt_fits_and_keeps_every_line( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch, line: str, compact: bool +) -> None: + monkeypatch.setenv("SKILLSPECTOR_COMPACT_PROMPTS", "1" if compact else "0") + content = line * 5000 + batches = analyzer.get_batches(["source.js"], {"source.js": content}) + source_lines = content.splitlines(keepends=True) + covered: set[int] = set() + for batch in batches: + prompt = analyzer.build_prompt(batch) + if analyzer.response_schema is not None: + prompt += lab.structured_output_instructions(analyzer.response_schema) + assert len(prompt.encode("utf-8")) + lab.REQUEST_FRAMING_TOKENS <= analyzer._input_budget + end = batch.end_line or len(source_lines) + assert batch.content == "".join(source_lines[batch.start_line - 1 : end]) + covered.update(range(batch.start_line, end + 1)) + assert covered == set(range(1, len(source_lines) + 1)) + assert len(batches) < len(source_lines) + + +def test_real_custom_findings_and_metadata_overhead_is_measured( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + original = analyzer.build_prompt + + def custom_prompt(batch: lab.Batch, **kwargs: object) -> str: + return ( + original(batch) + "\nMetadata: " + "m" * 1000 + "\n" + "detail " * len(batch.findings) + ) + + monkeypatch.setattr(analyzer, "build_prompt", custom_prompt) + finding = Finding( + rule_id="fixture", message="fixture", severity="LOW", file="source.js", start_line=1 + ) + batches = analyzer.get_batches(["source.js"], {"source.js": "line\n" * 4000}, [finding] * 100) + for batch in batches: + prompt = analyzer.build_prompt(batch) + lab.structured_output_instructions( + analyzer.response_schema + ) + assert len(prompt.encode("utf-8")) + lab.REQUEST_FRAMING_TOKENS <= analyzer._input_budget + + +def test_empty_file_stays_empty(analyzer: lab.LLMAnalyzerBase) -> None: + assert lab.chunk_file_by_lines("", 1) == [("", 1, 1)] + batches = analyzer.get_batches(["empty"], {"empty": ""}) + assert len(batches) == 1 + assert batches[0].content == "" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_unsplittable_line_is_an_explicit_failure_without_provider_call( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + provider = MagicMock(side_effect=AssertionError("provider must not be reached")) + monkeypatch.setattr(lab, "_invoke_with_usage", provider) + monkeypatch.setattr(lab, "_ainvoke_with_usage", provider) + content = "x" * 10000 + batches = analyzer.get_batches(["long"], {"long": content}) + assert len(batches) == 1 + assert batches[0].content == content + if asynchronous: + outcome = await analyzer.arun_batches_detailed(batches, max_concurrency=1) + else: + outcome = analyzer.run_batches_detailed(batches) + assert not outcome.successful + assert len(outcome.failures) == 1 + assert outcome.failures[0].error_class == "LLMInputBudgetError" + assert outcome.failures[0].reason is lab.LedgerReason.LLM_BATCH_FAILED + provider.assert_not_called() + + +def test_unknown_openai_codec_uses_conservative_bytes() -> None: + from langchain_openai import ChatOpenAI + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = None + model.model_name = "unknown-future-model" + counter, source = lab._input_token_counter(model) + assert source == "conservative_utf8_bytes" + assert counter("漢🙂") == len("漢🙂".encode()) + + +def test_recognized_and_explicitly_configured_codec_is_respected( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from langchain_openai import ChatOpenAI + from tiktoken import registry + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "custom-provider-deployment" + codec = MagicMock() + codec.encode.return_value = [1, 2, 3] + monkeypatch.setattr(registry, "ENCODINGS", {"o200k_base": codec}) + counter, source = lab._input_token_counter(model) + text = "All source text stays data: <|endoftext|>" + assert source == "recognized_tiktoken_codec" + assert counter(text) == 3 + codec.encode.assert_called_once_with(text, disallowed_special=()) + + +def test_explicit_custom_codec_is_respected_and_failure_is_conservative() -> None: + from langchain_openai import ChatOpenAI + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = lambda text: [1, 2, 3] + counter, source = lab._input_token_counter(model) + assert source == "configured_chat_model_tokenizer" + assert counter("source") == 3 + model.custom_get_token_ids = MagicMock(side_effect=ValueError("codec unavailable")) + counter, _ = lab._input_token_counter(model) + assert counter("漢🙂") == len("漢🙂".encode()) + + +def test_actual_request_client_is_rechecked_when_its_codec_differs( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + from langchain_openai import ChatOpenAI + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = None + model.model_name = "unknown-runtime-model" + monkeypatch.setattr(analyzer, "_model_for_call", lambda: (model, None)) + monkeypatch.setattr(analyzer, "_token_counter", lambda text: 1) + provider = MagicMock(side_effect=AssertionError("provider must not be reached")) + monkeypatch.setattr(lab, "_invoke_with_usage", provider) + batch = lab.Batch(file_path="source", content="x" * 10000) + with pytest.raises(lab.LLMInputBudgetError): + analyzer._invoke_batch(batch, analyzer.build_prompt(batch)) + provider.assert_not_called() + + +def test_unfit_fixed_wrapper_retains_one_batch_without_per_line_refinement( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + original = analyzer.build_prompt + rendered = MagicMock(side_effect=lambda batch, **kwargs: "metadata" * 10000 + original(batch)) + monkeypatch.setattr(analyzer, "build_prompt", rendered) + content = "source line\n" * 10000 + batches = analyzer.get_batches(["source"], {"source": content}) + assert len(batches) == 1 + assert batches[0].content == content + assert rendered.call_count <= 3 + with pytest.raises(lab.LLMInputBudgetError): + analyzer._check_prompt_budget(analyzer.build_prompt(batches[0])) + + +@pytest.mark.parametrize("failure", [OSError, ValueError, RuntimeError]) +def test_codec_name_lookup_failure_uses_conservative_bytes( + monkeypatch: pytest.MonkeyPatch, failure: type[Exception] +) -> None: + import tiktoken + from langchain_openai import ChatOpenAI + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "recognized-model" + loader = MagicMock(side_effect=failure("offline or invalid tokenizer vocabulary")) + monkeypatch.setattr(tiktoken, "encoding_name_for_model", loader) + counter, source = lab._input_token_counter(model) + assert source == "conservative_utf8_bytes" + assert counter("漢🙂") == len("漢🙂".encode()) + loader.assert_called_once_with("gpt-4o") + + +@pytest.mark.parametrize("cached", [False, True]) +def test_codec_selection_never_loads_or_downloads_vocabulary( + monkeypatch: pytest.MonkeyPatch, cached: bool +) -> None: + import requests + import tiktoken + from langchain_openai import ChatOpenAI + from tiktoken import registry + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "recognized-model" + codec = MagicMock() + codec.encode.return_value = [1, 2, 3] + monkeypatch.setattr(registry, "ENCODINGS", {"o200k_base": codec} if cached else {}) + loader = MagicMock(side_effect=AssertionError("Vocabulary loading/network is forbidden")) + for owner, name in [ + (tiktoken, "encoding_for_model"), + (tiktoken, "get_encoding"), + (requests, "get"), + ]: + monkeypatch.setattr(owner, name, loader) + counter, source = lab._input_token_counter(model) + text = "漢🙂 source" + assert counter(text) == (3 if cached else len(text.encode())) + assert source == ("recognized_tiktoken_codec" if cached else "conservative_utf8_bytes") + loader.assert_not_called() + + +def test_cached_only_constructor_cannot_consume_shared_deadline( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + import tiktoken + from langchain_openai import ChatOpenAI + from tiktoken import registry + + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "recognized-model" + monkeypatch.setattr(registry, "ENCODINGS", {}) + clock = [0.0] + + def unbounded_loader(_label): + clock[0] += 1000.0 + return MagicMock() + + loader = MagicMock(side_effect=unbounded_loader) + monkeypatch.setattr(tiktoken, "encoding_for_model", loader) + monkeypatch.setattr(lab, "get_chat_model", lambda **kwargs: model) + created = lab.LLMAnalyzerBase("Review all lines", "gpt-4o", timeout=lambda: 1.0 - clock[0]) + assert created._remaining_timeout() == 1.0 + assert created._token_counter_source == "conservative_utf8_bytes" + loader.assert_not_called() + + +def test_selected_fallback_is_reused_for_equivalent_request_clients( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + from langchain_openai import ChatOpenAI + from tiktoken import registry + + def client(): + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = None + model.tiktoken_model_name = "gpt-4o" + model.model_name = "recognized-model" + return model + + model = client() + monkeypatch.setattr(registry, "ENCODINGS", {}) + before = analyzer._count_prompt_tokens("漢🙂 source", chat_model=model) + codec = MagicMock() + codec.encode.return_value = [1] + registry.ENCODINGS["o200k_base"] = codec + assert analyzer._count_prompt_tokens("漢🙂 source", chat_model=model) == before + assert analyzer._count_prompt_tokens("漢🙂 source", chat_model=client()) == before + codec.encode.assert_not_called() + + +def test_large_aggregate_findings_do_not_create_one_request_per_line( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch +) -> None: + original_prompt = analyzer.build_prompt + + def prompt(batch: lab.Batch, **kwargs: object) -> str: + return ( + original_prompt(batch) + "\nFindings:\n" + "\n".join(f.message for f in batch.findings) + ) + + monkeypatch.setattr(analyzer, "build_prompt", prompt) + content = "original source line\n" * 5000 + findings = [ + Finding( + rule_id="fixture", + message="Retained finding detail " * 60, + severity="HIGH", + file="source", + start_line=line, + ) + for line in range(1, 5001, 80) + ] + assert sum(len(f.message) for f in findings) > 4 * analyzer._input_budget + batches = analyzer.get_batches(["source"], {"source": content}, findings) + assert len(batches) < 150, ( + "Range-specific findings fit; aggregate metadata must not force one-line requests" + ) + covered = set() + retained = set() + lines = content.splitlines(keepends=True) + for batch in batches: + end = batch.end_line or len(lines) + assert batch.content == "".join(lines[batch.start_line - 1 : end]) + covered.update(range(batch.start_line, end + 1)) + retained.update(f.finding_id for f in batch.findings) + analyzer._check_prompt_budget(analyzer.build_prompt(batch)) + assert covered == set(range(1, len(lines) + 1)) + assert retained == {f.finding_id for f in findings} + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("limit_source", ["serialized-overhead", "refreshed-client"]) +async def test_field_recovery_keeps_known_findings_when_serialized_request_is_refused( + analyzer: lab.LLMAnalyzerBase, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + limit_source: str, +) -> None: + """A completed answer can be retained without submitting an oversized recovery.""" + source = '{"findings":[{"rule_id":"SSD-1","severity":"HIGH","start_line":1}]}' + cause = lab.StructuredOutputParseError( + "response lacks a required field", + issue=lab.ResponseIssue.INVALID_SCHEMA, + recovery_text=source, + ) + invalid = lab._StructuredResponseValidationError(cause.issue) + invalid.__cause__ = cause + analyzer.recovery_rule_ids = ("SSD-1",) + if limit_source == "serialized-overhead": + # The initial prompt fits, while repeating it with recovery instructions + # exceeds the real byte/token bound despite the old chars/4 estimate. + prompt = "x" * (analyzer._input_budget - analyzer._count_prompt_tokens("") - 1) + assert analyzer._prompt_fits_budget(prompt) + else: + from langchain_openai import ChatOpenAI + + prompt = "Assess the original source." + monkeypatch.setattr(analyzer, "_token_counter", lambda text: 1) + model = MagicMock(spec=ChatOpenAI) + model.custom_get_token_ids = lambda text: range(analyzer._input_budget + 1) + model.tiktoken_model_name = None + model.model_name = "refreshed-recovery-client" + monkeypatch.setattr(analyzer, "_model_for_call", lambda: (model, None)) + monkeypatch.setattr(analyzer, "build_prompt", lambda batch, **kwargs: prompt) + rejected = AsyncMock(side_effect=invalid) if asynchronous else MagicMock(side_effect=invalid) + monkeypatch.setattr(analyzer, "_ainvoke_batch" if asynchronous else "_invoke_batch", rejected) + provider = AsyncMock() if asynchronous else MagicMock() + monkeypatch.setattr( + lab, "_ainvoke_with_usage" if asynchronous else "_invoke_with_usage", provider + ) + batch = lab.Batch(file_path="source.py", content="original source") + + outcome = ( + await analyzer.arun_batches_detailed([batch], max_concurrency=1) + if asynchronous + else analyzer.run_batches_detailed([batch]) + ) + + assert not outcome.successful and not outcome.warnings + assert len(outcome.failures) == 1 + assert outcome.failures[0].error_class == "LLMInputBudgetError" + assert outcome.failures[0].reason is lab.LedgerReason.LLM_BATCH_FAILED + assert len(outcome.recovered) == 1 + retained_batch, findings = outcome.recovered[0] + assert retained_batch is batch + assert len(findings) == 1 + assert (findings[0].rule_id, findings[0].severity) == ("SSD-1", "HIGH") + assert findings[0].evidence["assessment_complete"] is False + assert "llm-format-recovery" in findings[0].tags + provider.assert_not_called() + rejected.assert_called_once() + assert analyzer._request_ordinal == 0 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_in_budget_field_recovery_preserves_diagnostics_and_partial_warning( + analyzer: lab.LLMAnalyzerBase, monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + """Permitted raw recovery keeps the existing diagnostic and warning contracts.""" + from langchain_core.messages import AIMessage + + source = '{"findings":[{"rule_id":"SSD-1","severity":"HIGH","start_line":1}]}' + cause = lab.StructuredOutputParseError( + "response lacks a required field", + issue=lab.ResponseIssue.INVALID_SCHEMA, + recovery_text=source, + ) + invalid = lab._StructuredResponseValidationError(cause.issue) + invalid.__cause__ = cause + analyzer.recovery_rule_ids = ("SSD-1",) + monkeypatch.setattr(analyzer, "build_prompt", lambda batch, **kwargs: "Assess this source.") + rejected = AsyncMock(side_effect=invalid) if asynchronous else MagicMock(side_effect=invalid) + monkeypatch.setattr(analyzer, "_ainvoke_batch" if asynchronous else "_invoke_batch", rejected) + response = AIMessage(content="UNKNOWN", response_metadata={"finish_reason": "stop"}) + provider = ( + AsyncMock(return_value=response) if asynchronous else MagicMock(return_value=response) + ) + monkeypatch.setattr( + lab, "_ainvoke_with_usage" if asynchronous else "_invoke_with_usage", provider + ) + name = "_ainvoke_with_failure_diagnostic" if asynchronous else "_invoke_with_failure_diagnostic" + diagnostic = ( + AsyncMock(wraps=getattr(analyzer, name)) + if asynchronous + else MagicMock(wraps=getattr(analyzer, name)) + ) + monkeypatch.setattr(analyzer, name, diagnostic) + batch = lab.Batch(file_path="source.py", content="original source") + + outcome = ( + await analyzer.arun_batches_detailed([batch], max_concurrency=1) + if asynchronous + else analyzer.run_batches_detailed([batch]) + ) + + assert not outcome.successful and not outcome.failures + assert len(outcome.warnings) == 1 + assert outcome.warnings[0].reason is lab.LedgerReason.LLM_FORMAT_RECOVERY_EXHAUSTED + assert len(outcome.recovered) == 1 + assert outcome.recovered[0][0] is batch + assert outcome.recovered[0][1][0].severity == "HIGH" + assert outcome.recovered[0][1][0].evidence["assessment_complete"] is False + assert provider.call_count > 0 + assert diagnostic.call_count == provider.call_count == analyzer._request_ordinal diff --git a/tests/unit/test_structured_response.py b/tests/unit/test_structured_response.py index b6589501f..8080177dc 100644 --- a/tests/unit/test_structured_response.py +++ b/tests/unit/test_structured_response.py @@ -1150,10 +1150,14 @@ def test_completed_ambiguous_response_remains_partial_even_when_recovery_says_ze def test_formatting_rechecks_the_input_allowance(http_model): - requests = http_model([('{"findings": [],}', "stop")]) + def reduce_input_allowance(_request_number): + analyzer._input_budget = 1 + + requests = http_model([('{"findings": [],}', "stop")], after_response=reduce_input_allowance) analyzer = make_analyzer("discovery") - analyzer._input_budget = 1 outcome = analyzer.run_batches_detailed([Batch(file_path="a.py", content="x")]) + assert not outcome.successful + assert len(outcome.failures) == 1 assert outcome.failures[0].reason == LedgerReason.LLM_RESPONSE_LIMIT assert len(requests) == 1