From 3dbbc89fd274075fdaf5e1b85be130a9d00fc2dc Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 22 Mar 2025 10:53:47 -0700
Subject: [PATCH] test_openai_web_search

---
 tests/llm_translation/test_openai.py | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py
index 456190469a..97af8d93e4 100644
--- a/tests/llm_translation/test_openai.py
+++ b/tests/llm_translation/test_openai.py
@@ -17,6 +17,10 @@ import litellm
 from litellm import Choices, Message, ModelResponse
 from base_llm_unit_tests import BaseLLMChatTest
 import asyncio
+from litellm.types.llms.openai import (
+    ChatCompletionAnnotation,
+    ChatCompletionAnnotationURLCitation,
+)
 
 
 def test_openai_prediction_param():
@@ -393,8 +397,28 @@ def test_openai_chat_completion_streaming_handler_reasoning_content():
     assert response.choices[0].delta.reasoning_content == "."
 
 
+def validate_response_url_citation(url_citation: ChatCompletionAnnotationURLCitation):
+    assert "end_index" in url_citation
+    assert "start_index" in url_citation
+    assert "url" in url_citation
+
+
+def validate_model_response_contains_web_search_annotations(response: ModelResponse):
+    """validates litellm response contains web search annotations"""
+    message = response.choices[0].message
+    annotations: ChatCompletionAnnotation = message.annotations
+    print("annotations: ", annotations)
+    assert annotations is not None
+    assert isinstance(annotations, list)
+    for annotation in annotations:
+        assert annotation["type"] == "url_citation"
+        url_citation: ChatCompletionAnnotationURLCitation = annotation["url_citation"]
+        validate_response_url_citation(url_citation)
+
+
 def test_openai_web_search():
-    # litellm._turn_on_debug()
+    """Makes a simple web search request and validates the response contains web search annotations and all expected fields are present"""
+    litellm._turn_on_debug()
     response = litellm.completion(
         model="openai/gpt-4o-search-preview",
         messages=[
@@ -405,3 +429,4 @@ def test_openai_web_search():
         ],
     )
     print("litellm response: ", response.model_dump_json(indent=4))
+    validate_model_response_contains_web_search_annotations(response)