mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(ccr): guard empty/malformed OpenAI choices in _extract_assistant_message (#2389)
## Description
`CCRResponseHandler._extract_assistant_message` extracts the assistant
message from an upstream response while building the CCR
retrieval-continuation history. The OpenAI branch is not defensive about
an empty or malformed `choices` array:
```python
elif provider == "openai":
message = response.get("choices", [{}])[0].get("message", {})
```
`response.get("choices", [{}])` only falls back to `[{}]` when the key
is **absent**. When `choices` is present but empty (`[]`) or carries a
null first element (`[null]`), this raises on the success path:
- `choices: []` → `[][0]` → `IndexError`
- `choices: [null]` → `None.get(...)` → `AttributeError`
OpenAI-compatible gateways can return those shapes on content-filtered
or usage-only responses. The sibling **Google** branch a few lines below
already guards this (`candidates = response.get("candidates", []); if
candidates: ... else: parts = []`), and so does `ccr/tool_calls.py` (it
checks `isinstance(choices, list)`, non-empty, and
`isinstance(first_choice, dict)`). Only this OpenAI branch was missed.
## Fix
Guard the list and the first element the same way the siblings do:
```python
elif provider == "openai":
choices = response.get("choices")
first = choices[0] if isinstance(choices, list) and choices else {}
message = first.get("message", {}) if isinstance(first, dict) else {}
return {
"role": "assistant",
"content": message.get("content"),
"tool_calls": message.get("tool_calls"),
}
```
A well-formed response is unaffected; an empty/null/absent `choices` now
yields `{"role": "assistant", "content": None, "tool_calls": None}`
instead of raising.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- `headroom/ccr/response_handler.py`: guard empty/non-list `choices` and
a non-dict first element in the OpenAI branch of
`_extract_assistant_message`.
- `tests/test_ccr_response_handler.py`: add
`TestExtractAssistantMessageEdgeCases` (empty `choices`, `[null]`,
absent, and the normal case).
- `CHANGELOG.md`: Bug Fixes entry.
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/ccr/response_handler.py tests/test_ccr_response_handler.py
All checks passed!
$ uvx ruff@0.15.17 format --check headroom/ccr/response_handler.py tests/test_ccr_response_handler.py
2 files already formatted
# Verified against the REAL imported module (headroom.ccr.response_handler is
# light — no ML imports), so this ran locally in the project venv:
$ python -c "from headroom.ccr.response_handler import CCRResponseHandler as H; h=H(); \
assert h._extract_assistant_message({'choices': []}, 'openai') == {'role':'assistant','content':None,'tool_calls':None}"
# (no IndexError; normal case still extracts content/tool_calls)
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, project venv (`uv sync --extra
proxy`), `uvx ruff@0.15.17`.
- Exact command / steps: imported the real `CCRResponseHandler` and
called `_extract_assistant_message` with `{"choices": []}`, `{"choices":
[null]}`, `{}` (absent), and a normal `{"choices": [{"message":
{...}}]}`.
- Observed result: the OLD code raised `IndexError` on `[]` and
`AttributeError` on `[null]`; the NEW code returns `{"role":
"assistant", "content": None, "tool_calls": None}` for all three
malformed shapes and still extracts `content`/`tool_calls` from a
well-formed response. Because `response_handler` has no ML imports, this
ran against the actual module, not a replica.
- Not tested: a live CCR retrieval round trip through a gateway that
emits empty choices; the added unit tests drive
`_extract_assistant_message` directly.
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
## Checklist
- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
`headroom/ccr/response_handler.py` is a light module (no ML imports), so
unlike most of my recent PRs I verified the fix by importing the real
class in the project venv (output above), in addition to the added unit
tests. This aligns the OpenAI branch with the already-defensive Google
branch and `ccr/tool_calls.py`.
This commit is contained in:
parent
d6a1af40d5
commit
89319fbcad
2 changed files with 36 additions and 1 deletions
|
|
@ -379,7 +379,15 @@ class CCRResponseHandler:
|
|||
"content": response.get("content", []),
|
||||
}
|
||||
elif provider == "openai":
|
||||
message = response.get("choices", [{}])[0].get("message", {})
|
||||
# Guard an empty/malformed ``choices`` the same way the Google branch
|
||||
# below (and ccr/tool_calls.py) already do: ``response.get("choices",
|
||||
# [{}])`` only falls back when the key is absent, so a present-but-
|
||||
# empty ``choices: []`` (or ``[null]``) — which OpenAI-compatible
|
||||
# gateways can send on a content-filtered/usage-only response — made
|
||||
# ``[0]`` raise IndexError (or ``.get`` raise on a non-dict).
|
||||
choices = response.get("choices")
|
||||
first = choices[0] if isinstance(choices, list) and choices else {}
|
||||
message = first.get("message", {}) if isinstance(first, dict) else {}
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": message.get("content"),
|
||||
|
|
|
|||
|
|
@ -719,3 +719,30 @@ class TestExtractAssistantMessage:
|
|||
assert message["role"] == "assistant"
|
||||
assert message["content"] == "Hello"
|
||||
assert message["tool_calls"] == [{"id": "123"}]
|
||||
|
||||
|
||||
class TestExtractAssistantMessageEdgeCases:
|
||||
"""Regression: `_extract_assistant_message` must not crash on an empty or
|
||||
malformed OpenAI `choices` array (OpenAI-compatible gateways can send
|
||||
`choices: []` or `[null]` on content-filtered / usage-only responses)."""
|
||||
|
||||
def test_openai_empty_choices_does_not_crash(self):
|
||||
handler = CCRResponseHandler()
|
||||
msg = handler._extract_assistant_message({"choices": []}, "openai")
|
||||
assert msg == {"role": "assistant", "content": None, "tool_calls": None}
|
||||
|
||||
def test_openai_null_first_choice_does_not_crash(self):
|
||||
handler = CCRResponseHandler()
|
||||
msg = handler._extract_assistant_message({"choices": [None]}, "openai")
|
||||
assert msg == {"role": "assistant", "content": None, "tool_calls": None}
|
||||
|
||||
def test_openai_absent_choices_does_not_crash(self):
|
||||
handler = CCRResponseHandler()
|
||||
msg = handler._extract_assistant_message({}, "openai")
|
||||
assert msg == {"role": "assistant", "content": None, "tool_calls": None}
|
||||
|
||||
def test_openai_normal_choice_still_extracts(self):
|
||||
handler = CCRResponseHandler()
|
||||
resp = {"choices": [{"message": {"content": "hi", "tool_calls": [{"id": "1"}]}}]}
|
||||
msg = handler._extract_assistant_message(resp, "openai")
|
||||
assert msg == {"role": "assistant", "content": "hi", "tool_calls": [{"id": "1"}]}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue