diff --git a/tests/conftest.py b/tests/conftest.py index 221e511..4fdb767 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -139,28 +139,23 @@ def favorite_color(person: str): def assert_tools_sequential(chat_fun: ChatFun, total_calls: int, stream: bool = True): chat = chat_fun(system_prompt="Be very terse, not even punctuation.") - def get_current_year(): - """Gets the current year""" - return 2024 + def forecast(city: str): + """Gets the weather forecast for a city""" + return "rainy" if city == "New York" else "sunny" - chat.register_tool(get_current_year) + chat.register_tool(forecast) - def popular_name(year: int): - """Gets the most popular name for a given year""" - if isinstance(year, str): # Sometimes Google sends the year as a string? - year = int(year) - return "Susan" if year == 2024 else "I don't know" + def equipment(weather: str): + """Gets the equipment needed for a weather condition""" + return "umbrella" if weather == "rainy" else "sunscreen" - chat.register_tool(popular_name) + chat.register_tool(equipment) response = chat.chat( - """ - What was the most popular name this year? - Note that you have a tool available to you to find the current year. - """, + "What should I pack for New York this weekend?", stream=stream, ) - assert "Susan" in str(response) + assert "umbrella" in str(response).lower() assert len(chat.turns()) == total_calls