From 4f62eca047512aa5afa5fad1007dcf1af366289f Mon Sep 17 00:00:00 2001 From: rjambrecic <32619626+rjambrecic@users.noreply.github.com> Date: Fri, 5 Jul 2024 08:47:03 +0200 Subject: [PATCH] Create weatherapi toolbox from url and integrate with the team (#816) * wip * wip * Add tests for weather team tools * Weatherman team test wip * Refactor helper_test_init function and add test for the weather team * Add fixture for starting FastAPI app in the conftest.py * Add end2end test for weather team * Fix import from incorrect lib * Update dependecies --- captn/captn_agents/backend/teams/__init__.py | 2 + .../backend/teams/_gbb_initial_team.py | 4 +- .../backend/teams/_weather_team.py | 155 +++++++++++++++ .../backend/tools/_weather_team_tools.py | 38 ++++ pyproject.toml | 6 +- .../captn_agents/backend/teams/helpers.py | 18 +- .../backend/teams/test_brief_creation_team.py | 10 +- .../teams/test_campaign_creation_team.py | 10 +- .../backend/teams/test_gbb_initial_team.py | 14 +- .../backend/teams/test_google_ads_team.py | 12 +- .../backend/teams/test_weather_team.py | 61 ++++++ .../teams/test_weekly_analysis_team.py | 12 +- .../tools/fixtures/weather_openapi.json | 184 ++++++++++++++++++ .../backend/tools/test_weather_team_tools.py | 85 ++++++++ tests/ci/conftest.py | 68 ++++++- tests/ci/test_conftest.py | 21 ++ 16 files changed, 672 insertions(+), 28 deletions(-) create mode 100644 captn/captn_agents/backend/teams/_weather_team.py create mode 100644 captn/captn_agents/backend/tools/_weather_team_tools.py create mode 100644 tests/ci/captn/captn_agents/backend/teams/test_weather_team.py create mode 100644 tests/ci/captn/captn_agents/backend/tools/fixtures/weather_openapi.json create mode 100644 tests/ci/captn/captn_agents/backend/tools/test_weather_team_tools.py create mode 100644 tests/ci/test_conftest.py diff --git a/captn/captn_agents/backend/teams/__init__.py b/captn/captn_agents/backend/teams/__init__.py index ab6ac7af..1d597b09 100644 --- a/captn/captn_agents/backend/teams/__init__.py +++ b/captn/captn_agents/backend/teams/__init__.py @@ -3,6 +3,7 @@ from ._gbb_initial_team import GBBInitialTeam from ._google_ads_team import GoogleAdsTeam from ._team import Team +from ._weather_team import WeatherTeam from ._weekly_analysis_team import ( REACT_APP_API_URL, WeeklyAnalysisTeam, @@ -14,6 +15,7 @@ "CampaignCreationTeam", "GBBInitialTeam", "WeeklyAnalysisTeam", + "WeatherTeam", "GoogleAdsTeam", "REACT_APP_API_URL", "Team", diff --git a/captn/captn_agents/backend/teams/_gbb_initial_team.py b/captn/captn_agents/backend/teams/_gbb_initial_team.py index 97326244..8d29fc0a 100644 --- a/captn/captn_agents/backend/teams/_gbb_initial_team.py +++ b/captn/captn_agents/backend/teams/_gbb_initial_team.py @@ -6,6 +6,8 @@ from ._shared_prompts import REPLY_TO_CLIENT_COMMAND from ._team import Team +__all__ = ("GBBInitialTeam",) + @Team.register_team("gbb_initial_team") class GBBInitialTeam(BriefCreationTeam): @@ -65,7 +67,7 @@ def _guidelines(self) -> str: If you fail to choose the appropriate team, you will be penalized! 3. Here is a list of teams you can choose from after you determine which one is the most appropriate for the task: -{self.construct_team_names_and_descriptions_message(use_only_team_names={"campaign_creation_team"})} +{self.construct_team_names_and_descriptions_message(use_only_team_names={"weather_team"})} Guidelines SUMMARY: - Write a detailed step-by-step plan diff --git a/captn/captn_agents/backend/teams/_weather_team.py b/captn/captn_agents/backend/teams/_weather_team.py new file mode 100644 index 00000000..27330e78 --- /dev/null +++ b/captn/captn_agents/backend/teams/_weather_team.py @@ -0,0 +1,155 @@ +from typing import Any, Callable, Dict, List, Optional, Tuple + +from fastagency.openapi.client import Client + +from ..config import Config +from ..toolboxes import Toolbox +from ..tools._weather_team_tools import ( + create_weather_team_client, + create_weather_team_toolbox, +) +from ._shared_prompts import REPLY_TO_CLIENT_COMMAND +from ._team import Team + + +@Team.register_team("weather_team") +class WeatherTeam(Team): + _default_roles = [ + { + "Name": "Weather_forecaster", + "Description": """You are a weather forecaster. +Never introduce yourself when writing messages. E.g. do not write 'As a ...'""", + }, + { + "Name": "News_reporter", + "Description": """You are a news reporter. +You are also SOLELY responsible for communicating with the client. + +Based on the initial task, a number of proposed solutions will be suggested by the team. You must ask the team to write a detailed plan +including steps and expected outcomes. +Once the initial task given to the team is completed by implementing proposed solutions, you must write down the +accomplished work and execute the 'reply_to_client' command. That message will be forwarded to the client so make +sure it is understandable by non-experts. +Never introduce yourself when writing messages. E.g. do not write 'As an account manager' +""", + }, + ] + + _functions: Optional[List[Dict[str, Any]]] = [] + + def __init__( + self, + *, + task: str, + user_id: int, + conv_id: int, + work_dir: str = "weather_team", + max_round: int = 80, + seed: int = 42, + temperature: float = 0.2, + config_list: Optional[List[Dict[str, str]]] = None, + create_toolbox_func: Callable[ + [int, int], Toolbox + ] = create_weather_team_toolbox, + create_client_func: Callable[[str], Client] = create_weather_team_client, + openapi_url: str = "https://weather.tools.fastagency.ai/openapi.json", + ): + recommended_modifications_and_answer_list: List[ + Tuple[Dict[str, Any], Optional[str]] + ] = [] + function_map: Dict[str, Callable[[Any], Any]] = {} + + roles: List[Dict[str, str]] = self._default_roles + + super().__init__( + user_id=user_id, + conv_id=conv_id, + roles=roles, + task=task, + function_map=function_map, + work_dir=work_dir, + max_round=max_round, + seed=seed, + temperature=temperature, + recommended_modifications_and_answer_list=recommended_modifications_and_answer_list, + use_user_proxy=True, + ) + + if config_list is None: + config = Config() + config_list = config.config_list_gpt_4o + + self.llm_config = self._get_llm_config( + seed=seed, temperature=temperature, config_list=config_list + ) + self.create_toolbox_func = create_toolbox_func + self.create_client_func = create_client_func + self.openapi_url = openapi_url + + self._create_members() + + self._add_client() + self._add_tools() + + self._create_initial_message() + + def _add_client(self) -> None: + self.client = self.create_client_func(self.openapi_url) + + self.client.register_for_execution(self.user_proxy) + for agent in self.members: + # Add only for Weather_forecaster + if agent.name == self._default_roles[0]["Name"].lower(): + if agent.llm_config["tools"] is None: + agent.llm_config.pop("tools") + self.client.register_for_llm(agent) + + def _add_tools(self) -> None: + self.toolbox = self.create_toolbox_func( + self.user_id, + self.conv_id, + ) + # Add only for News_reporter + for agent in self.members: + if ( + agent != self.user_proxy + and agent.name != self._default_roles[0]["Name"].lower() + ): + self.toolbox.add_to_agent(agent, self.user_proxy) + + @property + def _task(self) -> str: + return f"""You are a team in charge of weather forecasting. + +Here is the current customers brief/information we have gathered for you as a starting point: +{self.task} +""" + + @property + def _guidelines(self) -> str: + return """### Guidelines +1. Do NOT repeat the content of the previous messages nor repeat your role. +""" + + @property + def _commands(self) -> str: + return f"""## Commands +Only News_reporter has access to the following command: +1. {REPLY_TO_CLIENT_COMMAND} +"smart_suggestions": {{ + 'suggestions': ['Give me suggestions what I can do based on the current weather forecast.'], + 'type': 'oneOf' +}} + +2. Only Weather_forecaster has access to weather API to get the weather forecast for the city. +""" + + @classmethod + def get_capabilities(cls) -> str: + return "Weather forecasting for any city." + + @classmethod + def get_brief_template(cls) -> str: + return ( + "We only need the name of the city for which you want the weather forecast." + ) diff --git a/captn/captn_agents/backend/tools/_weather_team_tools.py b/captn/captn_agents/backend/tools/_weather_team_tools.py new file mode 100644 index 00000000..d52a22fa --- /dev/null +++ b/captn/captn_agents/backend/tools/_weather_team_tools.py @@ -0,0 +1,38 @@ +import httpx +from fastagency.openapi.client import Client + +from ..toolboxes import Toolbox +from ._functions import REPLY_TO_CLIENT_DESCRIPTION, BaseContext, reply_to_client + +__all__ = ( + "create_weather_team_client", + "create_weather_team_toolbox", +) + + +def create_weather_team_client(openapi_url: str) -> Client: + with httpx.Client() as httpx_client: + response = httpx_client.get(openapi_url) + response.raise_for_status() + openapi_spec = response.text + + client = Client.create(openapi_spec) + + return client + + +def create_weather_team_toolbox( + user_id: int, + conv_id: int, +) -> Toolbox: + toolbox = Toolbox() + + context = BaseContext( + user_id=user_id, + conv_id=conv_id, + ) + toolbox.set_context(context) + + toolbox.add_function(REPLY_TO_CLIENT_DESCRIPTION)(reply_to_client) + + return toolbox diff --git a/pyproject.toml b/pyproject.toml index 6eb1cb87..23cc1cd6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,13 +84,13 @@ testing = [ ] benchmarking = [ - "typer==0.12.3", + "typer==0.9.4", # downgraded because fastagency "filelock==3.15.4", "tabulate==0.9.0", ] agents = [ - "fastapi==0.111.0", + "fastapi>=0.110.2,<0.111.0", # downgraded because fastagency "APScheduler==3.10.4", "prisma==0.13.1", "google-ads==24.1.0", @@ -111,6 +111,8 @@ agents = [ "opentelemetry-exporter-otlp==1.25.0", "openpyxl==3.1.5", "aiofiles==24.1.0", + "fastagency@git+https://github.com/airtai/fastagency.git@3a2346c", + "python-multipart==0.0.9", # remove after fastagency is updated ] dev = [ diff --git a/tests/ci/captn/captn_agents/backend/teams/helpers.py b/tests/ci/captn/captn_agents/backend/teams/helpers.py index cd1eb755..0980866d 100644 --- a/tests/ci/captn/captn_agents/backend/teams/helpers.py +++ b/tests/ci/captn/captn_agents/backend/teams/helpers.py @@ -1,4 +1,4 @@ -from typing import Type +from typing import Dict, Type from autogen.agentchat import UserProxyAgent @@ -7,23 +7,28 @@ def helper_test_init( team: Team, - number_of_team_members: int, - number_of_functions: int, + number_of_registered_executions: int, + agent_number_of_functions_dict: Dict[str, int], team_class: Type[Team], ) -> None: try: assert isinstance(team, team_class) - assert len(team.members) == number_of_team_members + assert len(team.members) == len(agent_number_of_functions_dict) for agent in team.members: # execution of the tools number_of_functions_in_function_map = len(agent.function_map) if isinstance(agent, UserProxyAgent): - assert number_of_functions_in_function_map == number_of_functions + print(agent.function_map) + assert ( + number_of_functions_in_function_map + == number_of_registered_executions + ) else: assert number_of_functions_in_function_map == 0 + number_of_functions = agent_number_of_functions_dict[agent.name] # specification of the tools llm_config = agent.llm_config if not isinstance(agent, UserProxyAgent): @@ -36,7 +41,8 @@ def helper_test_init( tool["function"]["name"] for tool in llm_config["tools"] ] - assert set(team.user_proxy.function_map.keys()) == set(function_names) + # Check if the agent's function names are in the user_proxy's function map + assert set(function_names) <= set(team.user_proxy.function_map.keys()) else: assert llm_config is False diff --git a/tests/ci/captn/captn_agents/backend/teams/test_brief_creation_team.py b/tests/ci/captn/captn_agents/backend/teams/test_brief_creation_team.py index 21df183f..a3528982 100644 --- a/tests/ci/captn/captn_agents/backend/teams/test_brief_creation_team.py +++ b/tests/ci/captn/captn_agents/backend/teams/test_brief_creation_team.py @@ -27,10 +27,16 @@ def test_init(self) -> None: task="do your magic", ) + agent_number_of_functions_dict = { + "digitial_marketing_strategist": 4, + "account_manager": 4, + "user_proxy": 0, + } + helper_test_init( team=brief_creation_team, - number_of_team_members=3, - number_of_functions=4, + number_of_registered_executions=4, + agent_number_of_functions_dict=agent_number_of_functions_dict, team_class=BriefCreationTeam, ) diff --git a/tests/ci/captn/captn_agents/backend/teams/test_campaign_creation_team.py b/tests/ci/captn/captn_agents/backend/teams/test_campaign_creation_team.py index 3238e5bb..0402848a 100644 --- a/tests/ci/captn/captn_agents/backend/teams/test_campaign_creation_team.py +++ b/tests/ci/captn/captn_agents/backend/teams/test_campaign_creation_team.py @@ -19,10 +19,16 @@ def test_init(self) -> None: task="do your magic", ) + agent_number_of_functions_dict = { + "copywriter": 7, + "account_manager": 7, + "user_proxy": 0, + } + helper_test_init( team=campaign_creation_team, - number_of_team_members=3, - number_of_functions=7, + number_of_registered_executions=7, + agent_number_of_functions_dict=agent_number_of_functions_dict, team_class=CampaignCreationTeam, ) diff --git a/tests/ci/captn/captn_agents/backend/teams/test_gbb_initial_team.py b/tests/ci/captn/captn_agents/backend/teams/test_gbb_initial_team.py index 59ef8e6e..2c79e652 100644 --- a/tests/ci/captn/captn_agents/backend/teams/test_gbb_initial_team.py +++ b/tests/ci/captn/captn_agents/backend/teams/test_gbb_initial_team.py @@ -17,15 +17,21 @@ def setup(self) -> Iterator[None]: yield def test_init(self) -> None: - brief_creation_team = GBBInitialTeam( + gbb_initial_team = GBBInitialTeam( user_id=123, conv_id=456, task="do your magic", ) + agent_number_of_functions_dict = { + "digitial_marketing_strategist": 3, + "account_manager": 3, + "user_proxy": 0, + } + helper_test_init( - team=brief_creation_team, - number_of_team_members=3, - number_of_functions=3, + team=gbb_initial_team, + number_of_registered_executions=3, + agent_number_of_functions_dict=agent_number_of_functions_dict, team_class=GBBInitialTeam, ) diff --git a/tests/ci/captn/captn_agents/backend/teams/test_google_ads_team.py b/tests/ci/captn/captn_agents/backend/teams/test_google_ads_team.py index fab33480..f717ad37 100644 --- a/tests/ci/captn/captn_agents/backend/teams/test_google_ads_team.py +++ b/tests/ci/captn/captn_agents/backend/teams/test_google_ads_team.py @@ -11,9 +11,17 @@ def test_init(self) -> None: task="do your magic", ) + agent_number_of_functions_dict = { + "google_ads_specialist": 21, + "copywriter": 21, + "digital_strategist": 21, + "account_manager": 21, + "user_proxy": 0, + } + helper_test_init( team=google_ads_team, - number_of_team_members=5, - number_of_functions=21, + number_of_registered_executions=21, + agent_number_of_functions_dict=agent_number_of_functions_dict, team_class=GoogleAdsTeam, ) diff --git a/tests/ci/captn/captn_agents/backend/teams/test_weather_team.py b/tests/ci/captn/captn_agents/backend/teams/test_weather_team.py new file mode 100644 index 00000000..3a0c684b --- /dev/null +++ b/tests/ci/captn/captn_agents/backend/teams/test_weather_team.py @@ -0,0 +1,61 @@ +from typing import Iterator + +import pytest + +from captn.captn_agents.backend.teams import ( + Team, + WeatherTeam, +) + +from .helpers import helper_test_init + + +class TestWeatherTeam: + @pytest.fixture(autouse=True) + def setup(self) -> Iterator[None]: + Team._teams.clear() + yield + + def test_init(self) -> None: + weather_team = WeatherTeam( + task="do your magic", + user_id=123, + conv_id=456, + ) + agent_number_of_functions_dict = { + "weather_forecaster": 1, + "news_reporter": 1, + "user_proxy": 0, + } + + helper_test_init( + team=weather_team, + number_of_registered_executions=2, + agent_number_of_functions_dict=agent_number_of_functions_dict, + team_class=WeatherTeam, + ) + + @pytest.mark.flaky + @pytest.mark.openai + @pytest.mark.fastapi_openapi_team + def test_weather_team_end2end(self, weather_fastapi_openapi_url: str) -> None: + weather_team = WeatherTeam( + task="What's the weather like in London?", + user_id=123, + conv_id=456, + openapi_url=weather_fastapi_openapi_url, + ) + + weather_team.initiate_chat() + + messages = weather_team.get_messages() + success = False + for message in messages: + if ( + "tool_responses" in message + and message["content"] == "Weather in London is sunny" + ): + success = True + break + + assert success, messages diff --git a/tests/ci/captn/captn_agents/backend/teams/test_weekly_analysis_team.py b/tests/ci/captn/captn_agents/backend/teams/test_weekly_analysis_team.py index b4a18de2..4877ba4d 100644 --- a/tests/ci/captn/captn_agents/backend/teams/test_weekly_analysis_team.py +++ b/tests/ci/captn/captn_agents/backend/teams/test_weekly_analysis_team.py @@ -1368,10 +1368,18 @@ def test_init(self) -> None: task="do your magic", ) + agent_number_of_functions_dict = { + "google_ads_specialist": 4, + "copywriter": 4, + "digital_strategist": 4, + "account_manager": 4, + "user_proxy": 0, + } + helper_test_init( team=weekly_analysis_team, - number_of_team_members=5, - number_of_functions=4, + number_of_registered_executions=4, + agent_number_of_functions_dict=agent_number_of_functions_dict, team_class=WeeklyAnalysisTeam, ) diff --git a/tests/ci/captn/captn_agents/backend/tools/fixtures/weather_openapi.json b/tests/ci/captn/captn_agents/backend/tools/fixtures/weather_openapi.json new file mode 100644 index 00000000..fe4bd591 --- /dev/null +++ b/tests/ci/captn/captn_agents/backend/tools/fixtures/weather_openapi.json @@ -0,0 +1,184 @@ +{ + "openapi": "3.1.0", + "info": { + "title": "WeatherAPI", + "version": "0.1.0" + }, + "servers": [ + { + "url": "https://weather.tools.fastagency.ai", + "description": "Weather app server" + } + ], + "paths": { + "/": { + "get": { + "summary": "Get Weather", + "description": "Get weather forecast for a given city", + "operationId": "get_weather__get", + "parameters": [ + { + "name": "city", + "in": "query", + "required": true, + "schema": { + "type": "string", + "description": "city for which forecast is requested", + "title": "City" + }, + "description": "city for which forecast is requested" + } + ], + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Weather" + } + } + } + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } + } + }, + "components": { + "schemas": { + "DailyForecast": { + "properties": { + "forecast_date": { + "type": "string", + "format": "date", + "title": "Forecast Date" + }, + "temperature": { + "type": "integer", + "title": "Temperature" + }, + "hourly_forecasts": { + "items": { + "$ref": "#/components/schemas/HourlyForecast" + }, + "type": "array", + "title": "Hourly Forecasts" + } + }, + "type": "object", + "required": [ + "forecast_date", + "temperature", + "hourly_forecasts" + ], + "title": "DailyForecast" + }, + "HTTPValidationError": { + "properties": { + "detail": { + "items": { + "$ref": "#/components/schemas/ValidationError" + }, + "type": "array", + "title": "Detail" + } + }, + "type": "object", + "title": "HTTPValidationError" + }, + "HourlyForecast": { + "properties": { + "forecast_time": { + "type": "string", + "format": "time", + "title": "Forecast Time" + }, + "temperature": { + "type": "integer", + "title": "Temperature" + }, + "description": { + "type": "string", + "title": "Description" + } + }, + "type": "object", + "required": [ + "forecast_time", + "temperature", + "description" + ], + "title": "HourlyForecast" + }, + "ValidationError": { + "properties": { + "loc": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "type": "array", + "title": "Location" + }, + "msg": { + "type": "string", + "title": "Message" + }, + "type": { + "type": "string", + "title": "Error Type" + } + }, + "type": "object", + "required": [ + "loc", + "msg", + "type" + ], + "title": "ValidationError" + }, + "Weather": { + "properties": { + "city": { + "type": "string", + "title": "City" + }, + "temperature": { + "type": "integer", + "title": "Temperature" + }, + "daily_forecasts": { + "items": { + "$ref": "#/components/schemas/DailyForecast" + }, + "type": "array", + "title": "Daily Forecasts" + } + }, + "type": "object", + "required": [ + "city", + "temperature", + "daily_forecasts" + ], + "title": "Weather" + } + } + } + } diff --git a/tests/ci/captn/captn_agents/backend/tools/test_weather_team_tools.py b/tests/ci/captn/captn_agents/backend/tools/test_weather_team_tools.py new file mode 100644 index 00000000..1f3b9555 --- /dev/null +++ b/tests/ci/captn/captn_agents/backend/tools/test_weather_team_tools.py @@ -0,0 +1,85 @@ +import unittest +from pathlib import Path +from typing import Iterator +from unittest.mock import MagicMock + +import pytest +from autogen.agentchat import AssistantAgent, UserProxyAgent + +from captn.captn_agents.backend.config import Config +from captn.captn_agents.backend.teams._team import Team +from captn.captn_agents.backend.tools._weather_team_tools import ( + create_weather_team_client, + create_weather_team_toolbox, +) + +from .helpers import check_llm_config_descriptions, check_llm_config_total_tools + + +class TestTools: + @pytest.fixture(autouse=True) + def setup(self) -> Iterator[None]: + self.llm_config = { + "config_list": Config().config_list_gpt_3_5, + } + + self.toolbox = create_weather_team_toolbox( + user_id=12345, + conv_id=67890, + ) + + yield + Team._teams.clear() + + def test_llm_config(self) -> None: + agent = AssistantAgent(name="agent", llm_config=self.llm_config) + user_proxy = UserProxyAgent(name="user_proxy", code_execution_config=False) + + self.toolbox.add_to_agent(agent, user_proxy) + + llm_config = agent.llm_config + name_desc_dict = { + "reply_to_client": "Respond to the client", + } + + check_llm_config_total_tools(llm_config, 1) + check_llm_config_descriptions(llm_config, name_desc_dict) + + +class TestClient: + @pytest.fixture(autouse=True) + def setup(self) -> Iterator[None]: + self.llm_config = { + "config_list": Config().config_list_gpt_3_5, + } + + path = Path(__file__).parent / "fixtures" / "weather_openapi.json" + with Path.open(path, "r") as f: + weather_openapi_text = f.read() + + with unittest.mock.patch("httpx.Client.get") as mock_httpx_get: + mock_response = MagicMock() + mock_response.text = weather_openapi_text + mock_response.raise_for_status.return_value = None + + mock_httpx_get.return_value = mock_response + + self.client = create_weather_team_client("https://weather.com/openapi.json") + + yield + Team._teams.clear() + + def test_llm_config(self) -> None: + agent = AssistantAgent(name="agent", llm_config=self.llm_config) + user_proxy = UserProxyAgent(name="user_proxy", code_execution_config=False) + + self.client.register_for_execution(user_proxy) + self.client.register_for_llm(agent) + + llm_config = agent.llm_config + name_desc_dict = { + "get_weather__get": "Get weather forecast for a given city", + } + + check_llm_config_total_tools(llm_config, 1) + check_llm_config_descriptions(llm_config, name_desc_dict) diff --git a/tests/ci/conftest.py b/tests/ci/conftest.py index d0160935..fc41052b 100644 --- a/tests/ci/conftest.py +++ b/tests/ci/conftest.py @@ -1,11 +1,65 @@ -from typing import Any +import contextlib +import socket +import threading +import time +from platform import system +from typing import Annotated, Iterator import pytest +import uvicorn +from fastapi import FastAPI, Path -@pytest.fixture() -def patch_envs(monkeypatch: Any) -> None: # noqa: PT004 - monkeypatch.setenv("AZURE_OPENAI_API_KEY", "dummy_key") - monkeypatch.setenv("AZURE_API_ENDPOINT", "dummy_endpoint") - monkeypatch.setenv("AZURE_GPT4_MODEL", "airt-gpt4") - monkeypatch.setenv("AZURE_GPT35_MODEL", "gpt-35-turbo-16k") +def create_weather_fastapi_app(host: str, port: int) -> FastAPI: + app = FastAPI( + title="Weather", + servers=[ + {"url": f"http://{host}:{port}", "description": "Local development server"} + ], + ) + + @app.get("/forecast/{city}", description="Get the weather forecast for a city") + def forecast( + city: Annotated[str, Path(description="name of the city")], + ) -> str: + return f"Weather in {city} is sunny" + + return app + + +class Server(uvicorn.Server): # type: ignore [misc] + def install_signal_handlers(self) -> None: + pass + + @contextlib.contextmanager + def run_in_thread(self) -> Iterator[None]: + thread = threading.Thread(target=self.run) + thread.start() + try: + while not self.started: + time.sleep(1e-3) + yield + finally: + self.should_exit = True + thread.join() + + +def find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] # type: ignore [no-any-return] + + +@pytest.fixture(scope="session") +def weather_fastapi_openapi_url() -> Iterator[str]: + host = "127.0.0.1" + port = find_free_port() + app = create_weather_fastapi_app(host, port) + openapi_url = f"http://{host}:{port}/openapi.json" + + config = uvicorn.Config(app, host=host, port=port, log_level="info") + server = Server(config=config) + with server.run_in_thread(): + time.sleep(1 if system() != "Windows" else 5) # let the server start + + yield openapi_url diff --git a/tests/ci/test_conftest.py b/tests/ci/test_conftest.py new file mode 100644 index 00000000..dfc97ba4 --- /dev/null +++ b/tests/ci/test_conftest.py @@ -0,0 +1,21 @@ +import httpx + +from .conftest import find_free_port + + +def test_find_free_port() -> None: + port = find_free_port() + assert isinstance(port, int) + assert 1024 <= port <= 65535 + + +def test_weather_fastapi_openapi(weather_fastapi_openapi_url: str) -> None: + assert isinstance(weather_fastapi_openapi_url, str) + + resp = httpx.get(weather_fastapi_openapi_url) + assert resp.status_code == 200 + resp_json = resp.json() + assert "openapi" in resp_json + assert "servers" in resp_json + assert len(resp_json["servers"]) == 1 + assert resp_json["info"]["title"] == "Weather"