Fix Gemini API content filter handling (#746)

## Summary
- avoid AttributeError when Gemini API returns `None` for chat message
- return empty output if message is filtered
- add regression test

## Testing
- `make format`
- `make lint`
- `make mypy`
- `make tests`

Towards #744
This commit is contained in:
Rohan Mehta 2025-05-23 13:00:30 -04:00 committed by GitHub
parent a96108e279
commit 6e078bf7a9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 51 additions and 5 deletions

View file

@ -71,12 +71,22 @@ class OpenAIChatCompletionsModel(Model):
stream=False,
)
first_choice = response.choices[0]
message = first_choice.message
if _debug.DONT_LOG_MODEL_DATA:
logger.debug("Received model response")
else:
logger.debug(
f"LLM resp:\n{json.dumps(response.choices[0].message.model_dump(), indent=2)}\n"
)
if message is not None:
logger.debug(
"LLM resp:\n%s\n",
json.dumps(message.model_dump(), indent=2),
)
else:
logger.debug(
"LLM resp had no message. finish_reason: %s",
first_choice.finish_reason,
)
usage = (
Usage(
@ -101,13 +111,15 @@ class OpenAIChatCompletionsModel(Model):
else Usage()
)
if tracing.include_data():
span_generation.span_data.output = [response.choices[0].message.model_dump()]
span_generation.span_data.output = (
[message.model_dump()] if message is not None else []
)
span_generation.span_data.usage = {
"input_tokens": usage.input_tokens,
"output_tokens": usage.output_tokens,
}
items = Converter.message_to_output_items(response.choices[0].message)
items = Converter.message_to_output_items(message) if message is not None else []
return ModelResponse(
output=items,

View file

@ -191,6 +191,40 @@ async def test_get_response_with_tool_call(monkeypatch) -> None:
assert fn_call_item.arguments == "{'x':1}"
@pytest.mark.allow_call_model_methods
@pytest.mark.asyncio
async def test_get_response_with_no_message(monkeypatch) -> None:
"""If the model returns no message, get_response should return an empty output."""
msg = ChatCompletionMessage(role="assistant", content="ignored")
choice = Choice(index=0, finish_reason="content_filter", message=msg)
choice.message = None # type: ignore[assignment]
chat = ChatCompletion(
id="resp-id",
created=0,
model="fake",
object="chat.completion",
choices=[choice],
usage=None,
)
async def patched_fetch_response(self, *args, **kwargs):
return chat
monkeypatch.setattr(OpenAIChatCompletionsModel, "_fetch_response", patched_fetch_response)
model = OpenAIProvider(use_responses=False).get_model("gpt-4")
resp: ModelResponse = await model.get_response(
system_instructions=None,
input="",
model_settings=ModelSettings(),
tools=[],
output_schema=None,
handoffs=[],
tracing=ModelTracing.DISABLED,
previous_response_id=None,
)
assert resp.output == []
@pytest.mark.asyncio
async def test_fetch_response_non_stream(monkeypatch) -> None:
"""