mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: OpenAI Assistants Chat Client and Agent (#288)
* Initial version of assistant client
* More updates to assistant client
* Finished assistant chat client implementation
* Small fixes and basic example
* Added code interpreter example
* More examples
* Added chat client example
* Small fixes
* Added tests
* Enabled telemetry
* Small fix
* Removed files temporarily
* Revert "Removed files temporarily"
This reverts commit 5cdfa0d299.
* Small fixes
* Addressed PR feedback
* Fixed tests
* Small update
This commit is contained in:
Dmytro Struk
committed by
GitHub
Unverified
GitHub
parent
624709e5d1
commit
75794b4687
@@ -130,7 +130,7 @@ class FoundryChatClient(ChatClientBase):
|
||||
nor agent_id is provided, both will be created and managed automatically.
|
||||
agent_name: The name to use when creating new agents.
|
||||
thread_id: Default thread ID to use for conversations. Can be overridden by
|
||||
conversation_id property from ChatOptions, when making a request.
|
||||
conversation_id property, when making a request.
|
||||
project_endpoint: The Azure AI Foundry project endpoint URL. Used if client is not provided.
|
||||
model_deployment_name: The model deployment name to use for agent creation.
|
||||
credential: Azure async credential to use for authentication. If not provided,
|
||||
@@ -289,6 +289,7 @@ class FoundryChatClient(ChatClientBase):
|
||||
# Get any active run for this thread
|
||||
thread_run = await self._get_active_thread_run(thread_id)
|
||||
|
||||
stream: AsyncAgentRunStream[AsyncAgentEventHandler[Any]] | AsyncAgentEventHandler[Any]
|
||||
handler: AsyncAgentEventHandler[Any] = AsyncAgentEventHandler()
|
||||
tool_run_id, tool_outputs = self._convert_function_results_to_tool_output(tool_results)
|
||||
|
||||
@@ -354,28 +355,21 @@ class FoundryChatClient(ChatClientBase):
|
||||
thread_id: str,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
"""Process events from the agent stream and yield ChatResponseUpdate objects."""
|
||||
response_id: str | None = None
|
||||
|
||||
if stream is not None:
|
||||
# Use 'async with' only if the stream supports async context management (main agent stream).
|
||||
# Tool output handlers only support async iteration, not context management.
|
||||
if isinstance(stream, contextlib.AbstractAsyncContextManager):
|
||||
async with stream as response_stream: # type: ignore
|
||||
async for update in self._process_stream_events_from_iterator(
|
||||
response_stream, thread_id, response_id
|
||||
):
|
||||
yield update
|
||||
else:
|
||||
async for update in self._process_stream_events_from_iterator(stream, thread_id, response_id):
|
||||
# Use 'async with' only if the stream supports async context management (main agent stream).
|
||||
# Tool output handlers only support async iteration, not context management.
|
||||
if isinstance(stream, contextlib.AbstractAsyncContextManager):
|
||||
async with stream as response_stream: # type: ignore
|
||||
async for update in self._process_stream_events_from_iterator(response_stream, thread_id):
|
||||
yield update
|
||||
else:
|
||||
async for update in self._process_stream_events_from_iterator(stream, thread_id):
|
||||
yield update
|
||||
|
||||
async def _process_stream_events_from_iterator(
|
||||
self,
|
||||
stream_iter: Any,
|
||||
thread_id: str,
|
||||
response_id: str | None,
|
||||
self, stream_iter: Any, thread_id: str
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
"""Process events from the stream iterator and yield ChatResponseUpdate objects."""
|
||||
response_id: str | None = None
|
||||
async for event_type, event_data, _ in stream_iter: # type: ignore
|
||||
if event_type == AgentStreamEvent.THREAD_RUN_CREATED and isinstance(event_data, ThreadRun):
|
||||
yield ChatResponseUpdate(
|
||||
@@ -496,7 +490,7 @@ class FoundryChatClient(ChatClientBase):
|
||||
if chat_options.tool_choice != "none" and chat_options.tools is not None:
|
||||
for tool in chat_options.tools:
|
||||
if isinstance(tool, AIFunction):
|
||||
tool_definitions.append(ai_function_to_json_schema_spec(tool))
|
||||
tool_definitions.append(ai_function_to_json_schema_spec(tool)) # type: ignore[reportUnknownArgumentType]
|
||||
elif isinstance(tool, HostedCodeInterpreterTool):
|
||||
tool_definitions.append(CodeInterpreterToolDefinition())
|
||||
elif isinstance(tool, MutableMapping):
|
||||
|
||||
@@ -45,6 +45,7 @@ def create_test_foundry_chat_client(
|
||||
thread_id=thread_id,
|
||||
_should_delete_agent=should_delete_agent,
|
||||
_foundry_settings=foundry_settings,
|
||||
credential=None,
|
||||
)
|
||||
|
||||
|
||||
@@ -95,6 +96,7 @@ def test_foundry_chat_client_init_auto_create_client(
|
||||
thread_id=None,
|
||||
_should_delete_agent=False,
|
||||
_foundry_settings=foundry_settings,
|
||||
credential=None,
|
||||
)
|
||||
|
||||
assert chat_client.client is mock_ai_project_client
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from ._assistants_client import * # noqa: F403
|
||||
from ._chat_client import * # noqa: F403
|
||||
from ._exceptions import * # noqa: F403
|
||||
from ._responses_client import * # noqa: F403
|
||||
|
||||
@@ -0,0 +1,469 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
import sys
|
||||
from collections.abc import AsyncIterable, Mapping, MutableMapping, MutableSequence
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.beta.threads import (
|
||||
ImageURLContentBlockParam,
|
||||
ImageURLParam,
|
||||
MessageContentPartParam,
|
||||
MessageDeltaEvent,
|
||||
Run,
|
||||
TextContentBlockParam,
|
||||
TextDeltaBlock,
|
||||
)
|
||||
from openai.types.beta.threads.run_create_params import AdditionalMessage
|
||||
from openai.types.beta.threads.run_submit_tool_outputs_params import ToolOutput
|
||||
from openai.types.beta.threads.runs import RunStep
|
||||
from pydantic import Field, PrivateAttr, SecretStr, ValidationError
|
||||
|
||||
from .._clients import ChatClientBase, ai_function_to_json_schema_spec, use_tool_calling
|
||||
from .._tools import AIFunction, HostedCodeInterpreterTool
|
||||
from .._types import (
|
||||
AIContents,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
ChatRole,
|
||||
ChatToolMode,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
TextContent,
|
||||
UriContent,
|
||||
UsageContent,
|
||||
UsageDetails,
|
||||
)
|
||||
from ..exceptions import ServiceInitializationError
|
||||
from ..telemetry import use_telemetry
|
||||
from ._shared import OpenAIConfigBase, OpenAISettings
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
from typing import Self # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import Self # pragma: no cover
|
||||
|
||||
__all__ = ["OpenAIAssistantsClient"]
|
||||
|
||||
|
||||
@use_telemetry
|
||||
@use_tool_calling
|
||||
class OpenAIAssistantsClient(OpenAIConfigBase, ChatClientBase):
|
||||
"""OpenAI Assistants client."""
|
||||
|
||||
MODEL_PROVIDER_NAME: ClassVar[str] = "openai" # type: ignore[reportIncompatibleVariableOverride, misc]
|
||||
assistant_id: str | None = Field(default=None)
|
||||
assistant_name: str | None = Field(default=None)
|
||||
thread_id: str | None = Field(default=None)
|
||||
_should_delete_assistant: bool = PrivateAttr(default=False) # Track whether we should delete the assistant
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ai_model_id: str | None = None,
|
||||
assistant_id: str | None = None,
|
||||
assistant_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
api_key: str | None = None,
|
||||
org_id: str | None = None,
|
||||
default_headers: Mapping[str, str] | None = None,
|
||||
async_client: AsyncOpenAI | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize an OpenAI Assistants client.
|
||||
|
||||
Args:
|
||||
ai_model_id: OpenAI model name, see
|
||||
https://platform.openai.com/docs/models
|
||||
assistant_id: The ID of an OpenAI assistant to use.
|
||||
If not provided, a new assistant will be created (and deleted after the request).
|
||||
assistant_name: The name to use when creating new assistants.
|
||||
thread_id: Default thread ID to use for conversations. Can be overridden by
|
||||
conversation_id property, when making a request.
|
||||
If not provided, a new thread will be created (and deleted after the request).
|
||||
api_key: The optional API key to use. If provided will override,
|
||||
the env vars or .env file value.
|
||||
org_id: The optional org ID to use. If provided will override,
|
||||
the env vars or .env file value.
|
||||
default_headers: The default headers mapping of string keys to
|
||||
string values for HTTP requests. (Optional)
|
||||
async_client: An existing client to use. (Optional)
|
||||
env_file_path: Use the environment settings file as a fallback
|
||||
to environment variables. (Optional)
|
||||
env_file_encoding: The encoding of the environment settings file. (Optional)
|
||||
"""
|
||||
try:
|
||||
openai_settings = OpenAISettings(
|
||||
api_key=SecretStr(api_key) if api_key else None,
|
||||
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 ServiceInitializationError("Failed to create OpenAI settings.", ex) from ex
|
||||
|
||||
if not async_client and not openai_settings.api_key:
|
||||
raise ServiceInitializationError("The OpenAI API key is required.")
|
||||
if not openai_settings.chat_model_id:
|
||||
raise ServiceInitializationError("The OpenAI model ID is required.")
|
||||
|
||||
super().__init__(
|
||||
ai_model_id=openai_settings.chat_model_id,
|
||||
assistant_id=assistant_id, # type: ignore[reportCallIssue]
|
||||
assistant_name=assistant_name, # type: ignore[reportCallIssue]
|
||||
thread_id=thread_id, # type: ignore[reportCallIssue]
|
||||
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,
|
||||
client=async_client,
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> "Self":
|
||||
"""Async context manager entry."""
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
|
||||
"""Async context manager exit - clean up any assistants we created."""
|
||||
await self.close()
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Clean up any assistants we created."""
|
||||
if self._should_delete_assistant and self.assistant_id is not None:
|
||||
await self.client.beta.assistants.delete(self.assistant_id)
|
||||
self.assistant_id = None
|
||||
self._should_delete_assistant = False
|
||||
|
||||
async def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
return await ChatResponse.from_chat_response_generator(
|
||||
updates=self._inner_get_streaming_response(messages=messages, chat_options=chat_options, **kwargs)
|
||||
)
|
||||
|
||||
async def _inner_get_streaming_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
# Extract necessary state from messages and options
|
||||
run_options, tool_results = self._create_run_options(messages, chat_options, **kwargs)
|
||||
|
||||
# Get the thread ID
|
||||
thread_id: str | None = (
|
||||
chat_options.conversation_id if chat_options.conversation_id is not None else self.thread_id
|
||||
)
|
||||
|
||||
if thread_id is None and tool_results is not None:
|
||||
raise ValueError("No thread ID was provided, but chat messages includes tool results.")
|
||||
|
||||
# Determine which assistant to use and create if needed
|
||||
assistant_id = await self._get_assistant_id_or_create()
|
||||
|
||||
# Create the streaming response
|
||||
stream, thread_id = await self._create_assistant_stream(thread_id, assistant_id, run_options, tool_results)
|
||||
|
||||
# Process and yield each update from the stream
|
||||
async for update in self._process_stream_events(stream, thread_id):
|
||||
yield update
|
||||
|
||||
async def _get_assistant_id_or_create(self) -> str:
|
||||
"""Determine which assistant to use and create if needed.
|
||||
|
||||
Returns:
|
||||
str: The assistant_id to use.
|
||||
"""
|
||||
# If no assistant is provided, create a temporary assistant
|
||||
if self.assistant_id is None:
|
||||
created_assistant = await self.client.beta.assistants.create(
|
||||
name=self.assistant_name, model=self.ai_model_id
|
||||
)
|
||||
|
||||
self.assistant_id = created_assistant.id
|
||||
self._should_delete_assistant = True
|
||||
|
||||
return self.assistant_id
|
||||
|
||||
async def _create_assistant_stream(
|
||||
self,
|
||||
thread_id: str | None,
|
||||
assistant_id: str,
|
||||
run_options: dict[str, Any],
|
||||
tool_results: list[FunctionResultContent] | None,
|
||||
) -> tuple[Any, str]:
|
||||
"""Create the assistant stream for processing.
|
||||
|
||||
Returns:
|
||||
tuple: (stream, final_thread_id)
|
||||
"""
|
||||
# Get any active run for this thread
|
||||
thread_run = await self._get_active_thread_run(thread_id)
|
||||
|
||||
tool_run_id, tool_outputs = self._convert_function_results_to_tool_output(tool_results)
|
||||
|
||||
if thread_run is not None and tool_run_id is not None and tool_run_id == thread_run.id and tool_outputs:
|
||||
# There's an active run and we have tool results to submit, so submit the results.
|
||||
stream = self.client.beta.threads.runs.submit_tool_outputs_stream( # type: ignore[reportDeprecated]
|
||||
run_id=tool_run_id, thread_id=thread_run.thread_id, tool_outputs=tool_outputs
|
||||
)
|
||||
final_thread_id = thread_run.thread_id
|
||||
else:
|
||||
# Handle thread creation or cancellation
|
||||
final_thread_id = await self._prepare_thread(thread_id, thread_run, run_options)
|
||||
|
||||
# Now create a new run and stream the results.
|
||||
stream = self.client.beta.threads.runs.stream( # type: ignore[reportDeprecated]
|
||||
assistant_id=assistant_id, thread_id=final_thread_id, **run_options
|
||||
)
|
||||
|
||||
return stream, final_thread_id
|
||||
|
||||
async def _get_active_thread_run(self, thread_id: str | None) -> Run | None:
|
||||
"""Get any active run for the given thread."""
|
||||
if thread_id is None:
|
||||
return None
|
||||
|
||||
async for run in self.client.beta.threads.runs.list(thread_id=thread_id, limit=1, order="desc"): # type: ignore[reportDeprecated]
|
||||
if run.status not in ["completed", "cancelled", "failed", "expired"]:
|
||||
return run
|
||||
return None
|
||||
|
||||
async def _prepare_thread(self, thread_id: str | None, thread_run: Run | None, run_options: dict[str, Any]) -> str:
|
||||
"""Prepare the thread for a new run, creating or cleaning up as needed."""
|
||||
if thread_id is None:
|
||||
# No thread ID was provided, so create a new thread.
|
||||
thread = await self.client.beta.threads.create( # type: ignore[reportDeprecated]
|
||||
messages=run_options["additional_messages"],
|
||||
tool_resources=run_options.get("tool_resources"),
|
||||
metadata=run_options.get("metadata"),
|
||||
)
|
||||
run_options["additional_messages"] = []
|
||||
return thread.id
|
||||
|
||||
if thread_run is not None:
|
||||
# There was an active run; we need to cancel it before starting a new run.
|
||||
await self.client.beta.threads.runs.cancel(run_id=thread_run.id, thread_id=thread_id) # type: ignore[reportDeprecated]
|
||||
|
||||
return thread_id
|
||||
|
||||
async def _process_stream_events(self, stream: Any, thread_id: str) -> AsyncIterable[ChatResponseUpdate]:
|
||||
response_id: str | None = None
|
||||
|
||||
async with stream as response_stream:
|
||||
async for response in response_stream:
|
||||
if response.event == "thread.run.created":
|
||||
yield ChatResponseUpdate(
|
||||
contents=[],
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
role=ChatRole.ASSISTANT,
|
||||
)
|
||||
elif response.event == "thread.run.step.created" and isinstance(response.data, RunStep):
|
||||
response_id = response.data.run_id
|
||||
elif response.event == "thread.message.delta" and isinstance(response.data, MessageDeltaEvent):
|
||||
delta = response.data.delta
|
||||
role = ChatRole.USER if delta.role == "user" else ChatRole.ASSISTANT
|
||||
|
||||
for delta_block in delta.content or []:
|
||||
if isinstance(delta_block, TextDeltaBlock) and delta_block.text and delta_block.text.value:
|
||||
yield ChatResponseUpdate(
|
||||
role=role,
|
||||
text=delta_block.text.value,
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
)
|
||||
elif response.event == "thread.run.requires_action" and isinstance(response.data, Run):
|
||||
contents = self._create_function_call_contents(response.data, response_id)
|
||||
if contents:
|
||||
yield ChatResponseUpdate(
|
||||
role=ChatRole.ASSISTANT,
|
||||
contents=contents,
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
)
|
||||
elif (
|
||||
response.event == "thread.run.completed"
|
||||
and isinstance(response.data, Run)
|
||||
and response.data.usage is not None
|
||||
):
|
||||
usage = response.data.usage
|
||||
usage_content = UsageContent(
|
||||
UsageDetails(
|
||||
input_token_count=usage.prompt_tokens,
|
||||
output_token_count=usage.completion_tokens,
|
||||
total_token_count=usage.total_tokens,
|
||||
)
|
||||
)
|
||||
yield ChatResponseUpdate(
|
||||
role=ChatRole.ASSISTANT,
|
||||
contents=[usage_content],
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
)
|
||||
else:
|
||||
yield ChatResponseUpdate(
|
||||
contents=[],
|
||||
conversation_id=thread_id,
|
||||
message_id=response_id,
|
||||
raw_representation=response.data,
|
||||
response_id=response_id,
|
||||
role=ChatRole.ASSISTANT,
|
||||
)
|
||||
|
||||
def _create_function_call_contents(self, event_data: Run, response_id: str | None) -> list[AIContents]:
|
||||
"""Create function call contents from a tool action event."""
|
||||
contents: list[AIContents] = []
|
||||
|
||||
if event_data.required_action is not None:
|
||||
for tool_call in event_data.required_action.submit_tool_outputs.tool_calls:
|
||||
call_id = json.dumps([response_id, tool_call.id])
|
||||
function_name = tool_call.function.name
|
||||
function_arguments = json.loads(tool_call.function.arguments)
|
||||
contents.append(FunctionCallContent(call_id=call_id, name=function_name, arguments=function_arguments))
|
||||
|
||||
return contents
|
||||
|
||||
def _create_run_options(
|
||||
self,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions | None,
|
||||
**kwargs: Any,
|
||||
) -> tuple[dict[str, Any], list[FunctionResultContent] | None]:
|
||||
run_options: dict[str, Any] = {**kwargs}
|
||||
|
||||
if chat_options is not None:
|
||||
run_options["max_completion_tokens"] = chat_options.max_tokens
|
||||
run_options["model"] = chat_options.ai_model_id
|
||||
run_options["top_p"] = chat_options.top_p
|
||||
run_options["temperature"] = chat_options.temperature
|
||||
|
||||
if chat_options.allow_multiple_tool_calls is not None:
|
||||
run_options["parallel_tool_calls"] = chat_options.allow_multiple_tool_calls
|
||||
|
||||
if chat_options.tool_choice is not None:
|
||||
tool_definitions: list[MutableMapping[str, Any]] = []
|
||||
if chat_options.tool_choice != "none" and chat_options.tools is not None:
|
||||
for tool in chat_options.tools:
|
||||
if isinstance(tool, AIFunction):
|
||||
tool_definitions.append(ai_function_to_json_schema_spec(tool)) # type: ignore[reportUnknownArgumentType]
|
||||
elif isinstance(tool, HostedCodeInterpreterTool):
|
||||
tool_definitions.append({"type": "code_interpreter"})
|
||||
elif isinstance(tool, MutableMapping):
|
||||
tool_definitions.append(tool)
|
||||
|
||||
if len(tool_definitions) > 0:
|
||||
run_options["tools"] = tool_definitions
|
||||
|
||||
if chat_options.tool_choice == "none" or chat_options.tool_choice == "auto":
|
||||
run_options["tool_choice"] = chat_options.tool_choice
|
||||
elif (
|
||||
isinstance(chat_options.tool_choice, ChatToolMode)
|
||||
and chat_options.tool_choice == "required"
|
||||
and chat_options.tool_choice.required_function_name is not None
|
||||
):
|
||||
run_options["tool_choice"] = {
|
||||
"type": "function",
|
||||
"function": {"name": chat_options.tool_choice.required_function_name},
|
||||
}
|
||||
|
||||
if chat_options.response_format is not None:
|
||||
run_options["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": chat_options.response_format.model_json_schema(),
|
||||
}
|
||||
|
||||
instructions: list[str] = []
|
||||
tool_results: list[FunctionResultContent] | None = None
|
||||
|
||||
additional_messages: list[AdditionalMessage] | None = None
|
||||
|
||||
# System/developer messages are turned into instructions,
|
||||
# since there is no such message roles in OpenAI Assistants.
|
||||
# All other messages are added 1:1.
|
||||
for chat_message in messages:
|
||||
if chat_message.role.value in ["system", "developer"]:
|
||||
for text_content in [content for content in chat_message.contents if isinstance(content, TextContent)]:
|
||||
instructions.append(text_content.text)
|
||||
|
||||
continue
|
||||
|
||||
message_contents: list[MessageContentPartParam] = []
|
||||
|
||||
for content in chat_message.contents:
|
||||
if isinstance(content, TextContent):
|
||||
message_contents.append(TextContentBlockParam(type="text", text=content.text))
|
||||
elif isinstance(content, UriContent) and content.has_top_level_media_type("image"):
|
||||
message_contents.append(
|
||||
ImageURLContentBlockParam(type="image_url", image_url=ImageURLParam(url=content.uri))
|
||||
)
|
||||
elif isinstance(content, FunctionResultContent):
|
||||
if tool_results is None:
|
||||
tool_results = []
|
||||
tool_results.append(content)
|
||||
|
||||
if len(message_contents) > 0:
|
||||
if additional_messages is None:
|
||||
additional_messages = []
|
||||
additional_messages.append(
|
||||
AdditionalMessage(
|
||||
role="assistant" if chat_message.role == ChatRole.ASSISTANT else "user",
|
||||
content=message_contents,
|
||||
)
|
||||
)
|
||||
|
||||
if additional_messages is not None:
|
||||
run_options["additional_messages"] = additional_messages
|
||||
|
||||
if len(instructions) > 0:
|
||||
run_options["instructions"] = "".join(instructions)
|
||||
|
||||
return run_options, tool_results
|
||||
|
||||
def _convert_function_results_to_tool_output(
|
||||
self,
|
||||
tool_results: list[FunctionResultContent] | None,
|
||||
) -> tuple[str | None, list[ToolOutput] | None]:
|
||||
run_id: str | None = None
|
||||
tool_outputs: list[ToolOutput] | None = None
|
||||
|
||||
if tool_results:
|
||||
for function_result_content in tool_results:
|
||||
# When creating the FunctionCallContent, we created it with a CallId == [runId, callId].
|
||||
# We need to extract the run ID and ensure that the ToolOutput we send back to Azure
|
||||
# is only the call ID.
|
||||
run_and_call_ids: list[str] = json.loads(function_result_content.call_id)
|
||||
|
||||
if (
|
||||
not run_and_call_ids
|
||||
or len(run_and_call_ids) != 2
|
||||
or not run_and_call_ids[0]
|
||||
or not run_and_call_ids[1]
|
||||
or (run_id is not None and run_id != run_and_call_ids[0])
|
||||
):
|
||||
continue
|
||||
|
||||
run_id = run_and_call_ids[0]
|
||||
call_id = run_and_call_ids[1]
|
||||
|
||||
if tool_outputs is None:
|
||||
tool_outputs = []
|
||||
tool_outputs.append(ToolOutput(tool_call_id=call_id, output=str(function_result_content.result)))
|
||||
|
||||
return run_id, tool_outputs
|
||||
@@ -0,0 +1,388 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import os
|
||||
from typing import Annotated
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from pydantic import Field
|
||||
|
||||
from agent_framework import (
|
||||
ChatClient,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
from agent_framework.openai import OpenAIAssistantsClient
|
||||
|
||||
skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
|
||||
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
|
||||
or os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"),
|
||||
reason="No real OPENAI_API_KEY provided; skipping integration tests."
|
||||
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
|
||||
else "Integration tests are disabled.",
|
||||
)
|
||||
|
||||
|
||||
def create_test_openai_assistants_client(
|
||||
mock_async_openai: MagicMock,
|
||||
ai_model_id: str | None = None,
|
||||
assistant_id: str | None = None,
|
||||
assistant_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
should_delete_assistant: bool = False,
|
||||
) -> OpenAIAssistantsClient:
|
||||
"""Helper function to create OpenAIAssistantsClient instances for testing, bypassing Pydantic validation."""
|
||||
return OpenAIAssistantsClient.model_construct(
|
||||
ai_model_id=ai_model_id or "gpt-4",
|
||||
assistant_id=assistant_id,
|
||||
assistant_name=assistant_name,
|
||||
thread_id=thread_id,
|
||||
api_key="test-api-key",
|
||||
org_id="test-org-id",
|
||||
client=mock_async_openai,
|
||||
_should_delete_assistant=should_delete_assistant,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_async_openai() -> MagicMock:
|
||||
"""Mock AsyncOpenAI client."""
|
||||
mock_client = MagicMock()
|
||||
|
||||
# Mock beta.assistants
|
||||
mock_client.beta.assistants.create = AsyncMock(return_value=MagicMock(id="test-assistant-id"))
|
||||
mock_client.beta.assistants.delete = AsyncMock()
|
||||
|
||||
# Mock beta.threads
|
||||
mock_client.beta.threads.create = AsyncMock(return_value=MagicMock(id="test-thread-id"))
|
||||
mock_client.beta.threads.delete = AsyncMock()
|
||||
|
||||
# Mock beta.threads.runs
|
||||
mock_client.beta.threads.runs.create = AsyncMock(return_value=MagicMock(id="test-run-id"))
|
||||
mock_client.beta.threads.runs.retrieve = AsyncMock()
|
||||
mock_client.beta.threads.runs.submit_tool_outputs = AsyncMock()
|
||||
|
||||
# Mock beta.threads.messages
|
||||
mock_client.beta.threads.messages.create = AsyncMock()
|
||||
mock_client.beta.threads.messages.list = AsyncMock(return_value=MagicMock(data=[]))
|
||||
|
||||
return mock_client
|
||||
|
||||
|
||||
def test_openai_assistants_client_init_with_client(mock_async_openai: MagicMock) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with existing client."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, ai_model_id="gpt-4", assistant_id="existing-assistant-id", thread_id="test-thread-id"
|
||||
)
|
||||
|
||||
assert chat_client.client is mock_async_openai
|
||||
assert chat_client.ai_model_id == "gpt-4"
|
||||
assert chat_client.assistant_id == "existing-assistant-id"
|
||||
assert chat_client.thread_id == "test-thread-id"
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
assert isinstance(chat_client, ChatClient)
|
||||
|
||||
|
||||
def test_openai_assistants_client_init_auto_create_client(
|
||||
openai_unit_test_env: dict[str, str],
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with auto-created client."""
|
||||
chat_client = OpenAIAssistantsClient.model_construct(
|
||||
ai_model_id=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
assistant_id=None,
|
||||
assistant_name="TestAssistant",
|
||||
thread_id=None,
|
||||
api_key=openai_unit_test_env["OPENAI_API_KEY"],
|
||||
org_id=openai_unit_test_env["OPENAI_ORG_ID"],
|
||||
client=mock_async_openai,
|
||||
_should_delete_assistant=False,
|
||||
)
|
||||
|
||||
assert chat_client.client is mock_async_openai
|
||||
assert chat_client.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert chat_client.assistant_id is None
|
||||
assert chat_client.assistant_name == "TestAssistant"
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
|
||||
|
||||
def test_openai_assistants_client_init_validation_fail() -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with validation failure."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
# Force failure by providing invalid model ID type - this should cause validation to fail
|
||||
OpenAIAssistantsClient(ai_model_id=123, api_key="valid-key") # type: ignore
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_openai_assistants_client_init_missing_model_id(openai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with missing model ID."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIAssistantsClient(
|
||||
api_key=openai_unit_test_env.get("OPENAI_API_KEY", "test-key"), env_file_path="nonexistent.env"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
|
||||
def test_openai_assistants_client_init_missing_api_key(openai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with missing API key."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIAssistantsClient(ai_model_id="gpt-4", env_file_path="nonexistent.env")
|
||||
|
||||
|
||||
def test_openai_assistants_client_init_with_default_headers(openai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with default headers."""
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
chat_client = OpenAIAssistantsClient(
|
||||
ai_model_id="gpt-4",
|
||||
api_key=openai_unit_test_env["OPENAI_API_KEY"],
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
assert chat_client.ai_model_id == "gpt-4"
|
||||
assert isinstance(chat_client, ChatClient)
|
||||
|
||||
# Assert that the default header we added is present in the client's default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in chat_client.client.default_headers
|
||||
assert chat_client.client.default_headers[key] == value
|
||||
|
||||
|
||||
async def test_openai_assistants_client_get_assistant_id_or_create_existing_assistant(
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test _get_assistant_id_or_create when assistant_id is already provided."""
|
||||
chat_client = create_test_openai_assistants_client(mock_async_openai, assistant_id="existing-assistant-id")
|
||||
|
||||
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
|
||||
|
||||
assert assistant_id == "existing-assistant-id"
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
mock_async_openai.beta.assistants.create.assert_not_called()
|
||||
|
||||
|
||||
async def test_openai_assistants_client_get_assistant_id_or_create_create_new(
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test _get_assistant_id_or_create when creating a new assistant."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, ai_model_id="gpt-4", assistant_name="TestAssistant"
|
||||
)
|
||||
|
||||
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
|
||||
|
||||
assert assistant_id == "test-assistant-id"
|
||||
assert chat_client._should_delete_assistant # type: ignore
|
||||
mock_async_openai.beta.assistants.create.assert_called_once()
|
||||
|
||||
|
||||
async def test_openai_assistants_client_aclose_should_not_delete(
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test close when assistant should not be deleted."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, assistant_id="assistant-to-keep", should_delete_assistant=False
|
||||
)
|
||||
|
||||
await chat_client.close() # type: ignore
|
||||
|
||||
# Verify assistant deletion was not called
|
||||
mock_async_openai.beta.assistants.delete.assert_not_called()
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
|
||||
|
||||
async def test_openai_assistants_client_aclose_should_delete(mock_async_openai: MagicMock) -> None:
|
||||
"""Test close method calls cleanup."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
|
||||
)
|
||||
|
||||
await chat_client.close()
|
||||
|
||||
# Verify assistant deletion was called
|
||||
mock_async_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
|
||||
|
||||
async def test_openai_assistants_client_async_context_manager(mock_async_openai: MagicMock) -> None:
|
||||
"""Test async context manager functionality."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
|
||||
)
|
||||
|
||||
# Test context manager
|
||||
async with chat_client:
|
||||
pass # Just test that we can enter and exit
|
||||
|
||||
# Verify cleanup was called on exit
|
||||
mock_async_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
|
||||
|
||||
|
||||
def test_openai_assistants_client_serialize(openai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test serialization of OpenAIAssistantsClient."""
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
# Test basic initialization and to_dict
|
||||
chat_client = OpenAIAssistantsClient(
|
||||
ai_model_id="gpt-4",
|
||||
assistant_id="test-assistant-id",
|
||||
assistant_name="TestAssistant",
|
||||
thread_id="test-thread-id",
|
||||
api_key=openai_unit_test_env["OPENAI_API_KEY"],
|
||||
org_id=openai_unit_test_env["OPENAI_ORG_ID"],
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
dumped_settings = chat_client.to_dict()
|
||||
|
||||
assert dumped_settings["ai_model_id"] == "gpt-4"
|
||||
assert dumped_settings["assistant_id"] == "test-assistant-id"
|
||||
assert dumped_settings["assistant_name"] == "TestAssistant"
|
||||
assert dumped_settings["thread_id"] == "test-thread-id"
|
||||
assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
|
||||
assert dumped_settings["org_id"] == openai_unit_test_env["OPENAI_ORG_ID"]
|
||||
|
||||
# Assert that the default header we added is present in the dumped_settings default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in dumped_settings["default_headers"]
|
||||
assert dumped_settings["default_headers"][key] == value
|
||||
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
|
||||
assert "User-Agent" not in dumped_settings["default_headers"]
|
||||
|
||||
|
||||
def get_weather(
|
||||
location: Annotated[str, Field(description="The location to get the weather for.")],
|
||||
) -> str:
|
||||
"""Get the weather for a given location."""
|
||||
return f"The weather in {location} is sunny with a high of 25°C."
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_get_response() -> None:
|
||||
"""Test OpenAI Assistants Client response."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="The weather in Seattle is currently sunny with a high of 25°C. "
|
||||
"It's a beautiful day for outdoor activities.",
|
||||
)
|
||||
)
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await openai_assistants_client.get_response(messages=messages)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert any(word in response.text.lower() for word in ["sunny", "25", "weather", "seattle"])
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_get_response_tools() -> None:
|
||||
"""Test OpenAI Assistants Client response with tools."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await openai_assistants_client.get_response(
|
||||
messages=messages,
|
||||
tools=[get_weather],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert any(word in response.text.lower() for word in ["sunny", "25", "weather"])
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_streaming() -> None:
|
||||
"""Test OpenAI Assistants Client streaming response."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="The weather in Seattle is currently sunny with a high of 25°C. "
|
||||
"It's a beautiful day for outdoor activities.",
|
||||
)
|
||||
)
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = openai_assistants_client.get_streaming_response(messages=messages)
|
||||
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
assert any(word in full_message.lower() for word in ["sunny", "25", "weather", "seattle"])
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_streaming_tools() -> None:
|
||||
"""Test OpenAI Assistants Client streaming response with tools."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = openai_assistants_client.get_streaming_response(
|
||||
messages=messages,
|
||||
tools=[get_weather],
|
||||
tool_choice="auto",
|
||||
)
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
assert any(word in full_message.lower() for word in ["sunny", "25", "weather"])
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_with_existing_assistant() -> None:
|
||||
"""Test OpenAI Assistants Client with existing assistant ID."""
|
||||
# First create an assistant to use in the test
|
||||
async with OpenAIAssistantsClient() as temp_client:
|
||||
# Get the assistant ID by triggering assistant creation
|
||||
messages = [ChatMessage(role="user", text="Hello")]
|
||||
await temp_client.get_response(messages=messages)
|
||||
assistant_id = temp_client.assistant_id
|
||||
|
||||
# Now test using the existing assistant
|
||||
async with OpenAIAssistantsClient(
|
||||
ai_model_id="gpt-4o-mini", assistant_id=assistant_id
|
||||
) as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
assert openai_assistants_client.assistant_id == assistant_id
|
||||
|
||||
messages = [ChatMessage(role="user", text="What can you do?")]
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await openai_assistants_client.get_response(messages=messages)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert len(response.text) > 0
|
||||
Reference in New Issue
Block a user