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:
parent
a96108e279
commit
6e078bf7a9
2 changed files with 51 additions and 5 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
Loading…
Reference in a new issue