Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 21 additions & 11 deletions src/agents/extensions/models/litellm_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
124 changes: 124 additions & 0 deletions tests/models/test_litellm_annotations.py
Original file line number Diff line number Diff line change
@@ -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
3 changes: 2 additions & 1 deletion tests/test_run_step_execution.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down