-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Isolate malformed LLM batch responses #265
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -33,7 +33,7 @@ | |
| from typing import Literal | ||
|
|
||
| from langchain_core.messages import BaseMessage | ||
| from pydantic import BaseModel, Field, field_validator | ||
| from pydantic import BaseModel, Field, ValidationError, field_validator | ||
|
|
||
| from skillspector.llm_utils import get_chat_model | ||
| from skillspector.logging_config import get_logger | ||
|
|
@@ -374,7 +374,9 @@ def run_batches( | |
|
|
||
| The element type of the inner list depends on the subclass: the default | ||
| :meth:`parse_response` returns :class:`Finding` objects; subclasses may | ||
| return dicts or other types. | ||
| return dicts or other types. Malformed structured responses are isolated | ||
| per batch and omitted from the result so callers can detect partial | ||
| analysis by comparing submitted and returned batches. | ||
| """ | ||
| results: list[tuple[Batch, list]] = [] | ||
| for batch in batches: | ||
|
|
@@ -386,11 +388,38 @@ def run_batches( | |
| len(batch.findings), | ||
| ) | ||
| if self._structured_llm: | ||
| response = self._structured_llm.invoke(prompt) | ||
| try: | ||
| response = self._structured_llm.invoke(prompt) | ||
| except ValidationError as exc: | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This carve-out is correct here, but the same fix is missing in |
||
| logger.warning("LLM batch failed for %s: %s", batch.file_label, exc) | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Dropping the batch with only a log warning gives callers no signal. Previously a |
||
| continue | ||
| except (ValueError, NotImplementedError): | ||
| raise | ||
| except Exception as exc: | ||
| logger.warning("LLM batch failed for %s: %s", batch.file_label, exc) | ||
| continue | ||
| else: | ||
| response = _message_text(self._llm.invoke(prompt)) | ||
| try: | ||
| response = _message_text(self._llm.invoke(prompt)) | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Inconsistency: this raw-mode branch has no |
||
| except ValidationError as exc: | ||
| logger.warning("LLM batch failed for %s: %s", batch.file_label, exc) | ||
| continue | ||
| except (ValueError, NotImplementedError): | ||
| raise | ||
| except Exception as exc: | ||
| logger.warning("LLM batch failed for %s: %s", batch.file_label, exc) | ||
| continue | ||
| logger.debug("LLM response for %s", batch.file_label) | ||
| parsed = self.parse_response(response, batch) | ||
| try: | ||
| parsed = self.parse_response(response, batch) | ||
| except ValidationError as exc: | ||
| logger.warning("LLM batch parse failed for %s: %s", batch.file_label, exc) | ||
| continue | ||
| except (ValueError, NotImplementedError): | ||
| raise | ||
| except Exception as exc: | ||
| logger.warning("LLM batch parse failed for %s: %s", batch.file_label, exc) | ||
| continue | ||
| results.append((batch, parsed)) | ||
| return results | ||
|
|
||
|
|
@@ -408,12 +437,12 @@ async def arun_batches( | |
| cross-chunk batches are parallelized in a single gather call. | ||
|
|
||
| Failures are isolated per batch: a transient error (timeout, 429, | ||
| oversized-chunk 400, ...) costs only its own batch, which is logged | ||
| and omitted from the result, so one bad call cannot cancel the rest | ||
| of the fan-out. Callers can detect partial results by comparing the | ||
| returned batches against the submitted ones. ``ValueError`` and | ||
| ``NotImplementedError`` signal misconfiguration rather than infra | ||
| trouble and keep propagating. | ||
| oversized-chunk 400, malformed structured response, ...) costs only | ||
| its own batch, which is logged and omitted from the result, so one bad | ||
| call cannot cancel the rest of the fan-out. Callers can detect partial | ||
| results by comparing the returned batches against the submitted ones. | ||
| ``ValueError`` and ``NotImplementedError`` signal misconfiguration | ||
| rather than infra trouble and keep propagating. | ||
|
|
||
| The return type mirrors :meth:`run_batches`. | ||
| """ | ||
|
|
@@ -438,6 +467,9 @@ async def _process(batch: Batch) -> tuple[Batch, list]: | |
| results = await asyncio.gather(*[_process(b) for b in batches], return_exceptions=True) | ||
| successful: list[tuple[Batch, list]] = [] | ||
| for batch, result in zip(batches, results, strict=True): | ||
| if isinstance(result, ValidationError): | ||
| logger.warning("LLM batch failed for %s: %s", batch.file_label, result) | ||
| continue | ||
| if isinstance(result, (ValueError, NotImplementedError)): | ||
| raise result | ||
| if isinstance(result, BaseException): | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -388,6 +388,50 @@ async def test_arun_batches_uses_message_text_for_content_blocks(self) -> None: | |
| assert results[0][1] == ["async chunk"] | ||
|
|
||
|
|
||
| # --------------------------------------------------------------------------- | ||
| # LLMAnalyzerBase.run_batches (sync execution) | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| class TestRunBatches: | ||
| MODEL = "nvidia/openai/gpt-oss-120b" | ||
|
|
||
| @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) | ||
| def test_malformed_structured_batch_does_not_abort_the_others(self) -> None: | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good regression test for the sync path. Please add the async counterpart: an |
||
| """A malformed structured response costs only its own batch.""" | ||
|
|
||
| def _invoke(prompt: str) -> LLMAnalysisResult: | ||
| if "b.py" in prompt: | ||
| return LLMAnalysisResult.model_validate({"findings": 'We{"findings":[]}'}) | ||
| return LLMAnalysisResult( | ||
| findings=[ | ||
| LLMFinding(rule_id="T-1", message="hit", severity="LOW", start_line=1), | ||
| ] | ||
| ) | ||
|
|
||
| analyzer = LLMAnalyzerBase(base_prompt="test", model=self.MODEL) | ||
| analyzer._structured_llm.invoke.side_effect = _invoke | ||
|
|
||
| batches = [ | ||
| Batch(file_path="a.py", content="code a"), | ||
| Batch(file_path="b.py", content="code b"), | ||
| Batch(file_path="c.py", content="code c"), | ||
| ] | ||
| results = analyzer.run_batches(batches) | ||
|
|
||
| assert {batch.file_path for batch, _ in results} == {"a.py", "c.py"} | ||
| assert [items[0].rule_id for _, items in results] == ["T-1", "T-1"] | ||
|
|
||
| @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) | ||
| def test_value_error_still_propagates(self) -> None: | ||
| """ValueError signals misconfiguration, not a malformed model response.""" | ||
| analyzer = LLMAnalyzerBase(base_prompt="test", model=self.MODEL) | ||
| analyzer._structured_llm.invoke.side_effect = ValueError("no API key") | ||
|
|
||
| with pytest.raises(ValueError, match="no API key"): | ||
| analyzer.run_batches([Batch(file_path="a.py", content="code")]) | ||
|
|
||
|
|
||
| # --------------------------------------------------------------------------- | ||
| # LLMAnalyzerBase.arun_batches (async parallel execution) | ||
| # --------------------------------------------------------------------------- | ||
|
|
@@ -577,6 +621,32 @@ async def _flaky_ainvoke(prompt: str) -> LLMAnalysisResult: | |
| results = await analyzer.arun_batches(batches) | ||
| assert {batch.file_path for batch, _ in results} == {"a.py", "c.py"} | ||
|
|
||
| @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) | ||
| async def test_malformed_structured_batch_does_not_abort_the_others(self) -> None: | ||
| """A malformed structured response is isolated even though it is a ValueError.""" | ||
|
|
||
| async def _ainvoke(prompt: str) -> LLMAnalysisResult: | ||
| if "b.py" in prompt: | ||
| return LLMAnalysisResult.model_validate({"findings": 'We{"findings":[]}'}) | ||
| return LLMAnalysisResult( | ||
| findings=[ | ||
| LLMFinding(rule_id="T-1", message="hit", severity="LOW", start_line=1), | ||
| ] | ||
| ) | ||
|
|
||
| analyzer = LLMAnalyzerBase(base_prompt="test", model=self.MODEL) | ||
| analyzer._structured_llm.ainvoke = _ainvoke | ||
|
|
||
| batches = [ | ||
| Batch(file_path="a.py", content="code a"), | ||
| Batch(file_path="b.py", content="code b"), | ||
| Batch(file_path="c.py", content="code c"), | ||
| ] | ||
| results = await analyzer.arun_batches(batches) | ||
|
|
||
| assert {batch.file_path for batch, _ in results} == {"a.py", "c.py"} | ||
| assert [items[0].rule_id for _, items in results] == ["T-1", "T-1"] | ||
|
|
||
| @patch(MOCK_PATCH_TARGET, _mock_get_chat_model) | ||
| async def test_all_batches_failed_returns_empty(self) -> None: | ||
| analyzer = LLMAnalyzerBase(base_prompt="test", model=self.MODEL) | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maintainability: this is the first of three near-identical try/except blocks in the loop. Extracting the policy into a small helper (skip on ValidationError/other Exceptions, propagate ValueError/NotImplementedError) would keep it in one place and let
arun_batchesshare it, preventing the sync/async drift that left the async path unfixed. Also, therun_batchesdocstring should document the new failure-isolation semantics the wayarun_batches's docstring does.