diff --git a/benchmark/run_benchmark.py b/benchmark/run_benchmark.py index 0c10171..f9a3c49 100644 --- a/benchmark/run_benchmark.py +++ b/benchmark/run_benchmark.py @@ -86,6 +86,30 @@ def parse_args() -> argparse.Namespace: default=5, help="Number of sample paragraphs to retain in the output preview.", ) + parser.add_argument( + "--record-artifacts", + action="store_true", + help="Record LLM extraction artifacts for debugging benchmark failures.", + ) + parser.add_argument( + "--artifact-dir", + help="Artifact output directory when --record-artifacts is enabled.", + ) + parser.add_argument( + "--record-raw-response-text", + action="store_true", + help="Include raw provider response text in recorded benchmark artifacts.", + ) + parser.add_argument( + "--record-paragraph-text", + action="store_true", + help="Include paragraph text in recorded benchmark artifacts.", + ) + parser.add_argument( + "--record-paragraph-metadata", + action="store_true", + help="Include paragraph metadata in recorded benchmark artifacts.", + ) return parser.parse_args() @@ -177,7 +201,12 @@ def build_config(args: argparse.Namespace) -> LabelGeneratorConfig: config.extraction.llm.max_output_tokens = args.max_output_tokens config.extraction.llm.timeout_seconds = args.timeout_seconds config.extraction.llm.max_concepts_per_paragraph = args.max_concepts_per_paragraph - config.extraction.llm.record_extraction_artifacts = False + config.extraction.llm.record_extraction_artifacts = args.record_artifacts + if args.artifact_dir is not None: + config.extraction.llm.artifact_dir = args.artifact_dir + config.extraction.llm.record_raw_response_text = args.record_raw_response_text + config.extraction.llm.record_paragraph_text = args.record_paragraph_text + config.extraction.llm.record_paragraph_metadata = args.record_paragraph_metadata return config diff --git a/src/labelgen/extraction/llm_extractor.py b/src/labelgen/extraction/llm_extractor.py index 0580e8d..e9de2c6 100644 --- a/src/labelgen/extraction/llm_extractor.py +++ b/src/labelgen/extraction/llm_extractor.py @@ -267,6 +267,9 @@ def _load_json_object(self, content: str) -> dict[str, Any]: else: start = content.find("{") if start == -1: + recovered = self._recover_single_paragraph_output(content) + if recovered is not None: + return recovered raise RuntimeError( "LLM extraction response did not contain valid JSON." ) from None @@ -275,7 +278,13 @@ def _load_json_object(self, content: str) -> dict[str, Any]: if recovered is not None: data = recovered else: - data, _ = decoder.raw_decode(content[start:]) + try: + data, _ = decoder.raw_decode(content[start:]) + except json.JSONDecodeError: + recovered = self._recover_single_paragraph_output(content) + if recovered is not None: + return recovered + raise if not isinstance(data, dict): raise RuntimeError("LLM extraction response must decode to a JSON object.") return cast(dict[str, Any], data) @@ -339,6 +348,99 @@ def _recover_partial_json_object( return None return cast(dict[str, Any], data) + def _recover_single_paragraph_output(self, content: str) -> dict[str, Any] | None: + """Recover one-paragraph concept output from lightly malformed local-model JSON.""" + + paragraphs_index = content.find('"paragraphs"') + if paragraphs_index == -1: + return None + array_start = content.find("[", paragraphs_index) + if array_start == -1: + return None + inner_start = content.find("[", array_start + 1) + if inner_start == -1: + array_end = content.find("]", array_start + 1) + if array_end == -1: + return None + return {"paragraphs": [[]]} + inner_end = self._find_matching_bracket(content, inner_start) + if inner_end is None: + return None + inner_payload = content[inner_start + 1 : inner_end] + concepts = self._extract_string_literals(inner_payload) + return {"paragraphs": [concepts]} + + def _find_matching_bracket(self, content: str, start_index: int) -> int | None: + """Find the matching closing bracket for one array literal.""" + + depth = 0 + in_string = False + escape = False + for index in range(start_index, len(content)): + char = content[index] + if in_string: + if escape: + escape = False + continue + if char == "\\": + escape = True + continue + if char == '"': + in_string = False + continue + if char == '"': + in_string = True + continue + if char == "[": + depth += 1 + continue + if char == "]": + depth -= 1 + if depth == 0: + return index + return None + + def _extract_string_literals(self, content: str) -> list[str]: + """Extract JSON-style string literals from a malformed list payload.""" + + literals: list[str] = [] + index = 0 + while index < len(content): + if content[index] != '"': + index += 1 + continue + cursor = index + 1 + escape = False + while cursor < len(content): + char = content[cursor] + if escape: + escape = False + cursor += 1 + continue + if char == "\\": + escape = True + cursor += 1 + continue + if char == '"': + literal = content[index : cursor + 1] + try: + decoded = json.loads(literal) + except json.JSONDecodeError: + # Skip past the closing quote we just found so malformed + # literals cannot trap the outer scanner on the same token. + index = cursor + 1 + break + if isinstance(decoded, str): + literals.append(decoded) + index = cursor + 1 + break + cursor += 1 + else: + break + if cursor >= len(content): + break + return literals + def _build_mentions(self, paragraph: Paragraph, concepts: list[str]) -> list[ConceptMention]: """Convert parsed concept text into mention models.""" diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 437b679..4ca77af 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -104,6 +104,11 @@ def test_summarize_run_reports_basic_benchmark_fields() -> None: max_output_tokens=512, timeout_seconds=30.0, max_concepts_per_paragraph=12, + record_artifacts=False, + artifact_dir=None, + record_raw_response_text=False, + record_paragraph_text=False, + record_paragraph_metadata=False, ) summary = run_benchmark.summarize_run( args=args, @@ -141,6 +146,11 @@ def test_summarize_run_builds_preview_from_cleaned_result_paragraphs() -> None: max_output_tokens=512, timeout_seconds=30.0, max_concepts_per_paragraph=12, + record_artifacts=False, + artifact_dir=None, + record_raw_response_text=False, + record_paragraph_text=False, + record_paragraph_metadata=False, ) summary = run_benchmark.summarize_run( args=args, @@ -203,6 +213,11 @@ def test_build_config_uses_conservative_default_batch_size_for_ollama() -> None: max_output_tokens=512, timeout_seconds=30.0, max_concepts_per_paragraph=12, + record_artifacts=False, + artifact_dir=None, + record_raw_response_text=False, + record_paragraph_text=False, + record_paragraph_metadata=False, ) config = run_benchmark.build_config(args) @@ -221,8 +236,40 @@ def test_build_config_keeps_default_cloud_batch_size() -> None: max_output_tokens=512, timeout_seconds=30.0, max_concepts_per_paragraph=12, + record_artifacts=False, + artifact_dir=None, + record_raw_response_text=False, + record_paragraph_text=False, + record_paragraph_metadata=False, ) config = run_benchmark.build_config(args) assert config.extraction.llm.batch_size == 8 + + +def test_build_config_can_enable_benchmark_artifacts() -> None: + args = argparse.Namespace( + extractor="llm", + provider="ollama", + model="qwen3.5:4b", + output_contract_mode="auto", + sample_preview=2, + batch_size=None, + max_output_tokens=512, + timeout_seconds=30.0, + max_concepts_per_paragraph=12, + record_artifacts=True, + artifact_dir="experiment/artifacts/benchmark-ollama", + record_raw_response_text=True, + record_paragraph_text=True, + record_paragraph_metadata=True, + ) + + config = run_benchmark.build_config(args) + + assert config.extraction.llm.record_extraction_artifacts is True + assert config.extraction.llm.artifact_dir == "experiment/artifacts/benchmark-ollama" + assert config.extraction.llm.record_raw_response_text is True + assert config.extraction.llm.record_paragraph_text is True + assert config.extraction.llm.record_paragraph_metadata is True diff --git a/tests/test_llm_extraction.py b/tests/test_llm_extraction.py index 8097b11..9a9bee3 100644 --- a/tests/test_llm_extraction.py +++ b/tests/test_llm_extraction.py @@ -203,6 +203,42 @@ def test_llm_extractor_recovers_missing_closing_delimiters(tmp_path: Path) -> No assert mentions == [] +def test_llm_extractor_recovers_single_paragraph_strings_without_commas( + tmp_path: Path, +) -> None: + config = LabelGeneratorConfig(extractor_mode="llm") + config.extraction.llm.model = "test-model" + config.extraction.llm.cache_dir = str(tmp_path) + client = FakeLLMProviderClient('{"paragraphs": [["OpenAI platform" "developer tooling"]]}') + extractor = LLMConceptExtractor(config.extraction, client=client) + + mentions = extractor.extract([Paragraph(id="p1", text="OpenAI builds developer tooling.")]) + + assert [mention.normalized for mention in mentions] == [ + "openai platform", + "developer tooling", + ] + + +def test_llm_extractor_skips_invalid_string_literals_without_hanging( + tmp_path: Path, +) -> None: + config = LabelGeneratorConfig(extractor_mode="llm") + config.extraction.llm.model = "test-model" + config.extraction.llm.cache_dir = str(tmp_path) + client = FakeLLMProviderClient( + '{"paragraphs": [["bad\\qescape" "OpenAI platform" "developer tooling"]]}' + ) + extractor = LLMConceptExtractor(config.extraction, client=client) + + mentions = extractor.extract([Paragraph(id="p1", text="OpenAI builds developer tooling.")]) + + assert [mention.normalized for mention in mentions] == [ + "openai platform", + "developer tooling", + ] + + def test_llm_extractor_can_record_structured_artifacts(tmp_path: Path) -> None: config = LabelGeneratorConfig(extractor_mode="llm") config.extraction.llm.model = "test-model"