Harden LLM benchmark stability for v0.2.2

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
 
 

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."""
 

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

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"