add tavily

This commit is contained in:
Dinesh Yeduguru 2024-12-20 21:01:41 -08:00
parent dcdf9da6ef
commit 9192a9bbb4
6 changed files with 163 additions and 2 deletions

View file

@ -149,7 +149,7 @@ async def create_agent_turn_with_search_tool(
tool_execution = tool_execution_events[0].event.payload.step_details
assert isinstance(tool_execution, ToolExecutionStep)
assert len(tool_execution.tool_calls) > 0
assert tool_execution.tool_calls[0].tool_name == BuiltinTool.brave_search
assert tool_execution.tool_calls[0].tool_name == tool_name
assert len(tool_execution.tool_responses) > 0
check_turn_complete_event(turn_response, session_id, search_query_messages)
@ -302,6 +302,20 @@ class TestAgents:
"brave_search",
)
@pytest.mark.asyncio
async def test_create_agent_turn_with_tavily_search(
self, agents_stack, search_query_messages, common_params
):
if "TAVILY_SEARCH_API_KEY" not in os.environ:
pytest.skip("TAVILY_SEARCH_API_KEY not set, skipping test")
await create_agent_turn_with_search_tool(
agents_stack,
search_query_messages,
common_params,
"tavily_search",
)
def check_event_types(turn_response):
event_types = [chunk.event.payload.event_type for chunk in turn_response]