feat: add support for exa /search outputschema

This commit is contained in:
noahpodgurski 2026-06-22 11:57:00 -04:00
parent e0c8a6b483
commit e548d53a7d
2 changed files with 72 additions and 1 deletions

View file

@ -46,6 +46,7 @@ class ExaAISearchRequest(_ExaAISearchRequestRequired, total=False):
context: Union[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):
@ -196,7 +197,10 @@ class ExaAISearchConfig(BaseSearchConfig):
)
results.append(search_result)
return SearchResponse(
search_response = SearchResponse(
results=results,
object="search",
)
if "output" in response_json:
search_response.output = response_json["output"]
return search_response

View file

@ -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"