fix web search tool

Author: Tonyhrule

src/cleanlab_tlm/utils/chat.py

@@ -755,32 +755,14 @@ def _responses_messages_to_string(messages: list[dict[str, Any]]) -> str:
                 )
 
             if message["action"]["type"] == "search":
-                if "annotations" in message:
-                    annotations = message["annotations"]
-                else:
-                    next_text_message = adjusted_messages[i + 1]
-                    if next_text_message["type"] != "message":
-                        continue
-                    next_text_content = next_text_message["content"][0]
-                    if next_text_content["type"] != "output_text":
-                        continue
-                    annotations = next_text_content["annotations"]
-
-                urls = list(
-                    {
-                        (annotation["url"], annotation["title"])
-                        for annotation in annotations
-                        if annotation["type"] == "url_citation"
-                    }
-                )
+                urls = list({source["url"] for source in message["action"]["sources"] if source["type"] == "url"})
 
                 with ThreadPoolExecutor() as executor:
 
-                    def extract_text(pair: tuple[str, str]) -> str:
+                    def extract_text(url: str) -> str:
                         fallback_text = "Response is not shown, but the LLM can still access it. Assume that whatever the LLM references in this URL is true."
 
                         try:
-                            url = pair[0]
                             if url in _url_cache:
                                 return _url_cache[url]
 
@@ -814,10 +796,9 @@ def extract_text(pair: tuple[str, str]) -> str:
                 websites = [
                     {
                         "url": url,
-                        "title": title,
                         "content": data,
                     }
-                    for (url, title), data in zip(urls, requests)
+                    for url, data in zip(urls, requests)
                 ]
 
                 tool_call = {

tests/test_chat.py

@@ -11,6 +11,7 @@
 )
 from openai.types.responses.response_function_web_search import (
     ActionSearch,
+    ActionSearchSource,
     ResponseFunctionWebSearch,
 )
 from openai.types.responses.response_output_message import ResponseOutputMessage
@@ -1308,6 +1309,12 @@ def test_form_prompt_string_responses_web_search() -> None:
                 action=ActionSearch(
                     query="Give me a positive news story from today",
                     type="search",
+                    sources=[
+                        ActionSearchSource(
+                            type="url",
+                            url="https://www.podego.com/insights/august-2025-good-news-ai-pfas-stories?utm_source=openai",
+                        ),
+                    ],
                 ),
                 status="completed",
                 type="web_search_call",
@@ -1394,7 +1401,6 @@ def test_form_prompt_string_responses_web_search() -> None:
   "output": [
     {
       "url": "https://www.podego.com/insights/august-2025-good-news-ai-pfas-stories?utm_source=openai",
-      "title": "Positive News Highlights | AI, PFAS Breakthroughs & More \\u2014 August 2025 \\u2014 Podego",
       "content": "MOCK CONTENT"
     }
   ]
@@ -2342,6 +2348,12 @@ def test_form_response_string_responses_web_search() -> None:
                 action=ActionSearch(
                     query="Give me a positive news story from today",
                     type="search",
+                    sources=[
+                        ActionSearchSource(
+                            type="url",
+                            url="https://www.podego.com/insights/august-2025-good-news-ai-pfas-stories?utm_source=openai",
+                        ),
+                    ],
                 ),
                 status="completed",
                 type="web_search_call",