diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 022622d7af7..3ca475848f6 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -44,6 +44,7 @@ class ExaAISearchRequest(_ExaAISearchRequestRequired, total=False): context: bool | dict # Optional - format results for LLMs moderation: bool # Optional - enable content moderation, default false contents: dict # Optional - content retrieval options + outputSchema: dict # Optional - output schema definition class ExaAISearchConfig(BaseSearchConfig): @@ -189,7 +190,12 @@ class ExaAISearchConfig(BaseSearchConfig): ) results.append(search_result) - return SearchResponse( + extra: dict = {} + if "output" in response_json: + extra["output"] = response_json["output"] + search_response = SearchResponse( results=results, object="search", + **extra, ) + return search_response diff --git a/tests/search_tests/test_exa_ai_search.py b/tests/search_tests/test_exa_ai_search.py index 7974668a61f..8461e90911e 100644 --- a/tests/search_tests/test_exa_ai_search.py +++ b/tests/search_tests/test_exa_ai_search.py @@ -1,7 +1,9 @@ import pytest import litellm from typing import List, Union +from unittest.mock import Mock +from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig from tests.search_tests.base_search_unit_tests import BaseSearchTest @@ -15,3 +17,68 @@ class TestExaAISearch(BaseSearchTest): Return search_provider for Exa AI Search. """ return "exa_ai" + + +class TestExaAISearchTransformation: + def test_should_include_output_when_exa_returns_output(self): + config = ExaAISearchConfig() + raw_response = Mock() + raw_response.json.return_value = { + "output": { + "content": "Nvidia announced the Vera Rubin platform.", + "grounding": [ + { + "field": "content", + "citations": [ + { + "url": "https://nvidianews.nvidia.com/news/test", + "title": "NVIDIA Newsroom", + } + ], + "confidence": "high", + } + ], + }, + "results": [ + { + "title": "NVIDIA Newsroom", + "url": "https://nvidianews.nvidia.com/news/test", + "text": "NVIDIA announced Vera Rubin.", + "publishedDate": "2026-06-22T00:00:00.000Z", + } + ], + } + + response = config.transform_search_response( + raw_response=raw_response, + logging_obj=None, + ) + + assert response.output == raw_response.json.return_value["output"] + assert ( + response.model_dump()["output"] + == raw_response.json.return_value["output"] + ) + assert response.results[0].title == "NVIDIA Newsroom" + + def test_should_not_require_output_when_exa_omits_output(self): + config = ExaAISearchConfig() + raw_response = Mock() + raw_response.json.return_value = { + "results": [ + { + "title": "NVIDIA Newsroom", + "url": "https://nvidianews.nvidia.com/news/test", + "text": "NVIDIA announced Vera Rubin.", + } + ], + } + + response = config.transform_search_response( + raw_response=raw_response, + logging_obj=None, + ) + + assert "output" not in response.model_dump() + assert response.object == "search" + assert response.results[0].title == "NVIDIA Newsroom"