From f3433b0af5f74df2095b3a46f8f3cd98b177e357 Mon Sep 17 00:00:00 2001 From: Evan Mattson <35585003+moonbox3@users.noreply.github.com> Date: Mon, 5 Aug 2024 12:57:44 -0400 Subject: [PATCH] Python: Refactor agent retrieve method to be classmethod. Update tests. (#7854) ### Motivation and Context In working with the agent's retrieve method, it felt awkward to have to create an agent object and then use that object to then retrieve the OpenAI assistant that was previously defined. ### Description Making the retrieve method a class method on the Agent class. Refactored some of the settings creation and client creation methods for use between multiple methods in the agent class. Updating unit tests. - Updates openai package to 1.39.0 ### Contribution Checklist - [X] The code builds clean without any errors or warnings - [X] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [X] All unit tests pass, and I have added new tests where possible - [ ] I didn't break anyone :smile: --- .../agents/open_ai/azure_assistant_agent.py | 94 +++++++++++++++---- .../agents/open_ai/open_ai_assistant_agent.py | 89 ++++++++++++++---- .../agents/open_ai/open_ai_assistant_base.py | 12 ++- .../unit/agents/test_azure_assistant_agent.py | 51 +++++----- .../agents/test_open_ai_assistant_agent.py | 52 +++++----- 5 files changed, 197 insertions(+), 101 deletions(-) diff --git a/python/semantic_kernel/agents/open_ai/azure_assistant_agent.py b/python/semantic_kernel/agents/open_ai/azure_assistant_agent.py index 28e085b9cc78..17032c3574e3 100644 --- a/python/semantic_kernel/agents/open_ai/azure_assistant_agent.py +++ b/python/semantic_kernel/agents/open_ai/azure_assistant_agent.py @@ -38,7 +38,7 @@ def __init__( service_id: str | None = None, deployment_name: str | None = None, api_key: str | None = None, - endpoint: str | None = None, + endpoint: HttpsUrl | None = None, api_version: str | None = None, ad_token: str | None = None, ad_token_provider: Callable[[], str | Awaitable[str]] | None = None, @@ -100,23 +100,21 @@ def __init__( Raises: AgentInitializationError: If the api_key is not provided in the configuration. """ - try: - azure_openai_settings = AzureOpenAISettings.create( - api_key=api_key, - endpoint=endpoint, - chat_deployment_name=deployment_name, - api_version=api_version, - env_file_path=env_file_path, - env_file_encoding=env_file_encoding, - ) - except ValidationError as ex: - raise AgentInitializationError("Failed to create Azure OpenAI settings.", ex) from ex + azure_openai_settings = AzureAssistantAgent._create_azure_openai_settings( + api_key=api_key, + endpoint=endpoint, + deployment_name=deployment_name, + api_version=api_version, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) if not azure_openai_settings.chat_deployment_name: raise AgentInitializationError("The Azure OpenAI chat_deployment_name is required.") if not azure_openai_settings.api_key and not ad_token and not ad_token_provider: raise AgentInitializationError("Please provide either api_key, ad_token or ad_token_provider.") + client = self._create_client( api_key=azure_openai_settings.api_key.get_secret_value() if azure_openai_settings.api_key else None, endpoint=azure_openai_settings.endpoint, @@ -165,7 +163,7 @@ async def create( service_id: str | None = None, deployment_name: str | None = None, api_key: str | None = None, - endpoint: str | None = None, + endpoint: HttpsUrl | None = None, api_version: str | None = None, ad_token: str | None = None, ad_token_provider: Callable[[], str | Awaitable[str]] | None = None, @@ -303,6 +301,42 @@ def _create_client( default_headers=merged_headers, ) + @staticmethod + def _create_azure_openai_settings( + api_key: str | None = None, + endpoint: HttpsUrl | None = None, + deployment_name: str | None = None, + api_version: str | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> AzureOpenAISettings: + """Create the Azure OpenAI settings. + + Args: + api_key: The Azure OpenAI API key. + endpoint: The Azure OpenAI endpoint. + deployment_name: The Azure OpenAI chat deployment name. + api_version: The Azure OpenAI API version. + env_file_path: The environment file path. + env_file_encoding: The environment file encoding. + + Returns: + An instance of the AzureOpenAISettings. + """ + try: + azure_openai_settings = AzureOpenAISettings.create( + api_key=api_key, + endpoint=endpoint, + chat_deployment_name=deployment_name, + api_version=api_version, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + except ValidationError as ex: + raise AgentInitializationError("Failed to create Azure OpenAI settings.", ex) from ex + + return azure_openai_settings + async def list_definitions(self) -> AsyncIterable[dict[str, Any]]: """List the assistant definitions. @@ -313,8 +347,10 @@ async def list_definitions(self) -> AsyncIterable[dict[str, Any]]: for assistant in assistants.data: yield self._create_open_ai_assistant_definition(assistant) + @classmethod async def retrieve( - self, + cls, + *, id: str, api_key: str | None = None, endpoint: HttpsUrl | None = None, @@ -324,6 +360,8 @@ async def retrieve( client: AsyncAzureOpenAI | None = None, kernel: "Kernel | None" = None, default_headers: dict[str, str] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, ) -> "AzureAssistantAgent": """Retrieve an assistant by ID. @@ -337,20 +375,36 @@ async def retrieve( client: The Azure OpenAI client. (optional) kernel: The Kernel instance. (optional) default_headers: The default headers. (optional) + env_file_path: The environment file path. (optional) + env_file_encoding: The environment file encoding. (optional) Returns: - An OpenAIAssistantAgent instance. + An AzureAssistantAgent instance. """ - client = self._create_client( + azure_openai_settings = AzureAssistantAgent._create_azure_openai_settings( api_key=api_key, endpoint=endpoint, api_version=api_version, - ad_token=ad_token, - ad_token_provider=ad_token_provider, - default_headers=default_headers, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, ) + + if not azure_openai_settings.chat_deployment_name: + raise AgentInitializationError("The Azure OpenAI chat_deployment_name is required.") + if not azure_openai_settings.api_key and not ad_token and not ad_token_provider: + raise AgentInitializationError("Please provide either api_key, ad_token or ad_token_provider.") + + if not client: + client = AzureAssistantAgent._create_client( + api_key=api_key, + endpoint=endpoint, + api_version=api_version, + ad_token=ad_token, + ad_token_provider=ad_token_provider, + default_headers=default_headers, + ) assistant = await client.beta.assistants.retrieve(id) - assistant_definition = self._create_open_ai_assistant_definition(assistant) + assistant_definition = OpenAIAssistantBase._create_open_ai_assistant_definition(assistant) return AzureAssistantAgent(kernel=kernel, **assistant_definition) # endregion diff --git a/python/semantic_kernel/agents/open_ai/open_ai_assistant_agent.py b/python/semantic_kernel/agents/open_ai/open_ai_assistant_agent.py index d588ab9b0801..6571fe56c3e3 100644 --- a/python/semantic_kernel/agents/open_ai/open_ai_assistant_agent.py +++ b/python/semantic_kernel/agents/open_ai/open_ai_assistant_agent.py @@ -94,21 +94,18 @@ def __init__( Raises: AgentInitializationError: If the api_key is not provided in the configuration. """ - try: - openai_settings = OpenAISettings.create( - api_key=api_key, - org_id=org_id, - chat_model_id=ai_model_id, - env_file_path=env_file_path, - env_file_encoding=env_file_encoding, - ) - except ValidationError as ex: - raise AgentInitializationError("Failed to create OpenAI settings.", ex) from ex + openai_settings = OpenAIAssistantAgent._create_open_ai_settings( + api_key=api_key, + org_id=org_id, + ai_model_id=ai_model_id, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) if not client and not openai_settings.api_key: raise AgentInitializationError("The OpenAI API key is required, if a client is not provided.") if not openai_settings.chat_model_id: - raise AgentInitializationError("The OpenAI model ID is required.") + raise AgentInitializationError("The OpenAI chat model ID is required.") if not client: client = self._create_client( @@ -271,6 +268,39 @@ def _create_client( default_headers=merged_headers, ) + @staticmethod + def _create_open_ai_settings( + api_key: str | None = None, + org_id: str | None = None, + ai_model_id: str | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> OpenAISettings: + """An internal method to create the OpenAI settings from the provided arguments. + + Args: + api_key: The OpenAI API key. + org_id: The OpenAI organization ID. (optional) + ai_model_id: The AI model ID. (optional) + env_file_path: The environment file path. (optional) + env_file_encoding: The environment file encoding. (optional) + + Returns: + An OpenAI settings instance. + """ + try: + openai_settings = OpenAISettings.create( + api_key=api_key, + org_id=org_id, + chat_model_id=ai_model_id, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + except ValidationError as ex: + raise AgentInitializationError("Failed to create OpenAI settings.", ex) from ex + + return openai_settings + async def list_definitions(self) -> AsyncIterable[dict[str, Any]]: """List the assistant definitions. @@ -281,30 +311,55 @@ async def list_definitions(self) -> AsyncIterable[dict[str, Any]]: for assistant in assistants.data: yield self._create_open_ai_assistant_definition(assistant) + @classmethod async def retrieve( - self, + cls, + *, id: str, - api_key: str, kernel: "Kernel | None" = None, + api_key: str | None = None, org_id: str | None = None, + ai_model_id: str | None = None, + client: AsyncOpenAI | None = None, default_headers: dict[str, str] | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, ) -> "OpenAIAssistantAgent": """Retrieve an assistant by ID. Args: id: The assistant ID. - api_key: The OpenAI API kernel: The Kernel instance. (optional) + api_key: The OpenAI API key. (optional) org_id: The OpenAI organization ID. (optional) + ai_model_id: The AI model ID. (optional) + client: The OpenAI client. (optional) default_headers: The default headers. (optional) - + env_file_path: The environment file path. (optional) + env_file_encoding: The environment file encoding. (optional Returns: An OpenAIAssistantAgent instance. """ - client = self._create_client(api_key=api_key, org_id=org_id, default_headers=default_headers) + openai_settings = OpenAIAssistantAgent._create_open_ai_settings( + api_key=api_key, + org_id=org_id, + ai_model_id=ai_model_id, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + if not client and not openai_settings.api_key: + raise AgentInitializationError("The OpenAI API key is required, if a client is not provided.") + if not openai_settings.chat_model_id: + raise AgentInitializationError("The OpenAI chat model ID is required.") + if not client: + client = OpenAIAssistantAgent._create_client( + api_key=openai_settings.api_key.get_secret_value() if openai_settings.api_key else None, + org_id=openai_settings.org_id, + default_headers=default_headers, + ) assistant = await client.beta.assistants.retrieve(id) - assistant_definition = self._create_open_ai_assistant_definition(assistant) + assistant_definition = OpenAIAssistantBase._create_open_ai_assistant_definition(assistant) return OpenAIAssistantAgent(kernel=kernel, **assistant_definition) # endregion diff --git a/python/semantic_kernel/agents/open_ai/open_ai_assistant_base.py b/python/semantic_kernel/agents/open_ai/open_ai_assistant_base.py index d34ee08ea32b..cd1d7bf3f332 100644 --- a/python/semantic_kernel/agents/open_ai/open_ai_assistant_base.py +++ b/python/semantic_kernel/agents/open_ai/open_ai_assistant_base.py @@ -54,7 +54,7 @@ class OpenAIAssistantBase(Agent): Manages the interaction with OpenAI Assistants. """ - _options_metadata_key: str = "__run_options" + _options_metadata_key: ClassVar[str] = "__run_options" ai_model_id: str client: AsyncOpenAI @@ -284,7 +284,8 @@ async def create_assistant( return self.assistant - def _create_open_ai_assistant_definition(self, assistant: "Assistant") -> dict[str, Any]: + @classmethod + def _create_open_ai_assistant_definition(cls, assistant: "Assistant") -> dict[str, Any]: """Create an OpenAI Assistant Definition from the provided assistant dictionary. Args: @@ -294,11 +295,11 @@ def _create_open_ai_assistant_definition(self, assistant: "Assistant") -> dict[s An OpenAI Assistant Definition. """ execution_settings = {} - if isinstance(assistant.metadata, dict) and self._options_metadata_key in assistant.metadata: - settings_data = assistant.metadata[self._options_metadata_key] + if isinstance(assistant.metadata, dict) and OpenAIAssistantBase._options_metadata_key in assistant.metadata: + settings_data = assistant.metadata[OpenAIAssistantBase._options_metadata_key] if isinstance(settings_data, str): settings_data = json.loads(settings_data) - assistant.metadata[self._options_metadata_key] = settings_data + assistant.metadata[OpenAIAssistantBase._options_metadata_key] = settings_data execution_settings = {key: value for key, value in settings_data.items()} file_ids: list[str] = [] @@ -737,6 +738,7 @@ def _generate_code_interpreter_content(self, agent_name: str, code: str) -> Chat role=AuthorRole.ASSISTANT, content=code, name=agent_name, + metadata={"code": True}, ) def _generate_annotation_content(self, annotation: "Annotation") -> AnnotationContent: diff --git a/python/tests/unit/agents/test_azure_assistant_agent.py b/python/tests/unit/agents/test_azure_assistant_agent.py index 3da6e3ea963a..2878509ea8bc 100644 --- a/python/tests/unit/agents/test_azure_assistant_agent.py +++ b/python/tests/unit/agents/test_azure_assistant_agent.py @@ -9,6 +9,7 @@ from pydantic import ValidationError from semantic_kernel.agents.open_ai.azure_assistant_agent import AzureAssistantAgent +from semantic_kernel.agents.open_ai.open_ai_assistant_base import OpenAIAssistantBase from semantic_kernel.exceptions.agent_exceptions import AgentInitializationError from semantic_kernel.kernel import Kernel @@ -159,35 +160,9 @@ async def test_retrieve_agent(kernel, azure_openai_unit_test_env): mock_client_instance.beta = MagicMock() mock_client_instance.beta.assistants = MagicMock() - mock_assistant = MagicMock() - mock_assistant.metadata = { - "__run_options": { - "max_completion_tokens": 100, - "max_prompt_tokens": 50, - "parallel_tool_calls_enabled": True, - "truncation_message_count": 10, - } - } - mock_assistant.model = "test_model" - mock_assistant.description = "test_description" - mock_assistant.id = "test_id" - mock_assistant.instructions = "test_instructions" - mock_assistant.name = "test_name" - mock_assistant.tools = ["code_interpreter", "file_search"] - mock_assistant.temperature = 0.7 - mock_assistant.top_p = 0.9 - mock_assistant.response_format = {"type": "json_object"} - mock_assistant.tool_resources = { - "code_interpreter": {"file_ids": ["file1", "file2"]}, - "file_search": {"vector_store_ids": ["vector_store1"]}, - } - - mock_client_instance.beta.assistants.retrieve = AsyncMock(return_value=mock_assistant) + mock_client_instance.beta.assistants.retrieve = AsyncMock(return_value=AsyncMock()) - agent = AzureAssistantAgent( - kernel=kernel, service_id="test_service", name="test_name", instructions="test_instructions", id="test_id" - ) - agent._create_open_ai_assistant_definition = MagicMock( + OpenAIAssistantBase._create_open_ai_assistant_definition = MagicMock( return_value={ "ai_model_id": "test_model", "description": "test_description", @@ -216,7 +191,7 @@ async def test_retrieve_agent(kernel, azure_openai_unit_test_env): } ) - retrieved_agent = await agent.retrieve("test_id", "test_api_key", kernel) + retrieved_agent = await AzureAssistantAgent.retrieve(id="test_id", api_key="test_api_key", kernel=kernel) assert retrieved_agent.model_dump( include={ "ai_model_id", @@ -264,7 +239,23 @@ async def test_retrieve_agent(kernel, azure_openai_unit_test_env): "truncation_message_count": 10, } mock_client_instance.beta.assistants.retrieve.assert_called_once_with("test_id") - agent._create_open_ai_assistant_definition.assert_called_once_with(mock_assistant) + OpenAIAssistantBase._create_open_ai_assistant_definition.assert_called_once() + + +@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True) +@pytest.mark.asyncio +async def test_retrieve_agent_missing_chat_deployment_name_throws(kernel, azure_openai_unit_test_env): + with pytest.raises(AgentInitializationError, match="The Azure OpenAI chat_deployment_name is required."): + _ = await AzureAssistantAgent.retrieve( + id="test_id", api_key="test_api_key", kernel=kernel, env_file_path="test.env" + ) + + +@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_API_KEY"]], indirect=True) +@pytest.mark.asyncio +async def test_retrieve_agent_missing_api_key_throws(kernel, azure_openai_unit_test_env): + with pytest.raises(AgentInitializationError, match="Please provide either api_key, ad_token or ad_token_provider."): + _ = await AzureAssistantAgent.retrieve(id="test_id", kernel=kernel, env_file_path="test.env") def test_open_ai_settings_create_throws(azure_openai_unit_test_env): diff --git a/python/tests/unit/agents/test_open_ai_assistant_agent.py b/python/tests/unit/agents/test_open_ai_assistant_agent.py index 0005f9ec76ab..1dd53711674c 100644 --- a/python/tests/unit/agents/test_open_ai_assistant_agent.py +++ b/python/tests/unit/agents/test_open_ai_assistant_agent.py @@ -14,6 +14,7 @@ from pydantic import ValidationError from semantic_kernel.agents.open_ai.open_ai_assistant_agent import OpenAIAssistantAgent +from semantic_kernel.agents.open_ai.open_ai_assistant_base import OpenAIAssistantBase from semantic_kernel.exceptions.agent_exceptions import AgentInitializationError from semantic_kernel.kernel import Kernel @@ -228,35 +229,12 @@ async def test_retrieve_agent(kernel, openai_unit_test_env): mock_client_instance.beta = MagicMock() mock_client_instance.beta.assistants = MagicMock() - mock_assistant = MagicMock() - mock_assistant.metadata = { - "__run_options": { - "max_completion_tokens": 100, - "max_prompt_tokens": 50, - "parallel_tool_calls_enabled": True, - "truncation_message_count": 10, - } - } - mock_assistant.model = "test_model" - mock_assistant.description = "test_description" - mock_assistant.id = "test_id" - mock_assistant.instructions = "test_instructions" - mock_assistant.name = "test_name" - mock_assistant.tools = ["code_interpreter", "file_search"] - mock_assistant.temperature = 0.7 - mock_assistant.top_p = 0.9 - mock_assistant.response_format = {"type": "json_object"} - mock_assistant.tool_resources = { - "code_interpreter": {"file_ids": ["file1", "file2"]}, - "file_search": {"vector_store_ids": ["vector_store1"]}, - } - - mock_client_instance.beta.assistants.retrieve = AsyncMock(return_value=mock_assistant) + mock_client_instance.beta.assistants.retrieve = AsyncMock(return_value=AsyncMock()) agent = OpenAIAssistantAgent( kernel=kernel, service_id="test_service", name="test_name", instructions="test_instructions", id="test_id" ) - agent._create_open_ai_assistant_definition = MagicMock( + OpenAIAssistantBase._create_open_ai_assistant_definition = MagicMock( return_value={ "ai_model_id": "test_model", "description": "test_description", @@ -285,7 +263,7 @@ async def test_retrieve_agent(kernel, openai_unit_test_env): } ) - retrieved_agent = await agent.retrieve("test_id", "test_api_key", kernel) + retrieved_agent = await agent.retrieve(id="test_id", api_key="test_api_key", kernel=kernel) assert retrieved_agent.model_dump( include={ "ai_model_id", @@ -333,7 +311,23 @@ async def test_retrieve_agent(kernel, openai_unit_test_env): "truncation_message_count": 10, } mock_client_instance.beta.assistants.retrieve.assert_called_once_with("test_id") - agent._create_open_ai_assistant_definition.assert_called_once_with(mock_assistant) + OpenAIAssistantBase._create_open_ai_assistant_definition.assert_called_once() + + +@pytest.mark.parametrize("exclude_list", [["OPENAI_CHAT_MODEL_ID"]], indirect=True) +@pytest.mark.asyncio +async def test_retrieve_agent_missing_chat_model_id_throws(kernel, openai_unit_test_env): + with pytest.raises(AgentInitializationError, match="The OpenAI chat model ID is required."): + _ = await OpenAIAssistantAgent.retrieve( + id="test_id", api_key="test_api_key", kernel=kernel, env_file_path="test.env" + ) + + +@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True) +@pytest.mark.asyncio +async def test_retrieve_agent_missing_api_key_throws(kernel, openai_unit_test_env): + with pytest.raises(AgentInitializationError, match="The OpenAI API key is required, if a client is not provided."): + _ = await OpenAIAssistantAgent.retrieve(id="test_id", kernel=kernel, env_file_path="test.env") def test_open_ai_settings_create_throws(openai_unit_test_env): @@ -348,14 +342,14 @@ def test_open_ai_settings_create_throws(openai_unit_test_env): @pytest.mark.parametrize("exclude_list", [["OPENAI_CHAT_MODEL_ID"]], indirect=True) def test_azure_openai_agent_create_missing_chat_model_id_throws(openai_unit_test_env): - with pytest.raises(AgentInitializationError, match="The OpenAI model ID is required."): + with pytest.raises(AgentInitializationError, match="The OpenAI chat model ID is required."): OpenAIAssistantAgent(service_id="test_service", env_file_path="test.env") @pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True) def test_azure_openai_agent_create_missing_api_key_throws(openai_unit_test_env): with pytest.raises(AgentInitializationError, match="The OpenAI API key is required, if a client is not provided."): - OpenAIAssistantAgent(service_id="test_service", env_file_path="test.env") + OpenAIAssistantAgent(env_file_path="test.env") def test_create_open_ai_assistant_definition(mock_assistant, openai_unit_test_env):