From 334bd80e686e7c75f93b63be89225ad301239d2f Mon Sep 17 00:00:00 2001 From: Devin Lai Date: Thu, 18 Jun 2026 11:44:19 +0800 Subject: [PATCH] fix: tolerate partial LiteLLM citations --- src/agents/extensions/models/litellm_model.py | 32 +++-- tests/models/test_litellm_annotations.py | 124 ++++++++++++++++++ tests/test_run_step_execution.py | 3 +- 3 files changed, 147 insertions(+), 12 deletions(-) create mode 100644 tests/models/test_litellm_annotations.py diff --git a/src/agents/extensions/models/litellm_model.py b/src/agents/extensions/models/litellm_model.py index 4e649feb08..bad58c708c 100644 --- a/src/agents/extensions/models/litellm_model.py +++ b/src/agents/extensions/models/litellm_model.py @@ -852,18 +852,28 @@ def convert_annotations_to_openai( if not annotations: return None - return [ - Annotation( - type="url_citation", - url_citation=AnnotationURLCitation( - start_index=annotation["url_citation"]["start_index"], - end_index=annotation["url_citation"]["end_index"], - url=annotation["url_citation"]["url"], - title=annotation["url_citation"]["title"], - ), + results: list[Annotation] = [] + for annotation in annotations: + url_citation = annotation.get("url_citation") + if not url_citation: + continue + + url = url_citation.get("url") or "" + if not url: + continue + + results.append( + Annotation( + type="url_citation", + url_citation=AnnotationURLCitation( + start_index=url_citation.get("start_index") or 0, + end_index=url_citation.get("end_index") or 0, + url=url, + title=url_citation.get("title") or "", + ), + ) ) - for annotation in annotations - ] + return results or None @classmethod def convert_tool_call_to_openai( diff --git a/tests/models/test_litellm_annotations.py b/tests/models/test_litellm_annotations.py new file mode 100644 index 0000000000..f58009b86c --- /dev/null +++ b/tests/models/test_litellm_annotations.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from typing import Any + +from litellm.types.utils import Message + +from agents.extensions.models.litellm_model import LitellmConverter + + +def _message_with_annotations(annotations: list[dict[str, Any]]) -> Message: + # Pass raw provider-shaped dicts intentionally; these mirror the partial payloads + # LiteLLM forwards and would not satisfy the strict ChatCompletionAnnotation type. + return Message.model_construct(role="assistant", content="hi", annotations=annotations) + + +def test_convert_annotations_returns_none_when_absent() -> None: + message = Message(role="assistant", content="hi") + assert LitellmConverter.convert_annotations_to_openai(message) is None + + +def test_convert_annotations_maps_full_citation() -> None: + message = _message_with_annotations( + [ + { + "type": "url_citation", + "url_citation": { + "start_index": 1, + "end_index": 4, + "url": "https://example.com", + "title": "Example", + }, + } + ] + ) + result = LitellmConverter.convert_annotations_to_openai(message) + assert result is not None + assert len(result) == 1 + citation = result[0].url_citation + assert citation.start_index == 1 + assert citation.end_index == 4 + assert citation.url == "https://example.com" + assert citation.title == "Example" + + +def test_convert_annotations_defaults_missing_title() -> None: + # Providers reached through LiteLLM may omit fields the OpenAI schema marks as + # required; hard indexing those fields previously raised KeyError and aborted + # the whole turn instead of degrading gracefully. + message = _message_with_annotations( + [ + { + "type": "url_citation", + "url_citation": { + "start_index": 0, + "end_index": 5, + "url": "https://example.com", + }, + } + ] + ) + result = LitellmConverter.convert_annotations_to_openai(message) + assert result is not None + assert len(result) == 1 + assert result[0].url_citation.url == "https://example.com" + assert result[0].url_citation.title == "" + + +def test_convert_annotations_defaults_missing_indices_and_title() -> None: + message = _message_with_annotations( + [ + { + "type": "url_citation", + "url_citation": { + "url": "https://example.com", + "title": None, + "start_index": None, + "end_index": None, + }, + } + ] + ) + result = LitellmConverter.convert_annotations_to_openai(message) + assert result is not None + assert len(result) == 1 + citation = result[0].url_citation + assert citation.start_index == 0 + assert citation.end_index == 0 + assert citation.url == "https://example.com" + assert citation.title == "" + + +def test_convert_annotations_skips_entries_without_url_citation_payload() -> None: + # LiteLLM enforces type == "url_citation" but allows the url_citation payload to be + # absent; such an entry carries no citation data, so it is skipped rather than + # emitted as an empty citation. + message = _message_with_annotations( + [ + {"type": "url_citation"}, + { + "type": "url_citation", + "url_citation": { + "start_index": 0, + "end_index": 2, + "url": "https://example.com", + "title": "Kept", + }, + }, + ] + ) + result = LitellmConverter.convert_annotations_to_openai(message) + assert result is not None + assert len(result) == 1 + assert result[0].url_citation.title == "Kept" + + +def test_convert_annotations_returns_none_when_no_usable_citations() -> None: + message = _message_with_annotations( + [ + {"type": "url_citation"}, + {"type": "url_citation", "url_citation": {"title": "Missing URL"}}, + ] + ) + + assert LitellmConverter.convert_annotations_to_openai(message) is None diff --git a/tests/test_run_step_execution.py b/tests/test_run_step_execution.py index 8163b5a099..73e0bb82f6 100644 --- a/tests/test_run_step_execution.py +++ b/tests/test_run_step_execution.py @@ -1463,7 +1463,8 @@ async def _second_tool() -> str: ] loop = asyncio.get_running_loop() previous_task_factory = loop.get_task_factory() - eager_task_factory = cast(Any, asyncio.eager_task_factory) + asyncio_module = cast(Any, asyncio) + eager_task_factory = asyncio_module.eager_task_factory loop.set_task_factory(eager_task_factory) try: