Prevent crashes from malformed LLM responses - #32
Conversation
There was a problem hiding this comment.
The intent (fail loudly on malformed Ollama responses instead of hitting an opaque AttributeError/KeyError) is good, and the isinstance(message, dict) check in _stream_chat is a real improvement. Two issues in the current implementation cut against that intent, though.
1. agent/core/llm.py:205-211 — a chunk with tool_calls but no content key now silently drops the tool calls.
Before, chunk["message"].get("content", "") tolerated a missing content and still reached the if "tool_calls" in ... branch. Now "content" in message gates the entire block, so a streamed chunk shaped like {"message": {"role": "assistant", "tool_calls": [...]}} is skipped outright and full_response["message"]["tool_calls"] is never set — the caller gets an empty assistant turn with no error. That is exactly the silent-failure mode this PR is trying to remove, and it is a regression relative to main. Suggested shape:
message = chunk["message"]
if isinstance(message, dict):
content = message.get("content") or ""
if content:
full_response["message"]["content"] += content
yield content
if "tool_calls" in message:
full_response["message"]["tool_calls"] = message["tool_calls"]2. agent/core/llm.py:58-59 / 207 — presence of content is checked, but not its type.
_validate_message_structure accepts {"message": {"content": None}}, and in the streaming path full_response["message"]["content"] += content still raises TypeError when content is None or a list. Since the stated goal is preventing crashes on malformed messages, the check should be on type (isinstance(content, str)) rather than key presence — that is where the actual crash lives.
3. agent/core/llm.py:184 — hard-failing a tool-call-only response may be stricter than callers want.
engine.py:94-97 deliberately tolerates a present-but-null/missing content (.get("content", "") or "", with a comment saying so). With this change, a non-streaming response carrying tool_calls and no content key raises OllamaError and aborts the whole turn at engine.py:88-92, where previously it worked. Consider allowing a missing content when tool_calls is present, or normalizing to "", so the two layers agree.
Minor: _validate_response(result, ["message"]) at line 183 already covers the "message" not in response case re-checked at line 51; message.get("content", "") at 206 has a now-dead default; and there are a few trailing-whitespace lines in the new docstring (43-49). Also worth adding tests under tests/ for the new validator and for the tool-calls-without-content streaming chunk, since that is the case that regressed.
…ence
- _stream_chat no longer gates the tool_calls branch on a 'content' key,
so a chunk carrying only tool_calls is recorded again.
- Validate that 'content' is a str; tolerate missing/null content when
'tool_calls' is present, matching engine.py's `.get("content", "") or ""`.
- Drop the redundant _validate_response(result, ["message"]) call.
- Add tests/test_llm.py covering the validator and streaming chunks.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Resolved the review. Summary of what changed and what did not. Fixed — streamed tool calls were dropped ( Fixed — presence checked instead of type ( Fixed — tool-call-only responses no longer abort the turn ( Also applied from the minor notes: dropped the redundant Not changed: nothing was pushed back on — each of the three points held up against the code. Checks: |
There was a problem hiding this comment.
Reviewed the diff in agent/core/llm.py and the new tests/test_llm.py. The new _validate_message_structure correctly replaces the presence-only _validate_response(result, ["message"]) check in chat(): it rejects a missing/non-dict message and non-str content, while deliberately tolerating absent/null content on a tool-call-only turn, which is a legitimate Ollama response. The streaming path change is consistent — non-str content raises (OllamaError is not caught by the surrounding except json.JSONDecodeError), and tool_calls are now preserved on content-less chunks, fixing a real gap in the old code that only read tool_calls after appending content.
Checked downstream impact: engine.py:94-95 already defends with .get("message", {}).get("content", "") or "" and catches OllamaError in process_message, so the stricter validation degrades to a user-facing error rather than a crash. _stream_chat has no production callers today, so the one behavior change there (empty-string chunks are no longer yielded) is inert. No injection, path-traversal, or secret-leakage surface — error strings echo the local Ollama response, not credentials — and no failures are swallowed. Scope is tight: one module plus its tests, matching the PR title and description.
What
Add validation that LLM response message contains expected structure before accessing content field.
Why
Prevents silent IndexError or KeyError crashes when Ollama returns malformed messages lacking required fields.