diff --git a/.github/ISSUE_TEMPLATE/1-bug_report.yml b/.github/ISSUE_TEMPLATE/1-bug_report.yml
index 81bae4057..5b106dd14 100644
--- a/.github/ISSUE_TEMPLATE/1-bug_report.yml
+++ b/.github/ISSUE_TEMPLATE/1-bug_report.yml
@@ -90,6 +90,7 @@ body:
multiple: false
options:
- "Python dev (main branch)"
+ - "Python 0.5.7"
- "Python 0.5.6"
- "Python 0.5.5"
- "Python 0.5.4"
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 0021d6afe..30df94bf1 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -33,7 +33,7 @@ jobs:
[
# For main use the workflow target
{ ref: "${{github.ref}}", dest-dir: dev, uv-version: "0.5.13", sphinx-release-override: "dev" },
- { ref: "python-v0.5.6", dest-dir: stable, uv-version: "0.5.13", sphinx-release-override: "stable" },
+ { ref: "python-v0.5.7", dest-dir: stable, uv-version: "0.5.13", sphinx-release-override: "stable" },
{ ref: "v0.4.0.post1", dest-dir: "0.4.0", uv-version: "0.5.13", sphinx-release-override: "" },
{ ref: "v0.4.1", dest-dir: "0.4.1", uv-version: "0.5.13", sphinx-release-override: "" },
{ ref: "v0.4.2", dest-dir: "0.4.2", uv-version: "0.5.13", sphinx-release-override: "" },
@@ -50,6 +50,7 @@ jobs:
{ ref: "python-v0.5.4", dest-dir: "0.5.4", uv-version: "0.5.13", sphinx-release-override: "" },
{ ref: "python-v0.5.5", dest-dir: "0.5.5", uv-version: "0.5.13", sphinx-release-override: "" },
{ ref: "python-v0.5.6", dest-dir: "0.5.6", uv-version: "0.5.13", sphinx-release-override: "" },
+ { ref: "python-v0.5.7", dest-dir: "0.5.7", uv-version: "0.5.13", sphinx-release-override: "" },
]
steps:
- name: Checkout
diff --git a/docs/switcher.json b/docs/switcher.json
index ea89227d3..3c36eef23 100644
--- a/docs/switcher.json
+++ b/docs/switcher.json
@@ -5,11 +5,16 @@
"url": "/autogen/dev/"
},
{
- "name": "0.5.6 (stable)",
+ "name": "0.5.7 (stable)",
"version": "stable",
"url": "/autogen/stable/",
"preferred": true
},
+ {
+ "name": "0.5.6",
+ "version": "0.5.6",
+ "url": "/autogen/0.5.6/"
+ },
{
"name": "0.5.5",
"version": "0.5.5",
@@ -90,4 +95,4 @@
"version": "0.2",
"url": "/autogen/0.2/"
}
-]
\ No newline at end of file
+]
diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props
index ea8f50875..5240e16ad 100644
--- a/dotnet/Directory.Packages.props
+++ b/dotnet/Directory.Packages.props
@@ -5,7 +5,8 @@
1.45.0
$(MicrosoftSemanticKernelStableVersion)-preview
$(MicrosoftSemanticKernelStableVersion)-alpha
- 9.3.0-preview.1.25161.3
+ 9.5.0
+ 9.5.0-preview.1.25265.7
9.0.0
9.0.3
9.0.0
@@ -64,9 +65,9 @@
-
-
-
+
+
+
@@ -135,4 +136,4 @@
-
\ No newline at end of file
+
diff --git a/dotnet/README.md b/dotnet/README.md
index 7647d227a..99b7ac3bf 100644
--- a/dotnet/README.md
+++ b/dotnet/README.md
@@ -40,7 +40,7 @@ var assistantAgent = new AssistantAgent(
// set human input mode to ALWAYS so that user always provide input
var userProxyAgent = new UserProxyAgent(
name: "user",
- humanInputMode: ConversableAgent.HumanInputMode.ALWAYS)
+ humanInputMode: HumanInputMode.ALWAYS)
.RegisterPrintMessage();
// start the conversation
diff --git a/dotnet/src/AutoGen.Core/Middleware/FunctionCallMiddleware.cs b/dotnet/src/AutoGen.Core/Middleware/FunctionCallMiddleware.cs
index 3c68e9545..60b9e7035 100644
--- a/dotnet/src/AutoGen.Core/Middleware/FunctionCallMiddleware.cs
+++ b/dotnet/src/AutoGen.Core/Middleware/FunctionCallMiddleware.cs
@@ -189,12 +189,12 @@ public class FunctionCallMiddleware : IStreamingMiddleware
}
}
- private Func> AIToolInvokeWrapper(Func>?, CancellationToken, Task> lambda)
+ private Func> AIToolInvokeWrapper(Func> lambda)
{
return async (string args) =>
{
var arguments = JsonSerializer.Deserialize>(args);
- var result = await lambda(arguments, CancellationToken.None);
+ var result = await lambda(new(arguments), CancellationToken.None);
return result switch
{
diff --git a/dotnet/src/Microsoft.AutoGen/Extensions/MEAI/ServiceCollectionChatCompletionExtensions.cs b/dotnet/src/Microsoft.AutoGen/Extensions/MEAI/ServiceCollectionChatCompletionExtensions.cs
index 2b398ac28..28e5b414e 100644
--- a/dotnet/src/Microsoft.AutoGen/Extensions/MEAI/ServiceCollectionChatCompletionExtensions.cs
+++ b/dotnet/src/Microsoft.AutoGen/Extensions/MEAI/ServiceCollectionChatCompletionExtensions.cs
@@ -89,7 +89,7 @@ public static class ServiceCollectionChatClientExtensions
.AddChatClient(service =>
{
var openAiClient = service.GetRequiredService();
- return openAiClient.AsChatClient(modelOrDeploymentName);
+ return openAiClient.GetChatClient(modelOrDeploymentName).AsIChatClient();
});
return services;
@@ -112,7 +112,7 @@ public static class ServiceCollectionChatClientExtensions
var endpoint = $"{serviceName}:Endpoint" ?? throw new InvalidOperationException($"No endpoint was specified for the Azure Inference Chat Client");
var endpointUri = string.IsNullOrEmpty(endpoint) ? null : new Uri(endpoint);
var token = Environment.GetEnvironmentVariable("AZURE_OPENAI_API_KEY") ?? throw new InvalidOperationException("No model access token was found in the environment variable AZURE_OPENAI_API_KEY");
- var chatClient = new ChatCompletionsClient(endpointUri, new AzureKeyCredential(token)).AsChatClient(modelOrDeploymentName);
+ var chatClient = new ChatCompletionsClient(endpointUri, new AzureKeyCredential(token)).AsIChatClient(modelOrDeploymentName);
hostBuilder.Services.AddChatClient(chatClient);
return hostBuilder.Services;
diff --git a/python/packages/autogen-agentchat/pyproject.toml b/python/packages/autogen-agentchat/pyproject.toml
index b331bb9e5..b7eba947e 100644
--- a/python/packages/autogen-agentchat/pyproject.toml
+++ b/python/packages/autogen-agentchat/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "autogen-agentchat"
-version = "0.5.6"
+version = "0.5.7"
license = {file = "LICENSE-CODE"}
description = "AutoGen agents and teams library"
readme = "README.md"
@@ -15,7 +15,7 @@ classifiers = [
"Operating System :: OS Independent",
]
dependencies = [
- "autogen-core==0.5.6",
+ "autogen-core==0.5.7",
]
[tool.ruff]
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py
index c189021a2..04acc9020 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py
@@ -131,17 +131,26 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
* If the model returns no tool call, then the response is immediately returned as a :class:`~autogen_agentchat.messages.TextMessage` or a :class:`~autogen_agentchat.messages.StructuredMessage` (when using structured output) in :attr:`~autogen_agentchat.base.Response.chat_message`.
* When the model returns tool calls, they will be executed right away:
- - When `reflect_on_tool_use` is False, the tool call results are returned as a :class:`~autogen_agentchat.messages.ToolCallSummaryMessage` in :attr:`~autogen_agentchat.base.Response.chat_message`. `tool_call_summary_format` can be used to customize the tool call summary.
+ - When `reflect_on_tool_use` is False, the tool call results are returned as a :class:`~autogen_agentchat.messages.ToolCallSummaryMessage` in :attr:`~autogen_agentchat.base.Response.chat_message`. You can customise the summary with either a static format string (`tool_call_summary_format`) **or** a callable (`tool_call_summary_formatter`); the callable is evaluated once per tool call.
- When `reflect_on_tool_use` is True, the another model inference is made using the tool calls and results, and final response is returned as a :class:`~autogen_agentchat.messages.TextMessage` or a :class:`~autogen_agentchat.messages.StructuredMessage` (when using structured output) in :attr:`~autogen_agentchat.base.Response.chat_message`.
- `reflect_on_tool_use` is set to `True` by default when `output_content_type` is set.
- `reflect_on_tool_use` is set to `False` by default when `output_content_type` is not set.
* If the model returns multiple tool calls, they will be executed concurrently. To disable parallel tool calls you need to configure the model client. For example, set `parallel_tool_calls=False` for :class:`~autogen_ext.models.openai.OpenAIChatCompletionClient` and :class:`~autogen_ext.models.openai.AzureOpenAIChatCompletionClient`.
.. tip::
- By default, the tool call results are returned as response when tool calls are made.
- So it is recommended to pay attention to the formatting of the tools return values,
- especially if another agent is expecting them in a specific format.
- Use `tool_call_summary_format` to customize the tool call summary, if needed.
+
+ By default, the tool call results are returned as the response when tool
+ calls are made, so pay close attention to how the tools’ return values
+ are formatted—especially if another agent expects a specific schema.
+
+ * Use **`tool_call_summary_format`** for a simple static template.
+ * Use **`tool_call_summary_formatter`** for full programmatic control
+ (e.g., “hide large success payloads, show full details on error”).
+
+ *Note*: `tool_call_summary_formatter` is **not serializable** and will
+ be ignored when an agent is loaded from, or exported to, YAML/JSON
+ configuration files.
+
**Hand off behavior:**
@@ -199,13 +208,22 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
If this is set, the agent will respond with a :class:`~autogen_agentchat.messages.StructuredMessage` instead of a :class:`~autogen_agentchat.messages.TextMessage`
in the final response, unless `reflect_on_tool_use` is `False` and a tool call is made.
output_content_type_format (str | None, optional): (Experimental) The format string used for the content of a :class:`~autogen_agentchat.messages.StructuredMessage` response.
- tool_call_summary_format (str, optional): The format string used to create the content for a :class:`~autogen_agentchat.messages.ToolCallSummaryMessage` response.
- The format string is used to format the tool call summary for every tool call result.
- Defaults to "{result}".
- When `reflect_on_tool_use` is `False`, a concatenation of all the tool call summaries, separated by a new line character ('\\n')
- will be returned as the response.
- Available variables: `{tool_name}`, `{arguments}`, `{result}`.
- For example, `"{tool_name}: {result}"` will create a summary like `"tool_name: result"`.
+ tool_call_summary_format (str, optional): Static format string applied to each tool call result when composing the :class:`~autogen_agentchat.messages.ToolCallSummaryMessage`.
+ Defaults to ``"{result}"``. Ignored if `tool_call_summary_formatter` is provided. When `reflect_on_tool_use` is ``False``, the summaries for all tool
+ calls are concatenated with a newline ('\\n') and returned as the response. Placeholders available in the template:
+ `{tool_name}`, `{arguments}`, `{result}`, `{is_error}`.
+ tool_call_summary_formatter (Callable[[FunctionCall, FunctionExecutionResult], str] | None, optional):
+ Callable that receives the ``FunctionCall`` and its ``FunctionExecutionResult`` and returns the summary string.
+ Overrides `tool_call_summary_format` when supplied and allows conditional logic — for example, emitting static string like
+ ``"Tool FooBar executed successfully."`` on success and a full payload (including all passed arguments etc.) only on failure.
+
+ **Limitation**: The callable is *not serializable*; values provided via YAML/JSON configs are ignored.
+
+ .. note::
+
+ `tool_call_summary_formatter` is intended for in-code use only. It cannot currently be saved or restored via
+ configuration files.
+
memory (Sequence[Memory] | None, optional): The memory store to use for the agent. Defaults to `None`.
metadata (Dict[str, str] | None, optional): Optional metadata for tracking.
@@ -652,6 +670,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
model_client_stream: bool = False,
reflect_on_tool_use: bool | None = None,
tool_call_summary_format: str = "{result}",
+ tool_call_summary_formatter: Callable[[FunctionCall, FunctionExecutionResult], str] | None = None,
output_content_type: type[BaseModel] | None = None,
output_content_type_format: str | None = None,
memory: Sequence[Memory] | None = None,
@@ -756,6 +775,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
stacklevel=2,
)
self._tool_call_summary_format = tool_call_summary_format
+ self._tool_call_summary_formatter = tool_call_summary_formatter
self._is_running = False
@property
@@ -803,6 +823,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
model_client_stream = self._model_client_stream
reflect_on_tool_use = self._reflect_on_tool_use
tool_call_summary_format = self._tool_call_summary_format
+ tool_call_summary_formatter = self._tool_call_summary_formatter
output_content_type = self._output_content_type
format_string = self._output_content_type_format
@@ -873,6 +894,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
model_client_stream=model_client_stream,
reflect_on_tool_use=reflect_on_tool_use,
tool_call_summary_format=tool_call_summary_format,
+ tool_call_summary_formatter=tool_call_summary_formatter,
output_content_type=output_content_type,
format_string=format_string,
):
@@ -976,6 +998,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
model_client_stream: bool,
reflect_on_tool_use: bool,
tool_call_summary_format: str,
+ tool_call_summary_formatter: Callable[[FunctionCall, FunctionExecutionResult], str] | None,
output_content_type: type[BaseModel] | None,
format_string: str | None = None,
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
@@ -1078,6 +1101,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
inner_messages=inner_messages,
handoffs=handoffs,
tool_call_summary_format=tool_call_summary_format,
+ tool_call_summary_formatter=tool_call_summary_formatter,
agent_name=agent_name,
)
@@ -1232,6 +1256,7 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
inner_messages: List[BaseAgentEvent | BaseChatMessage],
handoffs: Dict[str, HandoffBase],
tool_call_summary_format: str,
+ tool_call_summary_formatter: Callable[[FunctionCall, FunctionExecutionResult], str] | None,
agent_name: str,
) -> Response:
"""
@@ -1239,15 +1264,22 @@ class AssistantAgent(BaseChatAgent, Component[AssistantAgentConfig]):
"""
# Filter out calls which were actually handoffs
normal_tool_calls = [(call, result) for call, result in executed_calls_and_results if call.name not in handoffs]
- tool_call_summaries: List[str] = []
- for tool_call, tool_call_result in normal_tool_calls:
- tool_call_summaries.append(
- tool_call_summary_format.format(
- tool_name=tool_call.name,
- arguments=tool_call.arguments,
- result=tool_call_result.content,
- )
+
+ def default_tool_call_summary_formatter(call: FunctionCall, result: FunctionExecutionResult) -> str:
+ return tool_call_summary_format
+
+ summary_formatter = tool_call_summary_formatter or default_tool_call_summary_formatter
+
+ tool_call_summaries = [
+ summary_formatter(call, result).format(
+ tool_name=call.name,
+ arguments=call.arguments,
+ result=result.content,
+ is_error=result.is_error,
)
+ for call, result in normal_tool_calls
+ ]
+
tool_call_summary = "\n".join(tool_call_summaries)
return Response(
chat_message=ToolCallSummaryMessage(
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py
index 80e82068a..2efa6d5a7 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py
@@ -326,8 +326,8 @@ class CodeExecutorAgent(BaseChatAgent, Component[CodeExecutorAgentConfig]):
"""
DEFAULT_TERMINAL_DESCRIPTION = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks)."
- DEFAULT_AGENT_DESCRIPTION = "A Code Execution Agent that generates and executes Python and shell scripts based on user instructions. Python code should be provided in ```python code blocks, and sh shell scripts should be provided in ```sh code blocks for execution. It ensures correctness, efficiency, and minimal errors while gracefully handling edge cases."
- DEFAULT_SYSTEM_MESSAGE = "You are a Code Execution Agent. Your role is to generate and execute Python code based on user instructions, ensuring correctness, efficiency, and minimal errors. Handle edge cases gracefully."
+ DEFAULT_AGENT_DESCRIPTION = "A Code Execution Agent that generates and executes Python and shell scripts based on user instructions. It ensures correctness, efficiency, and minimal errors while gracefully handling edge cases."
+ DEFAULT_SYSTEM_MESSAGE = "You are a Code Execution Agent. Your role is to generate and execute Python code and shell scripts based on user instructions, ensuring correctness, efficiency, and minimal errors. Handle edge cases gracefully. Python code should be provided in ```python code blocks, and sh shell scripts should be provided in ```sh code blocks for execution."
NO_CODE_BLOCKS_FOUND_MESSAGE = "No code blocks found in the thread. Please provide at least one markdown-encoded code block to execute (i.e., quoting code in ```python or ```sh code blocks)."
component_config_schema = CodeExecutorAgentConfig
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat_manager.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat_manager.py
index 1ebe658c1..6c008bef6 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat_manager.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat_manager.py
@@ -2,7 +2,7 @@ import asyncio
from abc import ABC, abstractmethod
from typing import Any, List, Sequence
-from autogen_core import DefaultTopicId, MessageContext, event, rpc
+from autogen_core import CancellationToken, DefaultTopicId, MessageContext, event, rpc
from ...base import TerminationCondition
from ...messages import BaseAgentEvent, BaseChatMessage, MessageFactory, SelectSpeakerEvent, StopMessage
@@ -79,6 +79,7 @@ class BaseGroupChatManager(SequentialRoutedAgent, ABC):
self._current_turn = 0
self._message_factory = message_factory
self._emit_team_events = emit_team_events
+ self._active_speakers: List[str] = []
@rpc
async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> None:
@@ -122,64 +123,64 @@ class BaseGroupChatManager(SequentialRoutedAgent, ABC):
# Stop the group chat.
return
- # Select a speaker to start/continue the conversation
- speaker_name_future = asyncio.ensure_future(self.select_speaker(self._message_thread))
- # Link the select speaker future to the cancellation token.
- ctx.cancellation_token.link_future(speaker_name_future)
- speaker_name = await speaker_name_future
- if speaker_name not in self._participant_name_to_topic_type:
- raise RuntimeError(f"Speaker {speaker_name} not found in participant names.")
- await self._log_speaker_selection(speaker_name)
-
- # Send the message to the next speaker
- speaker_topic_type = self._participant_name_to_topic_type[speaker_name]
- await self.publish_message(
- GroupChatRequestPublish(),
- topic_id=DefaultTopicId(type=speaker_topic_type),
- cancellation_token=ctx.cancellation_token,
- )
-
- async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None:
- self._message_thread.extend(messages)
+ # Select speakers to start/continue the conversation
+ await self._transition_to_next_speakers(ctx.cancellation_token)
@event
async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: MessageContext) -> None:
try:
- # Append the message to the message thread and construct the delta.
+ # Construct the detla from the agent response.
delta: List[BaseAgentEvent | BaseChatMessage] = []
if message.agent_response.inner_messages is not None:
for inner_message in message.agent_response.inner_messages:
delta.append(inner_message)
delta.append(message.agent_response.chat_message)
+
+ # Append the messages to the message thread.
await self.update_message_thread(delta)
+ # Remove the agent from the active speakers list.
+ self._active_speakers.remove(message.agent_name)
+ if len(self._active_speakers) > 0:
+ # If there are still active speakers, return without doing anything.
+ return
+
# Check if the conversation should be terminated.
if await self._apply_termination_condition(delta, increment_turn_count=True):
# Stop the group chat.
return
- # Select a speaker to continue the conversation.
- speaker_name_future = asyncio.ensure_future(self.select_speaker(self._message_thread))
- # Link the select speaker future to the cancellation token.
- ctx.cancellation_token.link_future(speaker_name_future)
- speaker_name = await speaker_name_future
+ # Select speakers to continue the conversation.
+ await self._transition_to_next_speakers(ctx.cancellation_token)
+ except Exception as e:
+ # Handle the exception and signal termination with an error.
+ error = SerializableException.from_exception(e)
+ await self._signal_termination_with_error(error)
+ # Raise the exception to the runtime.
+ raise
+
+ async def _transition_to_next_speakers(self, cancellation_token: CancellationToken) -> None:
+ speaker_names_future = asyncio.ensure_future(self.select_speaker(self._message_thread))
+ # Link the select speaker future to the cancellation token.
+ cancellation_token.link_future(speaker_names_future)
+ speaker_names = await speaker_names_future
+ if isinstance(speaker_names, str):
+ # If only one speaker is selected, convert it to a list.
+ speaker_names = [speaker_names]
+ for speaker_name in speaker_names:
if speaker_name not in self._participant_name_to_topic_type:
raise RuntimeError(f"Speaker {speaker_name} not found in participant names.")
- await self._log_speaker_selection(speaker_name)
+ await self._log_speaker_selection(speaker_names)
- # Send the message to the next speakers
+ # Send request to publish message to the next speakers
+ for speaker_name in speaker_names:
speaker_topic_type = self._participant_name_to_topic_type[speaker_name]
await self.publish_message(
GroupChatRequestPublish(),
topic_id=DefaultTopicId(type=speaker_topic_type),
- cancellation_token=ctx.cancellation_token,
+ cancellation_token=cancellation_token,
)
- except Exception as e:
- # Handle the exception and signal termination with an error.
- error = SerializableException.from_exception(e)
- await self._signal_termination_with_error(error)
- # Raise the exception to the runtime.
- raise
+ self._active_speakers.append(speaker_name)
async def _apply_termination_condition(
self, delta: Sequence[BaseAgentEvent | BaseChatMessage], increment_turn_count: bool = False
@@ -216,9 +217,9 @@ class BaseGroupChatManager(SequentialRoutedAgent, ABC):
return True
return False
- async def _log_speaker_selection(self, speaker_name: str) -> None:
+ async def _log_speaker_selection(self, speaker_names: List[str]) -> None:
"""Log the selected speaker to the output message queue."""
- select_msg = SelectSpeakerEvent(content=[speaker_name], source=self._name)
+ select_msg = SelectSpeakerEvent(content=speaker_names, source=self._name)
if self._emit_team_events:
await self.publish_message(
GroupChatMessage(message=select_msg),
@@ -284,10 +285,26 @@ class BaseGroupChatManager(SequentialRoutedAgent, ABC):
"""
...
+ async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None:
+ """Update the message thread with the new messages.
+ This is called when the group chat receives a GroupChatStart or GroupChatAgentResponse event,
+ before calling the select_speakers method.
+ """
+ self._message_thread.extend(messages)
+
@abstractmethod
- async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
- """Select a speaker from the participants and return the
- topic type of the selected speaker."""
+ async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
+ """Select speakers from the participants and return the topic types of the selected speaker.
+ This is called when the group chat manager have received all responses from the participants
+ for a turn and is ready to select the next speakers for the next turn.
+
+ Args:
+ thread: The message thread of the group chat.
+
+ Returns:
+ A list of topic types of the selected speakers.
+ If only one speaker is selected, a single string is returned instead of a list.
+ """
...
@abstractmethod
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_chat_agent_container.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_chat_agent_container.py
index 69faeb491..f9ec8636f 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_chat_agent_container.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_chat_agent_container.py
@@ -89,7 +89,7 @@ class ChatAgentContainer(SequentialRoutedAgent):
# Publish the response to the group chat.
self._message_buffer.clear()
await self.publish_message(
- GroupChatAgentResponse(agent_response=response),
+ GroupChatAgentResponse(agent_response=response, agent_name=self._agent.name),
topic_id=DefaultTopicId(type=self._parent_topic_type),
cancellation_token=ctx.cancellation_token,
)
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_events.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_events.py
index ca07d87bb..febb69e60 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_events.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_events.py
@@ -48,6 +48,9 @@ class GroupChatAgentResponse(BaseModel):
agent_response: Response
"""The response from an agent."""
+ agent_name: str
+ """The name of the agent that produced the response."""
+
class GroupChatRequestPublish(BaseModel):
"""A request to publish a message to a group chat."""
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_digraph_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_digraph_group_chat.py
index 87b083b3d..94d133b20 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_digraph_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_graph/_digraph_group_chat.py
@@ -1,5 +1,6 @@
import asyncio
-from typing import Any, Callable, Dict, List, Literal, Mapping, Sequence, Set
+from collections import Counter, deque
+from typing import Any, Callable, Deque, Dict, List, Literal, Mapping, Sequence, Set
from autogen_core import AgentRuntime, CancellationToken, Component, ComponentModel
from pydantic import BaseModel
@@ -208,154 +209,70 @@ class GraphFlowManager(BaseGroupChatManager):
max_turns=max_turns,
message_factory=message_factory,
)
- self._graph = graph
- self._graph.graph_validate()
- self._graph_has_cycles = self._graph.get_has_cycles()
- if self._graph_has_cycles and self._termination_condition is None and self._max_turns is None:
+ graph.graph_validate()
+ if graph.get_has_cycles() and self._termination_condition is None and self._max_turns is None:
raise ValueError("A termination condition is required for cyclic graphs without a maximum turn limit.")
-
- self._use_default_start = self._graph.default_start_node is not None
- self._default_start_executed = False
- self._start_nodes = graph.get_start_nodes()
- self._leaf_nodes = graph.get_leaf_nodes()
- self._parents = graph.get_parents() # Parent node dependencies - helper dict to get all incoming edges
- self._active_nodes: Set[str] = set() # Currently executing nodes
- self._active_node_count: Dict[str, int] = {
- node: 0 for node in graph.nodes
- } # Number of times a node has been active
-
- # These are nodes next in line for execution as one or more of their parent nodes have started execution.
- # They execute when all their parent nodes have executed.
- # Nodes are added to this dict when at least one of their parent nodes becomes active.
- # Start nodes (no parents) are added to this dict at initialization as they are always ready to run.
- self._pending_execution: Dict[str, List[str]] = {node: [] for node in graph.get_start_nodes()}
-
- def _get_valid_target(self, node: DiGraphNode, content: str) -> str:
- """Check if a condition is met in the chat history."""
- for edge in node.edges:
- if edge.condition and edge.condition in content:
- return edge.target
-
- raise RuntimeError(f"Condition not met for node {node.name}. Content: {content}")
-
- def _is_node_ready(self, node_name: str) -> bool:
- """Check if a node is ready to execute based on its parent nodes.
- If activation is any then execute as soon as any parent has finished
- If activation is all then execute only when all parents have finished
- """
- node = self._graph.nodes[node_name]
- if node.activation == "any":
- return bool(self._pending_execution[node_name])
- return all(parent in self._pending_execution[node_name] for parent in self._parents[node_name])
-
- async def _select_speakers(self, thread: List[BaseAgentEvent | BaseChatMessage], many: bool = True) -> List[str]:
- """Select the next set of agents to execute based on DAG constraints."""
- next_speakers: Set[str] = set()
- source_node: DiGraphNode | None = None
- source: str | None = None
-
- if thread and isinstance(thread[-1], BaseChatMessage):
- source = thread[-1].source # name of the agent that just finished
- content = thread[-1].to_model_text()
-
- # Safety check: only an active node can send a response
- if source != "user":
- if source not in self._active_nodes:
- raise RuntimeError(f"Agent '{source}' is not currently active.")
-
- # Mark the node as no longer active (it just finished)
- self._active_node_count[source] -= 1
-
- if self._active_node_count[source] <= 0:
- self._active_nodes.remove(source)
-
- source_node = self._graph.nodes[source]
-
- if source_node.edges:
- # Case: conditional edges — only execute if condition is met
- target_nodes_names: List[str] = []
- if source_node.edges[0].condition is not None:
- target_nodes_names = [self._get_valid_target(source_node, content)]
- other_nodes = [
- edge.target for edge in source_node.edges if edge.target != target_nodes_names[0]
- ]
- for other_node in other_nodes:
- other_active_parents = [
- parent
- for parent in self._parents[other_node]
- if (parent != source and parent in self._active_nodes)
- ]
- if not other_active_parents:
- self._pending_execution.pop(other_node)
- else:
- self._pending_execution[other_node] = other_active_parents
-
- else:
- # Case: unconditional edges — mark this source as completed for all its children
- target_nodes_names = [edge.target for edge in source_node.edges]
-
- for target in target_nodes_names:
- self._pending_execution[target].append(source)
+ self._graph = graph
+ # Lookup table for incoming edges for each node.
+ self._parents = graph.get_parents()
+ # Lookup table for outgoing edges for each node.
+ self._edges: Dict[str, List[DiGraphEdge]] = {n: node.edges for n, node in graph.nodes.items()}
+ # Activation lookup table for each node.
+ self._activation: Dict[str, Literal["any", "all"]] = {n: node.activation for n, node in graph.nodes.items()}
+
+ # === Mutable states for the graph execution ===
+ # Count the number of remaining parents to activate each node.
+ self._remaining: Counter[str] = Counter({n: len(p) for n, p in self._parents.items()})
+ # Lookup table for nodes that have been enqueued through an any activation.
+ # This is used to prevent re-adding the same node multiple times.
+ self._enqueued_any: Dict[str, bool] = {n: False for n in graph.nodes}
+ # Ready queue for nodes that are ready to execute, starting with the start nodes.
+ self._ready: Deque[str] = deque([n for n in graph.get_start_nodes()])
+
+ async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None:
+ await super().update_message_thread(messages)
+
+ # Find the node that ran in the current turn.
+ message = messages[-1]
+ if message.source not in self._graph.nodes:
+ # Ignore messages from sources outside of the graph.
+ return
+ assert isinstance(message, BaseChatMessage)
+ source = message.source
+ content = message.to_model_text()
+
+ # Propagate the update to the children of the node.
+ for edge in self._edges[source]:
+ if edge.condition and edge.condition not in content:
+ continue
+ if self._activation[edge.target] == "all":
+ self._remaining[edge.target] -= 1
+ if self._remaining[edge.target] == 0:
+ # If all parents are done, add to the ready queue.
+ self._ready.append(edge.target)
else:
- # TODO: Check if there are any usecase where the User can decide on the next speaker
- pass
-
- # After updating _pending_execution, check which nodes are now unblocked
- for node_name in list(self._pending_execution):
- if self._use_default_start and not self._default_start_executed:
- if node_name == self._graph.default_start_node:
- next_speakers.add(node_name)
- self._default_start_executed = True
- break
-
- if self._is_node_ready(node_name):
- next_speakers.add(node_name)
- node = self._graph.nodes[node_name]
- if node.activation == "all":
- self._pending_execution.pop(node_name)
- else:
- # If activation is any, remove the parent that just finished
- if source is not None:
- self._pending_execution[node_name] = [
- parent for parent in self._pending_execution[node_name] if parent != source
- ]
-
- # If none of the other parents of this node are active, remove this node from pending execution
- node_parents = self._parents[node_name]
- if not any(parent in self._active_nodes for parent in node_parents):
- self._pending_execution.pop(node_name)
-
- if not many:
- break
-
- # Prepopulate children of next_speakers into _pending_execution
- for node_name in next_speakers:
- for edge in self._graph.nodes[node_name].edges:
- if edge.target not in self._pending_execution:
- self._pending_execution[edge.target] = []
-
- # Mark newly selected speakers as active
- for speaker in next_speakers:
- if speaker not in self._active_nodes:
- self._active_nodes.add(speaker)
-
- self._active_node_count[speaker] += 1
-
- if not self._pending_execution and not next_speakers and not self._active_nodes:
- next_speakers = set([_DIGRAPH_STOP_AGENT_NAME]) # Call the termination agent
-
- return list(next_speakers)
-
- async def select_speakers(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> List[str]:
- return await self._select_speakers(thread)
-
- async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
- """Select a speaker from the participants and return the
- topic type of the selected speaker."""
- speakers = await self._select_speakers(thread, many=False)
+ # If activation is any, add to the ready queue if not already enqueued.
+ if not self._enqueued_any[edge.target]:
+ self._ready.append(edge.target)
+ self._enqueued_any[edge.target] = True
+
+ async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str]:
+ # Drain the ready queue for the next set of speakers.
+ speakers: List[str] = []
+ while self._ready:
+ speaker = self._ready.popleft()
+ speakers.append(speaker)
+ # Reset the bookkeeping for the node that were selected.
+ if self._activation[speaker] == "any":
+ self._enqueued_any[speaker] = False
+ else:
+ self._remaining[speaker] = len(self._parents[speaker])
+
+ # If there are no speakers, trigger the stop agent.
if not speakers:
- raise RuntimeError("No available speakers found.")
- return speakers[0]
+ speakers = [_DIGRAPH_STOP_AGENT_NAME]
+
+ return speakers
async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
pass
@@ -365,10 +282,9 @@ class GraphFlowManager(BaseGroupChatManager):
state = {
"message_thread": [message.dump() for message in self._message_thread],
"current_turn": self._current_turn,
- "active_nodes": list(self._active_nodes),
- "pending_execution": self._pending_execution,
- "active_node_count": self._active_node_count,
- "default_start_executed": self._default_start_executed,
+ "remaining": dict(self._remaining),
+ "enqueued_any": dict(self._enqueued_any),
+ "ready": list(self._ready),
}
return state
@@ -376,10 +292,9 @@ class GraphFlowManager(BaseGroupChatManager):
"""Restore execution state from saved data."""
self._message_thread = [self._message_factory.create(msg) for msg in state["message_thread"]]
self._current_turn = state["current_turn"]
- self._active_nodes = set(state["active_nodes"])
- self._pending_execution = state["pending_execution"]
- self._active_node_count = state["active_node_count"]
- self._default_start_executed = state.get("default_start_executed", False)
+ self._remaining = Counter(state["remaining"])
+ self._enqueued_any = state["enqueued_any"]
+ self._ready = deque(state["ready"])
async def reset(self) -> None:
"""Reset execution state to the start of the graph."""
@@ -387,11 +302,9 @@ class GraphFlowManager(BaseGroupChatManager):
self._message_thread.clear()
if self._termination_condition:
await self._termination_condition.reset()
-
- self._active_nodes = set()
- self._active_node_count = {node: 0 for node in self._graph.nodes}
- self._pending_execution = {node: [] for node in self._start_nodes}
- self._default_start_executed = False
+ self._remaining = Counter({n: len(p) for n, p in self._parents.items()})
+ self._enqueued_any = {n: False for n in self._graph.nodes}
+ self._ready = deque([n for n in self._graph.get_start_nodes()])
class _StopAgent(BaseChatAgent):
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py
index 34d1df7cf..e92177001 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py
@@ -2,7 +2,7 @@ import asyncio
import json
import logging
import re
-from typing import Any, Dict, List, Mapping
+from typing import Any, Dict, List, Mapping, Sequence
from autogen_core import AgentId, CancellationToken, DefaultTopicId, MessageContext, event, rpc
from autogen_core.models import (
@@ -37,6 +37,7 @@ from .._events import (
GroupChatReset,
GroupChatStart,
GroupChatTermination,
+ SerializableException,
)
from ._prompts import (
ORCHESTRATOR_FINAL_ANSWER_PROMPT,
@@ -187,22 +188,29 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
@event
async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: MessageContext) -> None: # type: ignore
- delta: List[BaseAgentEvent | BaseChatMessage] = []
- if message.agent_response.inner_messages is not None:
- for inner_message in message.agent_response.inner_messages:
- delta.append(inner_message)
- await self.update_message_thread([message.agent_response.chat_message])
- delta.append(message.agent_response.chat_message)
-
- if self._termination_condition is not None:
- stop_message = await self._termination_condition(delta)
- if stop_message is not None:
- # Reset the termination conditions.
- await self._termination_condition.reset()
- # Signal termination.
- await self._signal_termination(stop_message)
- return
- await self._orchestrate_step(ctx.cancellation_token)
+ try:
+ delta: List[BaseAgentEvent | BaseChatMessage] = []
+ if message.agent_response.inner_messages is not None:
+ for inner_message in message.agent_response.inner_messages:
+ delta.append(inner_message)
+ await self.update_message_thread([message.agent_response.chat_message])
+ delta.append(message.agent_response.chat_message)
+
+ if self._termination_condition is not None:
+ stop_message = await self._termination_condition(delta)
+ if stop_message is not None:
+ # Reset the termination conditions.
+ await self._termination_condition.reset()
+ # Signal termination.
+ await self._signal_termination(stop_message)
+ return
+
+ await self._orchestrate_step(ctx.cancellation_token)
+ except Exception as e:
+ error = SerializableException.from_exception(e)
+ await self._signal_termination_with_error(error)
+ # Raise the error to the runtime.
+ raise
async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
pass
@@ -229,9 +237,9 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
self._n_rounds = orchestrator_state.n_rounds
self._n_stalls = orchestrator_state.n_stalls
- async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
+ async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
"""Not used in this orchestrator, we select next speaker in _orchestrate_step."""
- return ""
+ return [""]
async def reset(self) -> None:
"""Reset the group chat manager."""
@@ -275,7 +283,7 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
# Broadcast
await self.publish_message(
- GroupChatAgentResponse(agent_response=Response(chat_message=ledger_message)),
+ GroupChatAgentResponse(agent_response=Response(chat_message=ledger_message), agent_name=self._name),
topic_id=DefaultTopicId(type=self._group_topic_type),
)
@@ -389,7 +397,7 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
# Broadcast it
await self.publish_message( # Broadcast
- GroupChatAgentResponse(agent_response=Response(chat_message=message)),
+ GroupChatAgentResponse(agent_response=Response(chat_message=message), agent_name=self._name),
topic_id=DefaultTopicId(type=self._group_topic_type),
cancellation_token=cancellation_token,
)
@@ -470,7 +478,7 @@ class MagenticOneOrchestrator(BaseGroupChatManager):
# Broadcast
await self.publish_message(
- GroupChatAgentResponse(agent_response=Response(chat_message=message)),
+ GroupChatAgentResponse(agent_response=Response(chat_message=message), agent_name=self._name),
topic_id=DefaultTopicId(type=self._group_topic_type),
cancellation_token=cancellation_token,
)
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
index 6ea915511..d6b43afb2 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
@@ -1,5 +1,5 @@
import asyncio
-from typing import Any, Callable, List, Mapping
+from typing import Any, Callable, List, Mapping, Sequence
from autogen_core import AgentRuntime, Component, ComponentModel
from pydantic import BaseModel
@@ -69,8 +69,13 @@ class RoundRobinGroupChatManager(BaseGroupChatManager):
self._current_turn = round_robin_state.current_turn
self._next_speaker_index = round_robin_state.next_speaker_index
- async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
- """Select a speaker from the participants in a round-robin fashion."""
+ async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
+ """Select a speaker from the participants in a round-robin fashion.
+
+ .. note::
+
+ This method always returns a single speaker.
+ """
current_speaker_index = self._next_speaker_index
self._next_speaker_index = (current_speaker_index + 1) % len(self._participant_names)
current_speaker = self._participant_names[current_speaker_index]
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
index d54b7aea2..a0d9e7732 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
@@ -150,10 +150,14 @@ class SelectorGroupChatManager(BaseGroupChatManager):
base_chat_messages = [m for m in messages if isinstance(m, BaseChatMessage)]
await self._add_messages_to_context(self._model_context, base_chat_messages)
- async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
+ async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
"""Selects the next speaker in a group chat using a ChatCompletion client,
with the selector function as override if it returns a speaker name.
+ .. note::
+
+ This method always returns a single speaker name.
+
A key assumption is that the agent type is the same as the topic type, which we use as the agent name.
"""
# Use the selector function if provided.
@@ -171,7 +175,7 @@ class SelectorGroupChatManager(BaseGroupChatManager):
f"Expected one of: {self._participant_names}."
)
# Skip the model based selection.
- return speaker
+ return [speaker]
# Use the candidate function to filter participants if provided
if self._candidate_func is not None:
@@ -211,7 +215,7 @@ class SelectorGroupChatManager(BaseGroupChatManager):
agent_name = participants[0]
self._previous_speaker = agent_name
trace_logger.debug(f"Selected speaker: {agent_name}")
- return agent_name
+ return [agent_name]
def construct_message_history(self, message_history: List[LLMMessage]) -> str:
# Construct the history of the conversation.
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
index 7842e4c46..344994015 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
@@ -1,5 +1,5 @@
import asyncio
-from typing import Any, Callable, List, Mapping
+from typing import Any, Callable, List, Mapping, Sequence
from autogen_core import AgentRuntime, Component, ComponentModel
from pydantic import BaseModel
@@ -79,17 +79,22 @@ class SwarmGroupChatManager(BaseGroupChatManager):
await self._termination_condition.reset()
self._current_speaker = self._participant_names[0]
- async def select_speaker(self, thread: List[BaseAgentEvent | BaseChatMessage]) -> str:
+ async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
"""Select a speaker from the participants based on handoff message.
- Looks for the last handoff message in the thread to determine the next speaker."""
+ Looks for the last handoff message in the thread to determine the next speaker.
+
+ .. note::
+
+ This method always returns a single speaker.
+ """
if len(thread) == 0:
- return self._current_speaker
+ return [self._current_speaker]
for message in reversed(thread):
if isinstance(message, HandoffMessage):
self._current_speaker = message.target
# The latest handoff message should always target a valid participant.
assert self._current_speaker in self._participant_names
- return self._current_speaker
+ return [self._current_speaker]
return self._current_speaker
async def save_state(self) -> Mapping[str, Any]:
diff --git a/python/packages/autogen-agentchat/tests/test_assistant_agent.py b/python/packages/autogen-agentchat/tests/test_assistant_agent.py
index 666deb511..988707d95 100644
--- a/python/packages/autogen-agentchat/tests/test_assistant_agent.py
+++ b/python/packages/autogen-agentchat/tests/test_assistant_agent.py
@@ -32,7 +32,7 @@ from autogen_core.models import (
SystemMessage,
UserMessage,
)
-from autogen_core.models._model_client import ModelFamily
+from autogen_core.models._model_client import ModelFamily, ModelInfo
from autogen_core.tools import BaseTool, FunctionTool, StaticWorkbench
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_ext.models.replay import ReplayChatCompletionClient
@@ -56,10 +56,68 @@ async def _fail_function(input: str) -> str:
return "fail"
+async def _throw_function(input: str) -> str:
+ raise ValueError("Helpful debugging information what went wrong.")
+
+
async def _echo_function(input: str) -> str:
return input
+@pytest.fixture
+def model_info_all_capabilities() -> ModelInfo:
+ return {
+ "function_calling": True,
+ "vision": True,
+ "json_output": True,
+ "family": ModelFamily.GPT_4O,
+ "structured_output": True,
+ }
+
+
+@pytest.mark.asyncio
+async def test_run_with_tool_call_summary_format_function(model_info_all_capabilities: ModelInfo) -> None:
+ model_client = ReplayChatCompletionClient(
+ [
+ CreateResult(
+ finish_reason="function_calls",
+ content=[
+ FunctionCall(id="1", arguments=json.dumps({"input": "task"}), name="_pass_function"),
+ FunctionCall(id="2", arguments=json.dumps({"input": "task"}), name="_throw_function"),
+ ],
+ usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+ thought="Calling pass and fail function",
+ cached=False,
+ ),
+ ],
+ model_info=model_info_all_capabilities,
+ )
+
+ def conditional_string_templates(function_call: FunctionCall, function_call_result: FunctionExecutionResult) -> str:
+ if not function_call_result.is_error:
+ return "SUCCESS: {tool_name} with {arguments}"
+
+ else:
+ return "FAILURE: {result}"
+
+ agent = AssistantAgent(
+ "tool_use_agent",
+ model_client=model_client,
+ tools=[_pass_function, _throw_function],
+ tool_call_summary_formatter=conditional_string_templates,
+ )
+ result = await agent.run(task="task")
+
+ first_tool_call_summary = next((x for x in result.messages if isinstance(x, ToolCallSummaryMessage)), None)
+ if first_tool_call_summary is None:
+ raise AssertionError("Expected a ToolCallSummaryMessage but found none")
+
+ assert (
+ first_tool_call_summary.content
+ == 'SUCCESS: _pass_function with {"input": "task"}\nFAILURE: Helpful debugging information what went wrong.'
+ )
+
+
@pytest.mark.asyncio
async def test_run_with_tools(monkeypatch: pytest.MonkeyPatch) -> None:
model_client = ReplayChatCompletionClient(
diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py
index e7e8ba436..889b67d87 100644
--- a/python/packages/autogen-agentchat/tests/test_group_chat.py
+++ b/python/packages/autogen-agentchat/tests/test_group_chat.py
@@ -6,24 +6,6 @@ from typing import Any, AsyncGenerator, Dict, List, Mapping, Sequence
import pytest
import pytest_asyncio
-from autogen_core import AgentId, AgentRuntime, CancellationToken, FunctionCall, SingleThreadedAgentRuntime
-from autogen_core.model_context import BufferedChatCompletionContext
-from autogen_core.models import (
- AssistantMessage,
- CreateResult,
- FunctionExecutionResult,
- FunctionExecutionResultMessage,
- LLMMessage,
- RequestUsage,
- UserMessage,
-)
-from autogen_core.tools import FunctionTool
-from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor
-from autogen_ext.models.openai import OpenAIChatCompletionClient
-from autogen_ext.models.replay import ReplayChatCompletionClient
-from pydantic import BaseModel
-from utils import FileLogHandler
-
from autogen_agentchat import EVENT_LOGGER_NAME
from autogen_agentchat.agents import (
AssistantAgent,
@@ -57,6 +39,23 @@ from autogen_agentchat.teams._group_chat._round_robin_group_chat import RoundRob
from autogen_agentchat.teams._group_chat._selector_group_chat import SelectorGroupChatManager
from autogen_agentchat.teams._group_chat._swarm_group_chat import SwarmGroupChatManager
from autogen_agentchat.ui import Console
+from autogen_core import AgentId, AgentRuntime, CancellationToken, FunctionCall, SingleThreadedAgentRuntime
+from autogen_core.model_context import BufferedChatCompletionContext
+from autogen_core.models import (
+ AssistantMessage,
+ CreateResult,
+ FunctionExecutionResult,
+ FunctionExecutionResultMessage,
+ LLMMessage,
+ RequestUsage,
+ UserMessage,
+)
+from autogen_core.tools import FunctionTool
+from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor
+from autogen_ext.models.openai import OpenAIChatCompletionClient
+from autogen_ext.models.replay import ReplayChatCompletionClient
+from pydantic import BaseModel
+from utils import FileLogHandler
logger = logging.getLogger(EVENT_LOGGER_NAME)
logger.setLevel(logging.DEBUG)
diff --git a/python/packages/autogen-agentchat/tests/test_group_chat_graph.py b/python/packages/autogen-agentchat/tests/test_group_chat_graph.py
index 7c8baf4b1..2f381528c 100644
--- a/python/packages/autogen-agentchat/tests/test_group_chat_graph.py
+++ b/python/packages/autogen-agentchat/tests/test_group_chat_graph.py
@@ -1,6 +1,6 @@
import asyncio
-from typing import Any, AsyncGenerator, Callable, Dict, List, Sequence, Set
-from unittest.mock import AsyncMock, patch
+from typing import AsyncGenerator, List, Sequence
+from unittest.mock import patch
import pytest
import pytest_asyncio
@@ -12,9 +12,8 @@ from autogen_agentchat.agents import (
PerSourceFilter,
)
from autogen_agentchat.base import Response, TaskResult
-from autogen_agentchat.conditions import MaxMessageTermination
+from autogen_agentchat.conditions import MaxMessageTermination, SourceMatchTermination
from autogen_agentchat.messages import BaseChatMessage, ChatMessage, MessageFactory, StopMessage, TextMessage
-from autogen_agentchat.messages import BaseTextChatMessage as TextChatMessage
from autogen_agentchat.teams import (
DiGraphBuilder,
GraphFlow,
@@ -269,61 +268,6 @@ def test_validate_graph_mixed_conditions() -> None:
graph.graph_validate()
-def test_get_valid_target() -> None:
- node = DiGraphNode(
- name="A",
- edges=[DiGraphEdge(target="B", condition="approve"), DiGraphEdge(target="C", condition="reject")],
- )
- manager = GraphFlowManager.__new__(GraphFlowManager)
-
- assert manager._get_valid_target(node, "please approve this") == "B" # pyright: ignore[reportPrivateUsage]
- assert manager._get_valid_target(node, "i reject this") == "C" # pyright: ignore[reportPrivateUsage]
- with pytest.raises(RuntimeError):
- manager._get_valid_target(node, "unknown path") # pyright: ignore[reportPrivateUsage]
-
-
-def test_is_node_ready_all_and_any() -> None:
- graph = DiGraph(
- nodes={
- "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="C")]),
- "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]),
- "C": DiGraphNode(name="C", edges=[], activation="all"),
- }
- )
-
- manager = GraphFlowManager.__new__(GraphFlowManager)
- manager._graph = graph # pyright: ignore[reportPrivateUsage]
- manager._parents = graph.get_parents() # pyright: ignore[reportPrivateUsage]
-
- # === Test "all" activation ===
- # Case 1: No parent finished
- manager._pending_execution = {"C": []} # pyright: ignore[reportPrivateUsage]
- assert not manager._is_node_ready("C") # pyright: ignore[reportPrivateUsage]
-
- # Case 2: One parent finished
- manager._pending_execution = {"C": ["A"]} # pyright: ignore[reportPrivateUsage]
- assert not manager._is_node_ready("C") # pyright: ignore[reportPrivateUsage]
-
- # Case 3: All parents finished
- manager._pending_execution = {"C": ["A", "B"]} # pyright: ignore[reportPrivateUsage]
- assert manager._is_node_ready("C") # pyright: ignore[reportPrivateUsage]
-
- # === Test "any" activation ===
- graph.nodes["C"].activation = "any"
-
- # Case 1: No parent finished
- manager._pending_execution = {"C": []} # pyright: ignore[reportPrivateUsage]
- assert not manager._is_node_ready("C") # pyright: ignore[reportPrivateUsage]
-
- # Case 2: One parent finished
- manager._pending_execution = {"C": ["B"]} # pyright: ignore[reportPrivateUsage]
- assert manager._is_node_ready("C") # pyright: ignore[reportPrivateUsage]
-
- # Case 3: All parents finished
- manager._pending_execution = {"C": ["A", "B"]} # pyright: ignore[reportPrivateUsage]
- assert manager._is_node_ready("C") # pyright: ignore[reportPrivateUsage]
-
-
@pytest.mark.asyncio
async def test_invalid_digraph_manager_cycle_without_termination() -> None:
"""Test GraphManager raises error for cyclic graph without termination condition."""
@@ -359,170 +303,6 @@ async def test_invalid_digraph_manager_cycle_without_termination() -> None:
)
-@pytest.fixture
-def digraph_manager() -> Callable[..., GraphFlowManager]:
- @patch(
- "autogen_agentchat.teams._group_chat._base_group_chat_manager.BaseGroupChatManager.__init__", return_value=None
- )
- def _create(
- _: Any,
- graph: DiGraph,
- active_nodes: Set[str] | None = None,
- thread: List[BaseAgentEvent | BaseChatMessage] | None = None,
- pending: Dict[str, List[str]] | None = None,
- ) -> GraphFlowManager:
- manager = GraphFlowManager.__new__(GraphFlowManager)
- manager._graph = graph # pyright: ignore[reportPrivateUsage]
- manager._parents = graph.get_parents() # pyright: ignore[reportPrivateUsage]
- manager._start_nodes = graph.get_start_nodes() # pyright: ignore[reportPrivateUsage]
- manager._leaf_nodes = graph.get_leaf_nodes() # pyright: ignore[reportPrivateUsage]
- manager._active_nodes = set(active_nodes or []) # pyright: ignore[reportPrivateUsage]
- manager._active_node_count = {node: 0 for node in graph.nodes} # pyright: ignore[reportPrivateUsage]
- manager._message_factory = MessageFactory() # pyright: ignore[reportPrivateUsage]
- manager._message_thread = thread if thread is not None else [] # pyright: ignore[reportPrivateUsage]
- manager._pending_execution = pending if pending is not None else {node: [] for node in graph.get_start_nodes()} # pyright: ignore[reportPrivateUsage]
- manager._name = "test_manager" # pyright: ignore[reportPrivateUsage]
- manager._use_default_start = False # pyright: ignore[reportPrivateUsage]
- return manager
-
- return _create
-
-
-# -------------------- Test: Sequential Flow --------------------
-@pytest.mark.asyncio
-async def test_select_speakers_linear(digraph_manager: Callable[..., GraphFlowManager]) -> None:
- graph = DiGraph(
- nodes={
- "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]),
- "B": DiGraphNode(name="B", edges=[DiGraphEdge(target="C")]),
- "C": DiGraphNode(name="C", edges=[]),
- }
- )
- message_thread = [TextChatMessage(source="A", content="done", metadata={})]
- manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
- result = await manager.select_speakers(manager._message_thread) # pyright: ignore[reportPrivateUsage]
- assert result == ["B"]
- assert "B" in manager._active_nodes # pyright: ignore[reportPrivateUsage]
-
-
-# -------------------- Test: Parallel Fan-out --------------------
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_parallel(digraph_manager: Callable[..., GraphFlowManager]) -> None:
- graph = DiGraph(
- nodes={
- "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B"), DiGraphEdge(target="C")]),
- "B": DiGraphNode(name="B", edges=[]),
- "C": DiGraphNode(name="C", edges=[]),
- }
- )
- message_thread = [TextChatMessage(source="A", content="done", metadata={})]
- manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
- result = await manager.select_speakers(manager._message_thread) # pyright: ignore[reportPrivateUsage]
- assert set(result) == {"B", "C"}
- assert "B" in manager._active_nodes # pyright: ignore[reportPrivateUsage]
- assert "C" in manager._active_nodes # pyright: ignore[reportPrivateUsage]
-
-
-# -------------------- Test: Conditional Path --------------------
-@pytest.mark.asyncio
-async def test_select_speakers_conditional(digraph_manager: Callable[..., GraphFlowManager]) -> None:
- graph = DiGraph(
- nodes={
- "A": DiGraphNode(
- name="A", edges=[DiGraphEdge(target="B", condition="yes"), DiGraphEdge(target="C", condition="no")]
- ),
- "B": DiGraphNode(name="B", edges=[]),
- "C": DiGraphNode(name="C", edges=[]),
- }
- )
- message_thread = [TextChatMessage(source="A", content="no", metadata={})]
- manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
- result = await manager.select_speakers(manager._message_thread) # pyright: ignore[reportPrivateUsage]
- assert result == ["C"]
- assert "C" in manager._active_nodes # pyright: ignore[reportPrivateUsage]
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_from_start_nodes(digraph_manager: Callable[..., GraphFlowManager]) -> None:
- graph = DiGraph(
- nodes={
- "A": DiGraphNode(name="A", edges=[]),
- "B": DiGraphNode(name="B", edges=[]),
- }
- )
- # No prior message — both are start nodes
- manager = digraph_manager(graph=graph, active_nodes=set(), thread=[], pending={"A": [], "B": []})
- result = await manager.select_speakers([])
- assert set(result) == {"A", "B"}
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_termination(digraph_manager: Callable[..., GraphFlowManager]) -> None:
- graph = DiGraph(
- nodes={
- "A": DiGraphNode(name="A", edges=[]),
- }
- )
-
- # Create the manager and manually patch _signal_termination to track calls
- manager = digraph_manager(
- graph=graph, active_nodes={"A"}, thread=[TextChatMessage(source="A", content="done", metadata={})], pending={}
- )
- manager._signal_termination = AsyncMock() # type: ignore[assignment]
-
- result = await manager.select_speakers(manager._message_thread) # pyright: ignore[reportPrivateUsage]
-
- # No speakers left to run, so result should be empty
- assert result == [_DIGRAPH_STOP_AGENT_NAME]
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_conditional_all_activation(
- digraph_manager: Callable[..., GraphFlowManager],
-) -> None:
- graph = DiGraph(
- nodes={
- "A": DiGraphNode(
- name="A", edges=[DiGraphEdge(target="B", condition="yes"), DiGraphEdge(target="C", condition="no")]
- ),
- "B": DiGraphNode(name="B", edges=[], activation="all"),
- "C": DiGraphNode(name="C", edges=[], activation="all"),
- }
- )
- message_thread = [TextChatMessage(source="A", content="no", metadata={})]
- manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
- result = await manager.select_speakers(manager._message_thread) # pyright: ignore[reportPrivateUsage]
- assert result == ["C"]
- assert "C" in manager._active_nodes # pyright: ignore[reportPrivateUsage]
-
-
-@pytest.mark.asyncio
-async def test_select_speakers_conditional_any_activation(
- digraph_manager: Callable[..., GraphFlowManager],
-) -> None:
- graph = DiGraph(
- nodes={
- "A": DiGraphNode(
- name="A", edges=[DiGraphEdge(target="B", condition="yes"), DiGraphEdge(target="C", condition="no")]
- ),
- "B": DiGraphNode(name="B", edges=[], activation="any"),
- "C": DiGraphNode(name="C", edges=[], activation="any"),
- }
- )
- message_thread = [TextChatMessage(source="A", content="yes", metadata={})]
- manager = digraph_manager(graph=graph, active_nodes={"A"}, thread=message_thread, pending={"B": [], "C": []})
-
- result = await manager.select_speakers(manager._message_thread) # pyright: ignore[reportPrivateUsage]
- assert result == ["B"]
- assert "B" in manager._active_nodes # pyright: ignore[reportPrivateUsage]
-
-
class _EchoAgent(BaseChatAgent):
def __init__(self, name: str, description: str) -> None:
super().__init__(name, description)
@@ -1417,7 +1197,6 @@ async def test_graph_flow_serialize_deserialize() -> None:
participants=builder.get_participants(),
graph=builder.build(),
runtime=None,
- termination_condition=MaxMessageTermination(5),
)
serialized = team.dump_component()
@@ -1439,11 +1218,52 @@ async def test_graph_flow_serialize_deserialize() -> None:
assert results.messages[1].source == "A"
assert results.messages[1].content == "0"
assert isinstance(results.messages[2], TextMessage)
- assert results.messages[2].source == "A"
- assert results.messages[2].content == "1"
- assert isinstance(results.messages[3], TextMessage)
- assert results.messages[3].source == "B"
- assert results.messages[3].content == "0"
+ assert results.messages[2].source == "B"
+ assert results.messages[2].content == "0"
assert isinstance(results.messages[-1], StopMessage)
assert results.messages[-1].source == _DIGRAPH_STOP_AGENT_NAME
assert results.messages[-1].content == "Digraph execution is complete"
+
+
+@pytest.mark.asyncio
+async def test_graph_flow_stateful_pause_and_resume_with_termination() -> None:
+ client_a = ReplayChatCompletionClient(["A1", "A2"])
+ client_b = ReplayChatCompletionClient(["B1"])
+
+ a = AssistantAgent("A", model_client=client_a)
+ b = AssistantAgent("B", model_client=client_b)
+
+ builder = DiGraphBuilder()
+ builder.add_node(a).add_node(b)
+ builder.add_edge(a, b)
+ builder.set_entry_point(a)
+
+ team = GraphFlow(
+ participants=builder.get_participants(),
+ graph=builder.build(),
+ runtime=None,
+ termination_condition=SourceMatchTermination(sources=["A"]),
+ )
+
+ result = await team.run(task="Start")
+ assert len(result.messages) == 2
+ assert result.messages[0].source == "user"
+ assert result.messages[1].source == "A"
+ assert result.stop_reason is not None and result.stop_reason == "'A' answered"
+
+ # Export state.
+ state = await team.save_state()
+
+ # Load state into a new team.
+ new_team = GraphFlow(
+ participants=builder.get_participants(),
+ graph=builder.build(),
+ runtime=None,
+ )
+ await new_team.load_state(state)
+
+ # Resume.
+ result = await new_team.run()
+ assert len(result.messages) == 2
+ assert result.messages[0].source == "B"
+ assert result.messages[1].source == _DIGRAPH_STOP_AGENT_NAME
diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb
index cf719e10d..4be55fe4b 100644
--- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb
+++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/models.ipynb
@@ -427,6 +427,58 @@
"```"
]
},
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Llama API (experimental)\n",
+ "\n",
+ "[Llama API](https://llama.developer.meta.com?utm_source=partner-autogen&utm_medium=readme) is the Meta's first party API offering. It currently offers an [OpenAI compatible endpoint](https://llama.developer.meta.com/docs/features/compatibility).\n",
+ "So you can use the {py:class}`~autogen_ext.models.openai.OpenAIChatCompletionClient` with the Llama API.\n",
+ "\n",
+ "This endpoint fully supports the following OpenAI client library features:\n",
+ "* Chat completions\n",
+ "* Model selection\n",
+ "* Temperature/sampling\n",
+ "* Streaming\n",
+ "* Image understanding\n",
+ "* Structured output (JSON mode)\n",
+ "* Funciton calling (tools)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from autogen_core.models import UserMessage\n",
+ "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
+ "from autogen_core import Image\n",
+ "from pathlib import Path\n",
+ "\n",
+ "# Text\n",
+ "model_client = OpenAIChatCompletionClient(\n",
+ " model=\"Llama-4-Scout-17B-16E-Instruct-FP8\",\n",
+ " # api_key=\"LLAMA_API_KEY\"\n",
+ ")\n",
+ "\n",
+ "response = await model_client.create([UserMessage(content=\"Write me a poem\", source=\"user\")])\n",
+ "print(response)\n",
+ "await model_client.close()\n",
+ "\n",
+ "# Image\n",
+ "model_client = OpenAIChatCompletionClient(\n",
+ " model=\"Llama-4-Maverick-17B-128E-Instruct-FP8\",\n",
+ " # api_key=\"LLAMA_API_KEY\"\n",
+ ")\n",
+ "image = Image.from_file(Path(\"test.png\"))\n",
+ "\n",
+ "response = await model_client.create([UserMessage(content=[\"What is in this image\", image], source=\"user\")])\n",
+ "print(response)\n",
+ "await model_client.close()"
+ ]
+ },
{
"cell_type": "markdown",
"metadata": {},
diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb
index 05e799eb1..636841d08 100644
--- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb
+++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/state.ipynb
@@ -112,7 +112,7 @@
"source": [
"```{note}\n",
"For {py:class}`~autogen_agentchat.agents.AssistantAgent`, its state consists of the model_context.\n",
- "If your write your own custom agent, consider overriding the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.save_state` and {py:meth}`~autogen_agentchat.agents.BaseChatAgent.load_state` methods to customize the behavior. The default implementations save and load an empty state.\n",
+ "If you write your own custom agent, consider overriding the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.save_state` and {py:meth}`~autogen_agentchat.agents.BaseChatAgent.load_state` methods to customize the behavior. The default implementations save and load an empty state.\n",
"```"
]
},
diff --git a/python/packages/autogen-core/pyproject.toml b/python/packages/autogen-core/pyproject.toml
index f1f6ecf8b..7963a0907 100644
--- a/python/packages/autogen-core/pyproject.toml
+++ b/python/packages/autogen-core/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "autogen-core"
-version = "0.5.6"
+version = "0.5.7"
license = {file = "LICENSE-CODE"}
description = "Foundational interfaces and agent runtime implementation for AutoGen"
readme = "README.md"
@@ -69,7 +69,7 @@ dev = [
"pygments",
"sphinxext-rediraffe",
- "autogen_ext==0.5.6",
+ "autogen_ext==0.5.7",
# Documentation tooling
"diskcache",
diff --git a/python/packages/autogen-core/src/autogen_core/_agent.py b/python/packages/autogen-core/src/autogen_core/_agent.py
index 0f37b822f..e407fe137 100644
--- a/python/packages/autogen-core/src/autogen_core/_agent.py
+++ b/python/packages/autogen-core/src/autogen_core/_agent.py
@@ -1,9 +1,13 @@
-from typing import Any, Mapping, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Any, Mapping, Protocol, runtime_checkable
from ._agent_id import AgentId
from ._agent_metadata import AgentMetadata
from ._message_context import MessageContext
+# Forward declaration for type checking only
+if TYPE_CHECKING:
+ from ._agent_runtime import AgentRuntime
+
@runtime_checkable
class Agent(Protocol):
@@ -17,6 +21,15 @@ class Agent(Protocol):
"""ID of the agent."""
...
+ async def bind_id_and_runtime(self, id: AgentId, runtime: "AgentRuntime") -> None:
+ """Function used to bind an Agent instance to an `AgentRuntime`.
+
+ Args:
+ agent_id (AgentId): ID of the agent.
+ runtime (AgentRuntime): AgentRuntime instance to bind the agent to.
+ """
+ ...
+
async def on_message(self, message: Any, ctx: MessageContext) -> Any:
"""Message handler for the agent. This should only be called by the runtime, not by other agents.
diff --git a/python/packages/autogen-core/src/autogen_core/_agent_instantiation.py b/python/packages/autogen-core/src/autogen_core/_agent_instantiation.py
index 71921225c..a8904a42d 100644
--- a/python/packages/autogen-core/src/autogen_core/_agent_instantiation.py
+++ b/python/packages/autogen-core/src/autogen_core/_agent_instantiation.py
@@ -118,3 +118,9 @@ class AgentInstantiationContext:
raise RuntimeError(
"AgentInstantiationContext.agent_id() must be called within an instantiation context such as when the AgentRuntime is instantiating an agent. Mostly likely this was caused by directly instantiating an agent instead of using the AgentRuntime to do so."
) from e
+
+ @classmethod
+ def is_in_factory_call(cls) -> bool:
+ if cls._AGENT_INSTANTIATION_CONTEXT_VAR.get(None) is None:
+ return False
+ return True
diff --git a/python/packages/autogen-core/src/autogen_core/_agent_runtime.py b/python/packages/autogen-core/src/autogen_core/_agent_runtime.py
index 22493626f..d4bac4a9c 100644
--- a/python/packages/autogen-core/src/autogen_core/_agent_runtime.py
+++ b/python/packages/autogen-core/src/autogen_core/_agent_runtime.py
@@ -130,6 +130,60 @@ class AgentRuntime(Protocol):
"""
...
+ async def register_agent_instance(
+ self,
+ agent_instance: Agent,
+ agent_id: AgentId,
+ ) -> AgentId:
+ """Register an agent instance with the runtime. The type may be reused, but each agent_id must be unique. All agent instances within a type must be of the same object type. This API does not add any subscriptions.
+
+ .. note::
+
+ This is a low level API and usually the agent class's `register_instance` method should be used instead, as this also handles subscriptions automatically.
+
+ Example:
+
+ .. code-block:: python
+
+ from dataclasses import dataclass
+
+ from autogen_core import AgentId, AgentRuntime, MessageContext, RoutedAgent, event
+ from autogen_core.models import UserMessage
+
+
+ @dataclass
+ class MyMessage:
+ content: str
+
+
+ class MyAgent(RoutedAgent):
+ def __init__(self) -> None:
+ super().__init__("My core agent")
+
+ @event
+ async def handler(self, message: UserMessage, context: MessageContext) -> None:
+ print("Event received: ", message.content)
+
+
+ async def main() -> None:
+ runtime: AgentRuntime = ... # type: ignore
+ agent = MyAgent()
+ await runtime.register_agent_instance(
+ agent_instance=agent, agent_id=AgentId(type="my_agent", key="default")
+ )
+
+
+ import asyncio
+
+ asyncio.run(main())
+
+
+ Args:
+ agent_instance (Agent): A concrete instance of the agent.
+ agent_id (AgentId): The agent's identifier. The agent's type is `agent_id.type`.
+ """
+ ...
+
# TODO: uncomment out the following type ignore when this is fixed in mypy: https://github.com/python/mypy/issues/3737
async def try_get_underlying_agent_instance(self, id: AgentId, type: Type[T] = Agent) -> T: # type: ignore[assignment]
"""Try to get the underlying agent instance by name and namespace. This is generally discouraged (hence the long name), but can be useful in some cases.
diff --git a/python/packages/autogen-core/src/autogen_core/_base_agent.py b/python/packages/autogen-core/src/autogen_core/_base_agent.py
index bffb61b87..0ad0bc607 100644
--- a/python/packages/autogen-core/src/autogen_core/_base_agent.py
+++ b/python/packages/autogen-core/src/autogen_core/_base_agent.py
@@ -21,6 +21,7 @@ from ._subscription import Subscription, UnboundSubscription
from ._subscription_context import SubscriptionInstantiationContext
from ._topic import TopicId
from ._type_prefix_subscription import TypePrefixSubscription
+from ._type_subscription import TypeSubscription
T = TypeVar("T", bound=Agent)
@@ -82,20 +83,25 @@ class BaseAgent(ABC, Agent):
return AgentMetadata(key=self._id.key, type=self._id.type, description=self._description)
def __init__(self, description: str) -> None:
- try:
- runtime = AgentInstantiationContext.current_runtime()
- id = AgentInstantiationContext.current_agent_id()
- except LookupError as e:
- raise RuntimeError(
- "BaseAgent must be instantiated within the context of an AgentRuntime. It cannot be directly instantiated."
- ) from e
-
- self._runtime: AgentRuntime = runtime
- self._id: AgentId = id
+ if AgentInstantiationContext.is_in_factory_call():
+ self._runtime: AgentRuntime = AgentInstantiationContext.current_runtime()
+ self._id = AgentInstantiationContext.current_agent_id()
if not isinstance(description, str):
raise ValueError("Agent description must be a string")
self._description = description
+ async def bind_id_and_runtime(self, id: AgentId, runtime: AgentRuntime) -> None:
+ if hasattr(self, "_id"):
+ if self._id != id:
+ raise RuntimeError("Agent is already bound to a different ID")
+
+ if hasattr(self, "_runtime"):
+ if self._runtime != runtime:
+ raise RuntimeError("Agent is already bound to a different runtime")
+
+ self._id = id
+ self._runtime = runtime
+
@property
def type(self) -> str:
return self.id.type
@@ -155,6 +161,56 @@ class BaseAgent(ABC, Agent):
async def close(self) -> None:
pass
+ async def register_instance(
+ self,
+ runtime: AgentRuntime,
+ agent_id: AgentId,
+ *,
+ skip_class_subscriptions: bool = True,
+ skip_direct_message_subscription: bool = False,
+ ) -> AgentId:
+ """
+ This function is similar to `register` but is used for registering an instance of an agent. A subscription based on the agent ID is created and added to the runtime.
+ """
+ agent_id = await runtime.register_agent_instance(agent_instance=self, agent_id=agent_id)
+
+ id_subscription = TypeSubscription(topic_type=agent_id.key, agent_type=agent_id.type)
+ await runtime.add_subscription(id_subscription)
+
+ if not skip_class_subscriptions:
+ with SubscriptionInstantiationContext.populate_context(AgentType(agent_id.type)):
+ subscriptions: List[Subscription] = []
+ for unbound_subscription in self._unbound_subscriptions():
+ subscriptions_list_result = unbound_subscription()
+ if inspect.isawaitable(subscriptions_list_result):
+ subscriptions_list = await subscriptions_list_result
+ else:
+ subscriptions_list = subscriptions_list_result
+
+ subscriptions.extend(subscriptions_list)
+ for subscription in subscriptions:
+ await runtime.add_subscription(subscription)
+
+ if not skip_direct_message_subscription:
+ # Additionally adds a special prefix subscription for this agent to receive direct messages
+ try:
+ await runtime.add_subscription(
+ TypePrefixSubscription(
+ # The prefix MUST include ":" to avoid collisions with other agents
+ topic_type_prefix=agent_id.type + ":",
+ agent_type=agent_id.type,
+ )
+ )
+ except ValueError:
+ # We don't care if the subscription already exists
+ pass
+
+ # TODO: deduplication
+ for _message_type, serializer in self._handles_types():
+ runtime.add_message_serializer(serializer)
+
+ return agent_id
+
@classmethod
async def register(
cls,
diff --git a/python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py b/python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py
index 9610e7f54..5f5636441 100644
--- a/python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py
+++ b/python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py
@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import inspect
+import json
import logging
import sys
import uuid
@@ -265,6 +266,7 @@ class SingleThreadedAgentRuntime(AgentRuntime):
self._serialization_registry = SerializationRegistry()
self._ignore_unhandled_handler_exceptions = ignore_unhandled_exceptions
self._background_exception: BaseException | None = None
+ self._agent_instance_types: Dict[str, Type[Agent]] = {}
@property
def unprocessed_messages_count(
@@ -276,6 +278,55 @@ class SingleThreadedAgentRuntime(AgentRuntime):
def _known_agent_names(self) -> Set[str]:
return set(self._agent_factories.keys())
+ async def _create_otel_attributes(
+ self,
+ sender_agent_id: AgentId | None = None,
+ recipient_agent_id: AgentId | None = None,
+ message_context: MessageContext | None = None,
+ message: Any = None,
+ ) -> Mapping[str, str]:
+ """Create OpenTelemetry attributes for the given agent and message.
+
+ Args:
+ sender_agent (Agent, optional): The sender agent instance.
+ recipient_agent (Agent, optional): The recipient agent instance.
+ message (Any): The message instance.
+
+ Returns:
+ Attributes: A dictionary of OpenTelemetry attributes.
+ """
+ if not sender_agent_id and not recipient_agent_id and not message:
+ return {}
+ attributes: Dict[str, str] = {}
+ if sender_agent_id:
+ sender_agent = await self._get_agent(sender_agent_id)
+ attributes["sender_agent_type"] = sender_agent.id.type
+ attributes["sender_agent_class"] = sender_agent.__class__.__name__
+ if recipient_agent_id:
+ recipient_agent = await self._get_agent(recipient_agent_id)
+ attributes["recipient_agent_type"] = recipient_agent.id.type
+ attributes["recipient_agent_class"] = recipient_agent.__class__.__name__
+
+ if message_context:
+ serialized_message_context = {
+ "sender": str(message_context.sender),
+ "topic_id": str(message_context.topic_id),
+ "is_rpc": message_context.is_rpc,
+ "message_id": message_context.message_id,
+ }
+ attributes["message_context"] = json.dumps(serialized_message_context)
+
+ if message:
+ try:
+ serialized_message = self._try_serialize(message)
+ except Exception as e:
+ serialized_message = str(e)
+ else:
+ serialized_message = "No Message"
+ attributes["message"] = serialized_message
+
+ return attributes
+
# Returns the response of the message
async def send_message(
self,
@@ -311,6 +362,7 @@ class SingleThreadedAgentRuntime(AgentRuntime):
future = asyncio.get_event_loop().create_future()
if recipient.type not in self._known_agent_names:
future.set_exception(Exception("Recipient not found"))
+ return await future
content = message.__dict__ if hasattr(message, "__dict__") else message
logger.info(f"Sending message of type {type(message).__name__} to {recipient.type}: {content}")
@@ -440,7 +492,17 @@ class SingleThreadedAgentRuntime(AgentRuntime):
cancellation_token=message_envelope.cancellation_token,
message_id=message_envelope.message_id,
)
- with self._tracer_helper.trace_block("process", recipient_agent.id, parent=message_envelope.metadata):
+ with self._tracer_helper.trace_block(
+ "process",
+ recipient_agent.id,
+ parent=message_envelope.metadata,
+ attributes=await self._create_otel_attributes(
+ sender_agent_id=message_envelope.sender,
+ recipient_agent_id=recipient,
+ message_context=message_context,
+ message=message_envelope.message,
+ ),
+ ):
with MessageHandlerContext.populate_context(recipient_agent.id):
response = await recipient_agent.on_message(
message_envelope.message,
@@ -527,7 +589,17 @@ class SingleThreadedAgentRuntime(AgentRuntime):
agent = await self._get_agent(agent_id)
async def _on_message(agent: Agent, message_context: MessageContext) -> Any:
- with self._tracer_helper.trace_block("process", agent.id, parent=message_envelope.metadata):
+ with self._tracer_helper.trace_block(
+ "process",
+ agent.id,
+ parent=message_envelope.metadata,
+ attributes=await self._create_otel_attributes(
+ sender_agent_id=message_envelope.sender,
+ recipient_agent_id=agent.id,
+ message_context=message_context,
+ message=message_envelope.message,
+ ),
+ ):
with MessageHandlerContext.populate_context(agent.id):
try:
return await agent.on_message(
@@ -557,7 +629,16 @@ class SingleThreadedAgentRuntime(AgentRuntime):
# TODO if responses are given for a publish
async def _process_response(self, message_envelope: ResponseMessageEnvelope) -> None:
- with self._tracer_helper.trace_block("ack", message_envelope.recipient, parent=message_envelope.metadata):
+ with self._tracer_helper.trace_block(
+ "ack",
+ message_envelope.recipient,
+ parent=message_envelope.metadata,
+ attributes=await self._create_otel_attributes(
+ sender_agent_id=message_envelope.sender,
+ recipient_agent_id=message_envelope.recipient,
+ message=message_envelope.message,
+ ),
+ ):
content = (
message_envelope.message.__dict__
if hasattr(message_envelope.message, "__dict__")
@@ -830,6 +911,32 @@ class SingleThreadedAgentRuntime(AgentRuntime):
return type
+ async def register_agent_instance(
+ self,
+ agent_instance: Agent,
+ agent_id: AgentId,
+ ) -> AgentId:
+ def agent_factory() -> Agent:
+ raise RuntimeError(
+ "Agent factory was invoked for an agent instance that was not registered. This is likely due to the agent type being incorrectly subscribed to a topic. If this exception occurs when publishing a message to the DefaultTopicId, then it is likely that `skip_class_subscriptions` needs to be turned off when registering the agent."
+ )
+
+ if agent_id in self._instantiated_agents:
+ raise ValueError(f"Agent with id {agent_id} already exists.")
+
+ if agent_id.type not in self._agent_factories:
+ self._agent_factories[agent_id.type] = agent_factory
+ self._agent_instance_types[agent_id.type] = type_func_alias(agent_instance)
+ else:
+ if self._agent_factories[agent_id.type].__code__ != agent_factory.__code__:
+ raise ValueError("Agent factories and agent instances cannot be registered to the same type.")
+ if self._agent_instance_types[agent_id.type] != type_func_alias(agent_instance):
+ raise ValueError("Agent instances must be the same object type.")
+
+ await agent_instance.bind_id_and_runtime(id=agent_id, runtime=self)
+ self._instantiated_agents[agent_id] = agent_instance
+ return agent_id
+
async def _invoke_agent_factory(
self,
agent_factory: Callable[[], T | Awaitable[T]] | Callable[[AgentRuntime, AgentId], T | Awaitable[T]],
@@ -851,8 +958,7 @@ class SingleThreadedAgentRuntime(AgentRuntime):
raise ValueError("Agent factory must take 0 or 2 arguments.")
if inspect.isawaitable(agent):
- return cast(T, await agent)
-
+ agent = cast(T, await agent)
return agent
except BaseException as e:
diff --git a/python/packages/autogen-core/src/autogen_core/models/_model_client.py b/python/packages/autogen-core/src/autogen_core/models/_model_client.py
index b65a19853..7c8519ddf 100644
--- a/python/packages/autogen-core/src/autogen_core/models/_model_client.py
+++ b/python/packages/autogen-core/src/autogen_core/models/_model_client.py
@@ -37,9 +37,19 @@ class ModelFamily:
CLAUDE_3_5_HAIKU = "claude-3-5-haiku"
CLAUDE_3_5_SONNET = "claude-3-5-sonnet"
CLAUDE_3_7_SONNET = "claude-3-7-sonnet"
+ LLAMA_3_3_8B = "llama-3.3-8b"
+ LLAMA_3_3_70B = "llama-3.3-70b"
+ LLAMA_4_SCOUT = "llama-4-scout"
+ LLAMA_4_MAVERICK = "llama-4-maverick"
+ CODESRAL = "codestral"
+ OPEN_CODESRAL_MAMBA = "open-codestral-mamba"
+ MISTRAL = "mistral"
+ MINISTRAL = "ministral"
+ PIXTRAL = "pixtral"
UNKNOWN = "unknown"
ANY: TypeAlias = Literal[
+ # openai_models
"gpt-41",
"gpt-45",
"gpt-4o",
@@ -49,16 +59,30 @@ class ModelFamily:
"gpt-4",
"gpt-35",
"r1",
+ # google_models
"gemini-1.5-flash",
"gemini-1.5-pro",
"gemini-2.0-flash",
"gemini-2.5-pro",
+ # anthropic_models
"claude-3-haiku",
"claude-3-sonnet",
"claude-3-opus",
"claude-3-5-haiku",
"claude-3-5-sonnet",
"claude-3-7-sonnet",
+ # llama_models
+ "llama-3.3-8b",
+ "llama-3.3-70b",
+ "llama-4-scout",
+ "llama-4-maverick",
+ # mistral_models
+ "codestral",
+ "open-codestral-mamba",
+ "mistral",
+ "ministral",
+ "pixtral",
+ # unknown
"unknown",
]
@@ -98,6 +122,25 @@ class ModelFamily:
ModelFamily.GPT_35,
)
+ @staticmethod
+ def is_llama(family: str) -> bool:
+ return family in (
+ ModelFamily.LLAMA_3_3_8B,
+ ModelFamily.LLAMA_3_3_70B,
+ ModelFamily.LLAMA_4_SCOUT,
+ ModelFamily.LLAMA_4_MAVERICK,
+ )
+
+ @staticmethod
+ def is_mistral(family: str) -> bool:
+ return family in (
+ ModelFamily.CODESRAL,
+ ModelFamily.OPEN_CODESRAL_MAMBA,
+ ModelFamily.MISTRAL,
+ ModelFamily.MINISTRAL,
+ ModelFamily.PIXTRAL,
+ )
+
@deprecated("Use the ModelInfo class instead ModelCapabilities.")
class ModelCapabilities(TypedDict, total=False):
diff --git a/python/packages/autogen-core/src/autogen_core/tools/_static_workbench.py b/python/packages/autogen-core/src/autogen_core/tools/_static_workbench.py
index 2e762232e..1771e7e48 100644
--- a/python/packages/autogen-core/src/autogen_core/tools/_static_workbench.py
+++ b/python/packages/autogen-core/src/autogen_core/tools/_static_workbench.py
@@ -1,4 +1,5 @@
import asyncio
+import builtins
from typing import Any, Dict, List, Literal, Mapping
from pydantic import BaseModel
@@ -55,12 +56,12 @@ class StaticWorkbench(Workbench, Component[StaticWorkbenchConfig]):
try:
result_future = asyncio.ensure_future(tool.run_json(arguments, cancellation_token))
cancellation_token.link_future(result_future)
- result = await result_future
+ actual_tool_output = await result_future
is_error = False
+ result_str = tool.return_value_as_string(actual_tool_output)
except Exception as e:
- result = str(e)
+ result_str = self._format_errors(e)
is_error = True
- result_str = tool.return_value_as_string(result)
return ToolResult(name=tool.name, result=[TextResultContent(content=result_str)], is_error=is_error)
async def start(self) -> None:
@@ -90,3 +91,16 @@ class StaticWorkbench(Workbench, Component[StaticWorkbenchConfig]):
@classmethod
def _from_config(cls, config: StaticWorkbenchConfig) -> Self:
return cls(tools=[BaseTool.load_component(tool) for tool in config.tools])
+
+ def _format_errors(self, error: Exception) -> str:
+ """Recursively format errors into a string."""
+
+ error_message = ""
+ if hasattr(builtins, "ExceptionGroup") and isinstance(error, builtins.ExceptionGroup):
+ # ExceptionGroup is available in Python 3.11+.
+ # TODO: how to make this compatible with Python 3.10?
+ for sub_exception in error.exceptions: # type: ignore
+ error_message += self._format_errors(sub_exception) # type: ignore
+ else:
+ error_message += f"{str(error)}\n"
+ return error_message.strip()
diff --git a/python/packages/autogen-core/tests/test_base_agent.py b/python/packages/autogen-core/tests/test_base_agent.py
index 64bcf59d1..010bd0624 100644
--- a/python/packages/autogen-core/tests/test_base_agent.py
+++ b/python/packages/autogen-core/tests/test_base_agent.py
@@ -9,7 +9,7 @@ async def test_base_agent_create(mocker: MockerFixture) -> None:
runtime = mocker.Mock(spec=AgentRuntime)
# Shows how to set the context for the agent instantiation in a test context
- with AgentInstantiationContext.populate_context((runtime, AgentId("name", "namespace"))):
- agent = NoopAgent()
- assert agent.runtime == runtime
- assert agent.id == AgentId("name", "namespace")
+ with AgentInstantiationContext.populate_context((runtime, AgentId("name2", "namespace2"))):
+ agent2 = NoopAgent()
+ assert agent2.runtime == runtime
+ assert agent2.id == AgentId("name2", "namespace2")
diff --git a/python/packages/autogen-core/tests/test_runtime.py b/python/packages/autogen-core/tests/test_runtime.py
index 64a1cccf4..e93a57a6a 100644
--- a/python/packages/autogen-core/tests/test_runtime.py
+++ b/python/packages/autogen-core/tests/test_runtime.py
@@ -82,6 +82,60 @@ async def test_agent_type_must_be_unique() -> None:
await runtime.register_factory(type=AgentType("name2"), agent_factory=agent_factory, expected_class=NoopAgent)
+@pytest.mark.asyncio
+async def test_agent_type_register_instance() -> None:
+ runtime = SingleThreadedAgentRuntime()
+ agent1_id = AgentId(type="name", key="default")
+ agent2_id = AgentId(type="name", key="notdefault")
+ agent1 = NoopAgent()
+ agent1_dup = NoopAgent()
+ agent2 = NoopAgent()
+ await agent1.register_instance(runtime=runtime, agent_id=agent1_id)
+ await agent2.register_instance(runtime=runtime, agent_id=agent2_id)
+
+ assert await runtime.try_get_underlying_agent_instance(agent1_id, type=NoopAgent) == agent1
+ assert await runtime.try_get_underlying_agent_instance(agent2_id, type=NoopAgent) == agent2
+ with pytest.raises(ValueError):
+ await agent1_dup.register_instance(runtime=runtime, agent_id=agent1_id)
+
+
+@pytest.mark.asyncio
+async def test_agent_type_register_instance_different_types() -> None:
+ runtime = SingleThreadedAgentRuntime()
+ agent_id1 = AgentId(type="name", key="noop")
+ agent_id2 = AgentId(type="name", key="loopback")
+ agent1 = NoopAgent()
+ agent2 = LoopbackAgent()
+ await agent1.register_instance(runtime=runtime, agent_id=agent_id1)
+ with pytest.raises(ValueError):
+ await agent2.register_instance(runtime=runtime, agent_id=agent_id2)
+
+
+@pytest.mark.asyncio
+async def test_agent_type_register_instance_publish_new_source() -> None:
+ runtime = SingleThreadedAgentRuntime(ignore_unhandled_exceptions=False)
+ agent_id = AgentId(type="name", key="default")
+ agent1 = LoopbackAgent()
+ await agent1.register_instance(runtime=runtime, agent_id=agent_id)
+ await runtime.add_subscription(TypeSubscription("notdefault", "name"))
+
+ runtime.start()
+ with pytest.raises(RuntimeError):
+ await runtime.publish_message(MessageType(), TopicId("notdefault", "notdefault"))
+ await runtime.stop_when_idle()
+ await runtime.close()
+
+
+@pytest.mark.asyncio
+async def test_register_instance_factory() -> None:
+ runtime = SingleThreadedAgentRuntime()
+ agent1_id = AgentId(type="name", key="default")
+ agent1 = NoopAgent()
+ await agent1.register_instance(runtime=runtime, agent_id=agent1_id)
+ with pytest.raises(ValueError):
+ await NoopAgent.register(runtime, "name", lambda: NoopAgent())
+
+
@pytest.mark.asyncio
async def test_register_receives_publish(tracer_provider: TracerProvider) -> None:
runtime = SingleThreadedAgentRuntime(tracer_provider=tracer_provider)
diff --git a/python/packages/autogen-ext/pyproject.toml b/python/packages/autogen-ext/pyproject.toml
index 683f0814c..a86156591 100644
--- a/python/packages/autogen-ext/pyproject.toml
+++ b/python/packages/autogen-ext/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "autogen-ext"
-version = "0.5.6"
+version = "0.5.7"
license = {file = "LICENSE-CODE"}
description = "AutoGen extensions library"
readme = "README.md"
@@ -15,7 +15,7 @@ classifiers = [
"Operating System :: OS Independent",
]
dependencies = [
- "autogen-core==0.5.6",
+ "autogen-core==0.5.7",
]
[project.optional-dependencies]
@@ -32,7 +32,7 @@ docker = ["docker~=7.0", "asyncio_atexit>=1.0.1"]
ollama = ["ollama>=0.4.7", "tiktoken>=0.8.0"]
openai = ["openai>=1.66.5", "tiktoken>=0.8.0", "aiofiles"]
file-surfer = [
- "autogen-agentchat==0.5.6",
+ "autogen-agentchat==0.5.7",
"magika>=0.6.1rc2",
"markitdown[all]~=0.1.0a3",
]
@@ -44,21 +44,21 @@ llama-cpp = [
graphrag = ["graphrag>=1.0.1"]
chromadb = ["chromadb>=1.0.0"]
web-surfer = [
- "autogen-agentchat==0.5.6",
+ "autogen-agentchat==0.5.7",
"playwright>=1.48.0",
"pillow>=11.0.0",
"magika>=0.6.1rc2",
"markitdown[all]~=0.1.0a3",
]
magentic-one = [
- "autogen-agentchat==0.5.6",
+ "autogen-agentchat==0.5.7",
"magika>=0.6.1rc2",
"markitdown[all]~=0.1.0a3",
"playwright>=1.48.0",
"pillow>=11.0.0",
]
video-surfer = [
- "autogen-agentchat==0.5.6",
+ "autogen-agentchat==0.5.7",
"opencv-python>=4.5",
"ffmpeg-python",
"openai-whisper",
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py
index b34cec716..f31e7b1c0 100644
--- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py
+++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/__init__.py
@@ -1,4 +1,8 @@
-from ._anthropic_client import AnthropicChatCompletionClient, BaseAnthropicChatCompletionClient
+from ._anthropic_client import (
+ AnthropicBedrockChatCompletionClient,
+ AnthropicChatCompletionClient,
+ BaseAnthropicChatCompletionClient,
+)
from .config import (
AnthropicBedrockClientConfiguration,
AnthropicBedrockClientConfigurationConfigModel,
@@ -10,6 +14,7 @@ from .config import (
__all__ = [
"AnthropicChatCompletionClient",
+ "AnthropicBedrockChatCompletionClient",
"BaseAnthropicChatCompletionClient",
"AnthropicClientConfiguration",
"AnthropicBedrockClientConfiguration",
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py
index decafc94c..5dff5813c 100644
--- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py
+++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py
@@ -5,8 +5,6 @@ import json
import logging
import re
import warnings
-
-# from asyncio import Task
from typing import (
Any,
AsyncGenerator,
@@ -25,7 +23,7 @@ from typing import (
)
import tiktoken
-from anthropic import AnthropicBedrock, AsyncAnthropic, AsyncStream
+from anthropic import AsyncAnthropic, AsyncAnthropicBedrock, AsyncStream
from anthropic.types import (
Base64ImageSourceParam,
ContentBlock,
@@ -1121,7 +1119,7 @@ class AnthropicBedrockChatCompletionClient(
BaseAnthropicChatCompletionClient, Component[AnthropicBedrockClientConfigurationConfigModel]
):
"""
- Chat completion client for Anthropic's Claude models.
+ Chat completion client for Anthropic's Claude models on AWS Bedrock.
Args:
model (str): The Claude model to use (e.g., "claude-3-sonnet-20240229", "claude-3-opus-20240229")
@@ -1145,28 +1143,24 @@ class AnthropicBedrockChatCompletionClient(
.. code-block:: python
import asyncio
- from autogen_ext.models.anthropic import AnthropicBedrockChatCompletionClient
- from autogen_core.models import UserMessage
+ from autogen_ext.models.anthropic import AnthropicBedrockChatCompletionClient, BedrockInfo
+ from autogen_core.models import UserMessage, ModelInfo
async def main():
- config = {
- "model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
- "temperature": 0.1,
- "model_info": {
- "vision": True,
- "function_calling": True,
- "json_output": True,
- "family": ModelFamily.CLAUDE_3_5_SONNET,
- },
- "bedrock_info": {
- "aws_access_key": "",
- "aws_secret_key": "",
- "aws_session_token": "",
- "aws_region": "",
- },
- }
- anthropic_client = AnthropicBedrockChatCompletionClient(**config)
+ anthropic_client = AnthropicBedrockChatCompletionClient(
+ model="anthropic.claude-3-5-sonnet-20240620-v1:0",
+ temperature=0.1,
+ model_info=ModelInfo(
+ vision=False, function_calling=True, json_output=False, family="unknown", structured_output=True
+ ),
+ bedrock_info=BedrockInfo(
+ aws_access_key="",
+ aws_secret_key="",
+ aws_session_token="",
+ aws_region="",
+ ),
+ )
result = await anthropic_client.create([UserMessage(content="What is the capital of France?", source="user")]) # type: ignore
print(result)
@@ -1178,11 +1172,11 @@ class AnthropicBedrockChatCompletionClient(
component_type = "model"
component_config_schema = AnthropicBedrockClientConfigurationConfigModel
- component_provider_override = "autogen_ext.models.anthropic.AnthropicChatCompletionClient"
+ component_provider_override = "autogen_ext.models.anthropic.AnthropicBedrockChatCompletionClient"
def __init__(self, **kwargs: Unpack[AnthropicBedrockClientConfiguration]):
if "model" not in kwargs:
- raise ValueError("model is required for AnthropicChatCompletionClient")
+ raise ValueError("model is required for AnthropicBedrockChatCompletionClient")
self._raw_config: Dict[str, Any] = dict(kwargs).copy()
copied_args = dict(kwargs).copy()
@@ -1201,11 +1195,11 @@ class AnthropicBedrockChatCompletionClient(
# Handle bedrock_info as secretestr
aws_region = bedrock_info["aws_region"]
- aws_access_key = bedrock_info["aws_access_key"].get_secret_value()
- aws_secret_key = bedrock_info["aws_secret_key"].get_secret_value()
- aws_session_token = bedrock_info["aws_session_token"].get_secret_value()
+ aws_access_key = bedrock_info["aws_access_key"]
+ aws_secret_key = bedrock_info["aws_secret_key"]
+ aws_session_token = bedrock_info["aws_session_token"]
- client = AnthropicBedrock(
+ client = AsyncAnthropicBedrock(
aws_access_key=aws_access_key,
aws_secret_key=aws_secret_key,
aws_session_token=aws_session_token,
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py
index eb2e92b0a..22dfc067f 100644
--- a/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py
+++ b/python/packages/autogen-ext/src/autogen_ext/models/anthropic/config/__init__.py
@@ -26,11 +26,11 @@ class BedrockInfo(TypedDict):
"""
- aws_access_key: Required[SecretStr]
+ aws_access_key: Required[str]
"""Access key for the aws account to gain bedrock model access"""
- aws_secret_key: Required[SecretStr]
+ aws_secret_key: Required[str]
"""Access secret key for the aws account to gain bedrock model access"""
- aws_session_token: Required[SecretStr]
+ aws_session_token: Required[str]
"""aws session token for the aws account to gain bedrock model access"""
aws_region: Required[str]
"""aws region for the aws account to gain bedrock model access"""
@@ -83,5 +83,15 @@ class AnthropicClientConfigurationConfigModel(BaseAnthropicClientConfigurationCo
tool_choice: Union[Literal["auto", "any", "none"], Dict[str, Any]] | None = None
+class BedrockInfoConfigModel(TypedDict):
+ aws_access_key: Required[SecretStr]
+ """Access key for the aws account to gain bedrock model access"""
+ aws_session_token: Required[SecretStr]
+ """aws session token for the aws account to gain bedrock model access"""
+ aws_region: Required[str]
+ """aws region for the aws account to gain bedrock model access"""
+ aws_secret_key: Required[SecretStr]
+
+
class AnthropicBedrockClientConfigurationConfigModel(AnthropicClientConfigurationConfigModel):
- bedrock_info: BedrockInfo | None = None
+ bedrock_info: BedrockInfoConfigModel | None = None
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/ollama/_model_info.py b/python/packages/autogen-ext/src/autogen_ext/models/ollama/_model_info.py
index 158a0008f..6ab239c95 100644
--- a/python/packages/autogen-ext/src/autogen_ext/models/ollama/_model_info.py
+++ b/python/packages/autogen-ext/src/autogen_ext/models/ollama/_model_info.py
@@ -258,6 +258,13 @@ _MODEL_INFO: Dict[str, ModelInfo] = {
"family": ModelFamily.UNKNOWN,
"structured_output": True,
},
+ "qwen3": {
+ "vision": False,
+ "function_calling": True,
+ "json_output": True,
+ "family": ModelFamily.UNKNOWN,
+ "structured_output": True,
+ },
"snowflake-arctic-embed": {
"vision": False,
"function_calling": False,
@@ -351,6 +358,15 @@ _MODEL_TOKEN_LIMITS: Dict[str, int] = {
"qwen2.5-coder:0.5b": 32768,
"qwen2.5-coder:1.5b": 32768,
"qwen2.5-coder:3b": 32768,
+ "qwen3": 40960,
+ "qwen3:0.6b": 40960,
+ "qwen3:1.7b": 40960,
+ "qwen3:4b": 40960,
+ "qwen3:8b": 40960,
+ "qwen3:14b": 40960,
+ "qwen3:30b": 40960,
+ "qwen3:32b": 40960,
+ "qwen3:235b": 40960,
"snowflake-arctic-embed": 512,
"starcoder2": 16384,
"tinyllama": 2048,
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_message_transform.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_message_transform.py
index 04bd93718..d3f1f993b 100644
--- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_message_transform.py
+++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_message_transform.py
@@ -275,7 +275,6 @@ base_system_message_transformers: List[Callable[[LLMMessage, Dict[str, Any]], Di
base_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = [
_assert_valid_name,
- _set_name,
_set_role("user"),
]
@@ -293,6 +292,7 @@ system_message_transformers: List[Callable[[LLMMessage, Dict[str, Any]], Dict[st
single_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = (
base_user_transformer_funcs
+ [
+ _set_name,
_set_prepend_text_content,
]
)
@@ -300,6 +300,7 @@ single_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[
multimodal_user_transformer_funcs: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = (
base_user_transformer_funcs
+ [
+ _set_name,
_set_multimodal_content,
]
)
@@ -334,6 +335,19 @@ thought_assistant_transformer_funcs_gemini: List[Callable[[LLMMessage, Dict[str,
# === Specific message param functions ===
+single_user_transformer_funcs_mistral: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = (
+ base_user_transformer_funcs
+ + [
+ _set_prepend_text_content,
+ ]
+)
+
+multimodal_user_transformer_funcs_mistral: List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]] = (
+ base_user_transformer_funcs
+ + [
+ _set_multimodal_content,
+ ]
+)
# === Transformer maps ===
@@ -359,6 +373,8 @@ assistant_transformer_funcs: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]
"tools": tools_assistant_transformer_funcs,
"thought": thought_assistant_transformer_funcs,
}
+
+
assistant_transformer_constructors: Dict[str, Callable[..., Any]] = {
"text": ChatCompletionAssistantMessageParam,
"tools": ChatCompletionAssistantMessageParam,
@@ -403,6 +419,12 @@ assistant_transformer_funcs_claude: Dict[str, List[Callable[[LLMMessage, Dict[st
}
+user_transformer_funcs_mistral: Dict[str, List[Callable[[LLMMessage, Dict[str, Any]], Dict[str, Any]]]] = {
+ "text": single_user_transformer_funcs_mistral,
+ "multimodal": multimodal_user_transformer_funcs_mistral,
+}
+
+
def function_execution_result_message(message: LLMMessage, context: Dict[str, Any]) -> TrasformerReturnType:
assert isinstance(message, FunctionExecutionResultMessage)
return [
@@ -466,6 +488,24 @@ __CLAUDE_TRANSFORMER_MAP: TransformerMap = {
FunctionExecutionResultMessage: function_execution_result_message,
}
+__MISTRAL_TRANSFORMER_MAP: TransformerMap = {
+ SystemMessage: build_transformer_func(
+ funcs=system_message_transformers + [_set_empty_to_whitespace],
+ message_param_func=ChatCompletionSystemMessageParam,
+ ),
+ UserMessage: build_conditional_transformer_func(
+ funcs_map=user_transformer_funcs_mistral,
+ message_param_func_map=user_transformer_constructors,
+ condition_func=user_condition,
+ ),
+ AssistantMessage: build_conditional_transformer_func(
+ funcs_map=assistant_transformer_funcs,
+ message_param_func_map=assistant_transformer_constructors,
+ condition_func=assistant_condition,
+ ),
+ FunctionExecutionResultMessage: function_execution_result_message,
+}
+
# set openai models to use the transformer map
total_models = get_args(ModelFamily.ANY)
@@ -475,7 +515,16 @@ __claude_models = [model for model in total_models if ModelFamily.is_claude(mode
__gemini_models = [model for model in total_models if ModelFamily.is_gemini(model)]
-__unknown_models = list(set(total_models) - set(__openai_models) - set(__claude_models) - set(__gemini_models))
+__llama_models = [model for model in total_models if ModelFamily.is_llama(model)]
+
+__unknown_models = list(
+ set(total_models) - set(__openai_models) - set(__claude_models) - set(__gemini_models) - set(__llama_models)
+)
+__mistral_models = [model for model in total_models if ModelFamily.is_mistral(model)]
+
+__unknown_models = list(
+ set(total_models) - set(__openai_models) - set(__claude_models) - set(__gemini_models) - set(__mistral_models)
+)
for model in __openai_models:
register_transformer("openai", model, __BASE_TRANSFORMER_MAP)
@@ -486,6 +535,12 @@ for model in __claude_models:
for model in __gemini_models:
register_transformer("openai", model, __GEMINI_TRANSFORMER_MAP)
+for model in __llama_models:
+ register_transformer("openai", model, __BASE_TRANSFORMER_MAP)
+
+for model in __mistral_models:
+ register_transformer("openai", model, __MISTRAL_TRANSFORMER_MAP)
+
for model in __unknown_models:
register_transformer("openai", model, __BASE_TRANSFORMER_MAP)
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py
index dfaf87cf6..c5b8b0495 100644
--- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py
+++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py
@@ -36,6 +36,11 @@ _MODEL_POINTERS = {
"claude-3-5-haiku": "claude-3-5-haiku-20241022",
"claude-3-5-sonnet": "claude-3-5-sonnet-20241022",
"claude-3-7-sonnet": "claude-3-7-sonnet-20250219",
+ # Llama models
+ "llama-3.3-8b": "Llama-3.3-8B-Instruct",
+ "llama-3.3-70b": "Llama-3.3-70B-Instruct",
+ "llama-4-scout": "Llama-4-Scout-17B-16E-Instruct-FP8",
+ "llama-4-maverick": "Llama-4-Maverick-17B-128E-Instruct-FP8",
}
_MODEL_INFO: Dict[str, ModelInfo] = {
@@ -351,6 +356,38 @@ _MODEL_INFO: Dict[str, ModelInfo] = {
"structured_output": False,
"multiple_system_messages": True,
},
+ "Llama-3.3-8B-Instruct": {
+ "vision": False,
+ "function_calling": True,
+ "json_output": True,
+ "family": ModelFamily.LLAMA_3_3_8B,
+ "structured_output": False,
+ "multiple_system_messages": True,
+ },
+ "Llama-3.3-70B-Instruct": {
+ "vision": False,
+ "function_calling": True,
+ "json_output": True,
+ "family": ModelFamily.LLAMA_3_3_70B,
+ "structured_output": False,
+ "multiple_system_messages": True,
+ },
+ "Llama-4-Scout-17B-16E-Instruct-FP8": {
+ "vision": True,
+ "function_calling": True,
+ "json_output": True,
+ "family": ModelFamily.LLAMA_4_SCOUT,
+ "structured_output": True,
+ "multiple_system_messages": True,
+ },
+ "Llama-4-Maverick-17B-128E-Instruct-FP8": {
+ "vision": True,
+ "function_calling": True,
+ "json_output": True,
+ "family": ModelFamily.LLAMA_4_MAVERICK,
+ "structured_output": True,
+ "multiple_system_messages": True,
+ },
}
_MODEL_TOKEN_LIMITS: Dict[str, int] = {
@@ -391,10 +428,15 @@ _MODEL_TOKEN_LIMITS: Dict[str, int] = {
"claude-3-5-haiku-20241022": 50000,
"claude-3-5-sonnet-20241022": 40000,
"claude-3-7-sonnet-20250219": 20000,
+ "Llama-3.3-8B-Instruct": 128000,
+ "Llama-3.3-70B-Instruct": 128000,
+ "Llama-4-Scout-17B-16E-Instruct-FP8": 128000,
+ "Llama-4-Maverick-17B-128E-Instruct-FP8": 128000,
}
GEMINI_OPENAI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
ANTHROPIC_OPENAI_BASE_URL = "https://api.anthropic.com/v1/"
+LLAMA_API_BASE_URL = "https://api.llama.com/compat/v1/"
def resolve_model(model: str) -> str:
diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py
index ffe816e59..936b3a6c7 100644
--- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py
+++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py
@@ -637,6 +637,7 @@ class BaseOpenAIChatCompletionClient(ChatCompletionClient):
response=result.model_dump(),
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
+ tools=create_params.tools,
)
)
@@ -1369,6 +1370,11 @@ class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component[OpenA
copied_args["base_url"] = _model_info.ANTHROPIC_OPENAI_BASE_URL
if "api_key" not in copied_args and "ANTHROPIC_API_KEY" in os.environ:
copied_args["api_key"] = os.environ["ANTHROPIC_API_KEY"]
+ if copied_args["model"].startswith("Llama-"):
+ if "base_url" not in copied_args:
+ copied_args["base_url"] = _model_info.LLAMA_API_BASE_URL
+ if "api_key" not in copied_args and "LLAMA_API_KEY" in os.environ:
+ copied_args["api_key"] = os.environ["LLAMA_API_KEY"]
client = _openai_client_from_config(copied_args)
create_args = _create_args_from_config(copied_args)
diff --git a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py
index 504838740..6a3963586 100644
--- a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py
+++ b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py
@@ -251,6 +251,7 @@ class GrpcWorkerAgentRuntime(AgentRuntime):
self._subscription_manager = SubscriptionManager()
self._serialization_registry = SerializationRegistry()
self._extra_grpc_config = extra_grpc_config or []
+ self._agent_instance_types: Dict[str, Type[Agent]] = {}
if payload_serialization_format not in {JSON_DATA_CONTENT_TYPE, PROTOBUF_DATA_CONTENT_TYPE}:
raise ValueError(f"Unsupported payload serialization format: {payload_serialization_format}")
@@ -701,6 +702,14 @@ class GrpcWorkerAgentRuntime(AgentRuntime):
except BaseException as e:
logger.error("Error handling event", exc_info=e)
+ async def _register_agent_type(self, agent_type: str) -> None:
+ if self._host_connection is None:
+ raise RuntimeError("Host connection is not set.")
+ message = agent_worker_pb2.RegisterAgentTypeRequest(type=agent_type)
+ _response: agent_worker_pb2.RegisterAgentTypeResponse = await self._host_connection.stub.RegisterAgent(
+ message, metadata=self._host_connection.metadata
+ )
+
async def register_factory(
self,
type: str | AgentType,
@@ -729,14 +738,38 @@ class GrpcWorkerAgentRuntime(AgentRuntime):
return agent_instance
self._agent_factories[type.type] = factory_wrapper
-
# Send the registration request message to the host.
- message = agent_worker_pb2.RegisterAgentTypeRequest(type=type.type)
- _response: agent_worker_pb2.RegisterAgentTypeResponse = await self._host_connection.stub.RegisterAgent(
- message, metadata=self._host_connection.metadata
- )
+ await self._register_agent_type(type.type)
+
return type
+ async def register_agent_instance(
+ self,
+ agent_instance: Agent,
+ agent_id: AgentId,
+ ) -> AgentId:
+ def agent_factory() -> Agent:
+ raise RuntimeError(
+ "Agent factory was invoked for an agent instance that was not registered. This is likely due to the agent type being incorrectly subscribed to a topic. If this exception occurs when publishing a message to the DefaultTopicId, then it is likely that `skip_class_subscriptions` needs to be turned off when registering the agent."
+ )
+
+ if agent_id in self._instantiated_agents:
+ raise ValueError(f"Agent with id {agent_id} already exists.")
+
+ if agent_id.type not in self._agent_factories:
+ self._agent_factories[agent_id.type] = agent_factory
+ await self._register_agent_type(agent_id.type)
+ self._agent_instance_types[agent_id.type] = type_func_alias(agent_instance)
+ else:
+ if self._agent_factories[agent_id.type].__code__ != agent_factory.__code__:
+ raise ValueError("Agent factories and agent instances cannot be registered to the same type.")
+ if self._agent_instance_types[agent_id.type] != type_func_alias(agent_instance):
+ raise ValueError("Agent instances must be the same object type.")
+
+ await agent_instance.bind_id_and_runtime(id=agent_id, runtime=self)
+ self._instantiated_agents[agent_id] = agent_instance
+ return agent_id
+
async def _invoke_agent_factory(
self,
agent_factory: Callable[[], T | Awaitable[T]] | Callable[[AgentRuntime, AgentId], T | Awaitable[T]],
@@ -757,7 +790,7 @@ class GrpcWorkerAgentRuntime(AgentRuntime):
raise ValueError("Agent factory must take 0 or 2 arguments.")
if inspect.isawaitable(agent):
- return cast(T, await agent)
+ agent = cast(T, await agent)
return agent
diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/azure/__init__.py b/python/packages/autogen-ext/src/autogen_ext/tools/azure/__init__.py
index 733ba05d1..12851185e 100644
--- a/python/packages/autogen-ext/src/autogen_ext/tools/azure/__init__.py
+++ b/python/packages/autogen-ext/src/autogen_ext/tools/azure/__init__.py
@@ -4,6 +4,7 @@ from ._ai_search import (
SearchQuery,
SearchResult,
SearchResults,
+ VectorizableTextQuery,
)
from ._config import AzureAISearchConfig
@@ -14,4 +15,5 @@ __all__ = [
"SearchResult",
"SearchResults",
"AzureAISearchConfig",
+ "VectorizableTextQuery",
]
diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/azure/_ai_search.py b/python/packages/autogen-ext/src/autogen_ext/tools/azure/_ai_search.py
index cf0570d82..447a716b6 100644
--- a/python/packages/autogen-ext/src/autogen_ext/tools/azure/_ai_search.py
+++ b/python/packages/autogen-ext/src/autogen_ext/tools/azure/_ai_search.py
@@ -1,112 +1,87 @@
+from __future__ import annotations
+
+import asyncio
import logging
import time
from abc import ABC, abstractmethod
from contextvars import ContextVar
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, TypeVar, Union, cast, overload
-
-from autogen_core import CancellationToken, ComponentModel
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Dict,
+ List,
+ Literal,
+ Optional,
+ Protocol,
+ Union,
+)
+
+from autogen_core import CancellationToken, Component
from autogen_core.tools import BaseTool, ToolSchema
-from azure.core.credentials import AzureKeyCredential, TokenCredential
+from azure.core.credentials import AzureKeyCredential
+from azure.core.credentials_async import AsyncTokenCredential
from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
from azure.search.documents.aio import SearchClient
from pydantic import BaseModel, Field
-if TYPE_CHECKING:
- from azure.search.documents.models import VectorizableTextQuery
-
-_has_retry_policy = False
-try:
- from azure.core.pipeline.policies import RetryPolicy # type: ignore[assignment]
-
- _has_retry_policy = True
-except ImportError:
-
- class RetryPolicy: # type: ignore
- def __init__(self, retry_mode: str = "fixed", retry_total: int = 3, **kwargs: Any) -> None:
- pass
-
- _has_retry_policy = False
-
-HAS_RETRY_POLICY = _has_retry_policy
-
-has_azure_search = False
-
-if not TYPE_CHECKING:
- try:
- from azure.search.documents.models import VectorizableTextQuery
-
- has_azure_search = True
- except ImportError:
+from ._config import (
+ DEFAULT_API_VERSION,
+ AzureAISearchConfig,
+)
- class VectorizableTextQuery:
- """Fallback implementation when Azure SDK is not installed."""
+SearchDocument = Dict[str, Any]
+MetadataDict = Dict[str, Any]
+ContentDict = Dict[str, Any]
- def __init__(self, text: str, k: int, fields: Union[str, List[str]]) -> None:
- self.text = text
- self.k = k
- self.fields = fields if isinstance(fields, str) else ",".join(fields)
-
-
-class _FallbackAzureAISearchConfig:
- """Fallback configuration class for Azure AI Search when the main config module is not available.
+if TYPE_CHECKING:
+ from azure.search.documents.aio import AsyncSearchItemPaged
- This class provides a simple dictionary-based configuration object that mimics the behavior
- of the AzureAISearchConfig from the _config module. It's used as a fallback when the main
- configuration module cannot be imported.
+ SearchResultsIterable = AsyncSearchItemPaged[SearchDocument]
+else:
+ SearchResultsIterable = Any
- Args:
- **kwargs (Any): Keyword arguments containing configuration values
- """
+logger = logging.getLogger(__name__)
- def __init__(self, **kwargs: Any):
- self.name = kwargs.get("name", "")
- self.description = kwargs.get("description", "")
- self.endpoint = kwargs.get("endpoint", "")
- self.index_name = kwargs.get("index_name", "")
- self.credential = kwargs.get("credential", None)
- self.api_version = kwargs.get("api_version", "")
- self.query_type = kwargs.get("query_type", "simple")
- self.search_fields = kwargs.get("search_fields", None)
- self.select_fields = kwargs.get("select_fields", None)
- self.vector_fields = kwargs.get("vector_fields", None)
- self.filter = kwargs.get("filter", None)
- self.top = kwargs.get("top", None)
- self.retry_enabled = kwargs.get("retry_enabled", False)
- self.retry_mode = kwargs.get("retry_mode", "fixed")
- self.retry_max_attempts = kwargs.get("retry_max_attempts", 3)
- self.enable_caching = kwargs.get("enable_caching", False)
- self.cache_ttl_seconds = kwargs.get("cache_ttl_seconds", 300)
-
-
-AzureAISearchConfig: Any
+if TYPE_CHECKING:
+ from azure.search.documents.models import (
+ VectorizableTextQuery,
+ VectorizedQuery,
+ VectorQuery,
+ )
try:
- from ._config import AzureAISearchConfig
-except ImportError:
- import importlib.util
- import os
-
- current_dir = os.path.dirname(os.path.abspath(__file__))
- config_path = os.path.join(current_dir, "_config.py")
- config_module = None
+ from azure.search.documents.models import VectorizableTextQuery, VectorizedQuery, VectorQuery
- spec_config = importlib.util.spec_from_file_location("config_module", config_path)
- if spec_config is not None:
- config_module = importlib.util.module_from_spec(spec_config)
- loader = getattr(spec_config, "loader", None)
- if loader is not None:
- loader.exec_module(config_module)
+ has_azure_search = True
+except ImportError:
+ has_azure_search = False
+ logger.error(
+ "The 'azure-search-documents' package is required for this tool but was not found. "
+ "Please install it with: uv add install azure-search-documents"
+ )
- if config_module is not None and hasattr(config_module, "AzureAISearchConfig"):
- AzureAISearchConfig = config_module.AzureAISearchConfig
- else:
- AzureAISearchConfig = _FallbackAzureAISearchConfig
+if TYPE_CHECKING:
+ from typing import Protocol
+
+ class SearchClientProtocol(Protocol):
+ async def search(self, **kwargs: Any) -> SearchResultsIterable: ...
+ async def close(self) -> None: ...
+else:
+ SearchClientProtocol = Any
+
+__all__ = [
+ "AzureAISearchTool",
+ "BaseAzureAISearchTool",
+ "SearchQuery",
+ "SearchResults",
+ "SearchResult",
+ "VectorizableTextQuery",
+ "VectorizedQuery",
+ "VectorQuery",
+]
logger = logging.getLogger(__name__)
-T = TypeVar("T", bound="BaseAzureAISearchTool")
-ExpectedType = TypeVar("ExpectedType")
-
class SearchQuery(BaseModel):
"""Search query parameters.
@@ -127,13 +102,13 @@ class SearchResult(BaseModel):
Args:
score (float): The search score.
- content (Dict[str, Any]): The document content.
- metadata (Dict[str, Any]): Additional metadata about the document.
+ content (ContentDict): The document content.
+ metadata (MetadataDict): Additional metadata about the document.
"""
score: float = Field(description="The search score")
- content: Dict[str, Any] = Field(description="The document content")
- metadata: Dict[str, Any] = Field(description="Additional metadata about the document")
+ content: ContentDict = Field(description="The document content")
+ metadata: MetadataDict = Field(description="Additional metadata about the document")
class SearchResults(BaseModel):
@@ -146,7 +121,100 @@ class SearchResults(BaseModel):
results: List[SearchResult] = Field(description="List of search results")
-class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
+class EmbeddingProvider(Protocol):
+ """Protocol defining the interface for embedding generation."""
+
+ async def _get_embedding(self, query: str) -> List[float]:
+ """Generate embedding vector for the query text."""
+ ...
+
+
+class EmbeddingProviderMixin:
+ """Mixin class providing embedding generation functionality."""
+
+ search_config: AzureAISearchConfig
+
+ async def _get_embedding(self, query: str) -> List[float]:
+ """Generate embedding vector for the query text."""
+ if not hasattr(self, "search_config"):
+ raise ValueError("Host class must have a search_config attribute")
+
+ search_config = self.search_config
+ embedding_provider = getattr(search_config, "embedding_provider", None)
+ embedding_model = getattr(search_config, "embedding_model", None)
+
+ if not embedding_provider or not embedding_model:
+ raise ValueError(
+ "Client-side embedding is not configured. `embedding_provider` and `embedding_model` must be set."
+ ) from None
+
+ if embedding_provider.lower() == "azure_openai":
+ try:
+ from azure.identity import DefaultAzureCredential
+ from openai import AsyncAzureOpenAI
+ except ImportError:
+ raise ImportError(
+ "Azure OpenAI SDK is required for client-side embedding generation. "
+ "Please install it with: uv add openai azure-identity"
+ ) from None
+
+ api_key = getattr(search_config, "openai_api_key", None)
+ api_version = getattr(search_config, "openai_api_version", "2023-11-01")
+ endpoint = getattr(search_config, "openai_endpoint", None)
+
+ if not endpoint:
+ raise ValueError(
+ "Azure OpenAI endpoint (`openai_endpoint`) must be provided for client-side Azure OpenAI embeddings."
+ ) from None
+
+ if api_key:
+ azure_client = AsyncAzureOpenAI(api_key=api_key, api_version=api_version, azure_endpoint=endpoint)
+ else:
+
+ def get_token() -> str:
+ credential = DefaultAzureCredential()
+ token = credential.get_token("https://cognitiveservices.azure.com/.default")
+ if not token or not token.token:
+ raise ValueError("Failed to acquire token using DefaultAzureCredential for Azure OpenAI.")
+ return token.token
+
+ azure_client = AsyncAzureOpenAI(
+ azure_ad_token_provider=get_token, api_version=api_version, azure_endpoint=endpoint
+ )
+
+ try:
+ response = await azure_client.embeddings.create(model=embedding_model, input=query)
+ return response.data[0].embedding
+ except Exception as e:
+ raise ValueError(f"Failed to generate embeddings with Azure OpenAI: {str(e)}") from e
+
+ elif embedding_provider.lower() == "openai":
+ try:
+ from openai import AsyncOpenAI
+ except ImportError:
+ raise ImportError(
+ "OpenAI SDK is required for client-side embedding generation. "
+ "Please install it with: uv add openai"
+ ) from None
+
+ api_key = getattr(search_config, "openai_api_key", None)
+ openai_client = AsyncOpenAI(api_key=api_key)
+
+ try:
+ response = await openai_client.embeddings.create(model=embedding_model, input=query)
+ return response.data[0].embedding
+ except Exception as e:
+ raise ValueError(f"Failed to generate embeddings with OpenAI: {str(e)}") from e
+ else:
+ raise ValueError(
+ f"Unsupported client-side embedding provider: {embedding_provider}. "
+ "Currently supported providers are 'azure_openai' and 'openai'."
+ )
+
+
+class BaseAzureAISearchTool(
+ BaseTool[SearchQuery, SearchResults], Component[AzureAISearchConfig], EmbeddingProvider, ABC
+):
"""Abstract base class for Azure AI Search tools.
This class defines the common interface and functionality for all Azure AI Search tools.
@@ -161,15 +229,18 @@ class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
Use concrete implementations or the factory methods in AzureAISearchTool.
"""
+ component_config_schema = AzureAISearchConfig
+ component_provider_override = "autogen_ext.tools.azure.BaseAzureAISearchTool"
+
def __init__(
self,
name: str,
endpoint: str,
index_name: str,
- credential: Union[AzureKeyCredential, TokenCredential, Dict[str, str]],
+ credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]],
description: Optional[str] = None,
- api_version: str = "2023-11-01",
- query_type: Literal["keyword", "fulltext", "vector", "semantic"] = "keyword",
+ api_version: str = DEFAULT_API_VERSION,
+ query_type: Literal["simple", "full", "semantic", "vector"] = "simple",
search_fields: Optional[List[str]] = None,
select_fields: Optional[List[str]] = None,
vector_fields: Optional[List[str]] = None,
@@ -178,6 +249,11 @@ class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
semantic_config_name: Optional[str] = None,
enable_caching: bool = False,
cache_ttl_seconds: int = 300,
+ embedding_provider: Optional[str] = None,
+ embedding_model: Optional[str] = None,
+ openai_api_key: Optional[str] = None,
+ openai_api_version: Optional[str] = None,
+ openai_endpoint: Optional[str] = None,
):
"""Initialize the Azure AI Search tool.
@@ -185,10 +261,10 @@ class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
name (str): The name of this tool instance
endpoint (str): The full URL of your Azure AI Search service
index_name (str): Name of the search index to query
- credential (Union[AzureKeyCredential, TokenCredential, Dict[str, str]]): Azure credential for authentication (API key or token)
+ credential (Union[AzureKeyCredential, TokenCredential, Dict[str, str]]): Azure credential for authentication
description (Optional[str]): Optional description explaining the tool's purpose
- api_version (str): Azure AI Search API version to use
- query_type (Literal["keyword", "fulltext", "vector", "semantic"]): Type of search to perform
+ api_version (Optional[str]): Azure AI Search API version to use
+ query_type (Literal["simple", "full", "semantic", "vector"]): Type of search to perform
search_fields (Optional[List[str]]): Fields to search within documents
select_fields (Optional[List[str]]): Fields to return in search results
vector_fields (Optional[List[str]]): Fields to use for vector search
@@ -197,6 +273,11 @@ class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
semantic_config_name (Optional[str]): Semantic configuration name for enhanced results
enable_caching (bool): Whether to cache search results
cache_ttl_seconds (int): How long to cache results in seconds
+ embedding_provider (Optional[str]): Name of embedding provider for client-side embeddings
+ embedding_model (Optional[str]): Model name for client-side embeddings
+ openai_api_key (Optional[str]): API key for OpenAI/Azure OpenAI embeddings
+ openai_api_version (Optional[str]): API version for Azure OpenAI embeddings
+ openai_endpoint (Optional[str]): Endpoint URL for Azure OpenAI embeddings
"""
if not has_azure_search:
raise ImportError(
@@ -217,12 +298,14 @@ class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
description=description,
)
- self.search_config = AzureAISearchConfig(
+ processed_credential = self._process_credential(credential)
+
+ self.search_config: AzureAISearchConfig = AzureAISearchConfig(
name=name,
description=description,
endpoint=endpoint,
index_name=index_name,
- credential=self._process_credential(credential),
+ credential=processed_credential,
api_version=api_version,
query_type=query_type,
search_fields=search_fields,
@@ -233,35 +316,77 @@ class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
semantic_config_name=semantic_config_name,
enable_caching=enable_caching,
cache_ttl_seconds=cache_ttl_seconds,
+ embedding_provider=embedding_provider,
+ embedding_model=embedding_model,
+ openai_api_key=openai_api_key,
+ openai_api_version=openai_api_version,
+ openai_endpoint=openai_endpoint,
)
self._endpoint = endpoint
self._index_name = index_name
- self._credential = credential
+ self._credential = processed_credential
self._api_version = api_version
+
self._client: Optional[SearchClient] = None
self._cache: Dict[str, Dict[str, Any]] = {}
+ if self.search_config.api_version == "2023-11-01" and self.search_config.vector_fields:
+ warning_message = (
+ f"When explicitly setting api_version='{self.search_config.api_version}' for vector search: "
+ f"If client-side embedding is NOT configured (e.g., `embedding_model` is not set), "
+ f"this tool defaults to service-side vectorization (VectorizableTextQuery), which may fail or have limitations with this API version. "
+ f"If client-side embedding IS configured, the tool will use VectorizedQuery, which is generally compatible. "
+ f"For robust vector search, consider omitting api_version (recommended to use SDK default) or use a newer API version."
+ )
+ logger.warning(warning_message)
+
async def close(self) -> None:
- """Explicitly close the Azure SearchClient if needed (for cleanup in long-running apps/tests)."""
+ """Explicitly close the Azure SearchClient if needed (for cleanup)."""
if self._client is not None:
- await self._client.close()
- self._client = None
+ try:
+ await self._client.close()
+ except Exception:
+ pass
+ finally:
+ self._client = None
def _process_credential(
- self, credential: Union[AzureKeyCredential, TokenCredential, Dict[str, str]]
- ) -> Union[AzureKeyCredential, TokenCredential]:
- """Process credential to ensure it's the correct type."""
+ self, credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]]
+ ) -> Union[AzureKeyCredential, AsyncTokenCredential]:
+ """Process credential to ensure it's the correct type for async SearchClient.
+
+ Converts dictionary credentials with 'api_key' to AzureKeyCredential objects.
+
+ Args:
+ credential: The credential in either object or dictionary form
+
+ Returns:
+ A properly formatted credential object
+
+ Raises:
+ ValueError: If the credential dictionary doesn't contain an 'api_key'
+ TypeError: If the credential is not of a supported type
+ """
if isinstance(credential, dict):
if "api_key" in credential:
return AzureKeyCredential(credential["api_key"])
- raise ValueError(
- "If credential is a dict, it must contain an 'api_key' key with your API key as the value"
- ) from None
- return credential
+ raise ValueError("If credential is a dict, it must contain an 'api_key' key")
+
+ if isinstance(credential, (AzureKeyCredential, AsyncTokenCredential)):
+ return credential
+
+ raise TypeError("Credential must be AzureKeyCredential, AsyncTokenCredential, or a valid dict")
async def _get_client(self) -> SearchClient:
- """Get the search client for the configured index."""
+ """Get the search client for the configured index.
+
+ Returns:
+ SearchClient: Initialized search client
+
+ Raises:
+ ValueError: If index doesn't exist or authentication fails
+ """
if self._client is not None:
return self._client
@@ -272,23 +397,14 @@ class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
credential=self.search_config.credential,
api_version=self.search_config.api_version,
)
-
- assert self._client is not None
return self._client
except ResourceNotFoundError as e:
- raise ValueError(
- f"Index '{self.search_config.index_name}' not found. "
- f"Please check if the index exists in your Azure AI Search service at {self.search_config.endpoint}"
- ) from e
+ raise ValueError(f"Index '{self.search_config.index_name}' not found in Azure AI Search service.") from e
except HttpResponseError as e:
- if "401" in str(e):
- raise ValueError(
- f"Authentication failed. Please check your API key or credentials. Error: {str(e)}"
- ) from e
- elif "403" in str(e):
- raise ValueError(
- f"Permission denied. Please check that your credentials have access to this index. Error: {str(e)}"
- ) from e
+ if e.status_code == 401:
+ raise ValueError("Authentication failed. Please check your credentials.") from e
+ elif e.status_code == 403:
+ raise ValueError("Permission denied to access this index.") from e
else:
raise ValueError(f"Error connecting to Azure AI Search: {str(e)}") from e
except Exception as e:
@@ -304,311 +420,170 @@ class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
cancellation_token: Optional token to cancel the operation
Returns:
- Search results
- """
- if isinstance(args, str) and not args.strip():
- raise ValueError("Invalid search query format: Query cannot be empty")
+ SearchResults: Container with search results and metadata
+ Raises:
+ ValueError: If the search query is empty or invalid
+ ValueError: If there is an authentication error or other search issue
+ asyncio.CancelledError: If the operation is cancelled
+ """
if isinstance(args, str):
+ if not args.strip():
+ raise ValueError("Search query cannot be empty")
search_query = SearchQuery(query=args)
elif isinstance(args, dict) and "query" in args:
search_query = SearchQuery(query=args["query"])
elif isinstance(args, SearchQuery):
search_query = args
else:
- raise ValueError(f"Invalid search query format: {args}. Expected string, dict with 'query', or SearchQuery")
+ raise ValueError("Invalid search query format. Expected string, dict with 'query', or SearchQuery")
+
+ if cancellation_token is not None and cancellation_token.is_cancelled():
+ raise asyncio.CancelledError("Operation cancelled")
+
+ cache_key = ""
+ if self.search_config.enable_caching:
+ cache_key_parts = [
+ search_query.query,
+ str(self.search_config.top),
+ self.search_config.query_type,
+ ",".join(sorted(self.search_config.search_fields or [])),
+ ",".join(sorted(self.search_config.select_fields or [])),
+ ",".join(sorted(self.search_config.vector_fields or [])),
+ str(self.search_config.filter or ""),
+ str(self.search_config.semantic_config_name or ""),
+ ]
+ cache_key = ":".join(filter(None, cache_key_parts))
+ if cache_key in self._cache:
+ cache_entry = self._cache[cache_key]
+ cache_age = time.time() - cache_entry["timestamp"]
+ if cache_age < self.search_config.cache_ttl_seconds:
+ logger.debug(f"Using cached results for query: {search_query.query}")
+ return SearchResults(
+ results=[
+ SearchResult(score=r.score, content=r.content, metadata=r.metadata)
+ for r in cache_entry["results"]
+ ]
+ )
try:
- if cancellation_token is not None and cancellation_token.is_cancelled():
- raise Exception("Operation cancelled")
+ search_kwargs: Dict[str, Any] = {}
- if self.search_config.enable_caching:
- cache_key = f"{search_query.query}:{self.search_config.top}"
- if cache_key in self._cache:
- cache_entry = self._cache[cache_key]
- cache_age = time.time() - cache_entry["timestamp"]
- if cache_age < self.search_config.cache_ttl_seconds:
- logger.debug(f"Using cached results for query: {search_query.query}")
- return SearchResults(
- results=[
- SearchResult(score=r.score, content=r.content, metadata=r.metadata)
- for r in cache_entry["results"]
- ]
- )
+ if self.search_config.query_type != "vector":
+ search_kwargs["search_text"] = search_query.query
+ search_kwargs["query_type"] = self.search_config.query_type
+
+ if self.search_config.search_fields:
+ search_kwargs["search_fields"] = self.search_config.search_fields # type: ignore[assignment]
- search_options: Dict[str, Any] = {}
- search_options["query_type"] = self.search_config.query_type
+ if self.search_config.query_type == "semantic" and self.search_config.semantic_config_name:
+ search_kwargs["semantic_configuration_name"] = self.search_config.semantic_config_name
if self.search_config.select_fields:
- search_options["select"] = self.search_config.select_fields
+ search_kwargs["select"] = self.search_config.select_fields # type: ignore[assignment]
+ if self.search_config.filter:
+ search_kwargs["filter"] = str(self.search_config.filter)
+ if self.search_config.top is not None:
+ search_kwargs["top"] = self.search_config.top # type: ignore[assignment]
- if self.search_config.search_fields:
- search_options["search_fields"] = self.search_config.search_fields
+ if self.search_config.vector_fields and len(self.search_config.vector_fields) > 0:
+ if not search_query.query:
+ raise ValueError("Query text cannot be empty for vector search operations")
- if self.search_config.filter:
- search_options["filter"] = self.search_config.filter
+ use_client_side_embeddings = bool(
+ self.search_config.embedding_model and self.search_config.embedding_provider
+ )
- if self.search_config.top is not None:
- search_options["top"] = self.search_config.top
-
- if self.search_config.query_type == "fulltext" and self.search_config.semantic_config_name is not None:
- search_options["query_type"] = "semantic"
- search_options["semantic_configuration_name"] = self.search_config.semantic_config_name
-
- text_query = search_query.query
- if self.search_config.query_type == "vector" or (
- self.search_config.vector_fields and len(self.search_config.vector_fields) > 0
- ):
- if self.search_config.vector_fields:
- vector_fields_list = self.search_config.vector_fields
- search_options["vector_queries"] = [
- VectorizableTextQuery(
- text=search_query.query, k_nearest_neighbors=int(self.search_config.top or 5), fields=field
+ vector_queries: List[Union[VectorizedQuery, VectorizableTextQuery]] = []
+ if use_client_side_embeddings:
+ from azure.search.documents.models import VectorizedQuery
+
+ embedding_vector: List[float] = await self._get_embedding(search_query.query)
+ for field_spec in self.search_config.vector_fields:
+ fields = field_spec if isinstance(field_spec, str) else ",".join(field_spec)
+ vector_queries.append(
+ VectorizedQuery(
+ vector=embedding_vector,
+ k_nearest_neighbors=self.search_config.top or 5,
+ fields=fields,
+ kind="vector",
+ )
+ )
+ else:
+ from azure.search.documents.models import VectorizableTextQuery
+
+ for field in self.search_config.vector_fields:
+ fields = field if isinstance(field, str) else ",".join(field)
+ vector_queries.append(
+ VectorizableTextQuery( # type: ignore
+ text=search_query.query,
+ k_nearest_neighbors=self.search_config.top or 5,
+ fields=fields,
+ kind="vectorizable",
+ )
)
- for field in vector_fields_list
- ]
-
- client = await self._get_client()
- results: List[SearchResult] = []
- # Use the persistent client directly. Do NOT close after each operation.
- # WARNING: The SearchClient must live as long as the tool/agent is in use.
- search_future = client.search(text_query, **search_options) # type: ignore
+ search_kwargs["vector_queries"] = vector_queries # type: ignore[assignment]
if cancellation_token is not None:
- import asyncio
-
- # Using explicit type ignores to handle Azure SDK type complexity
- async def awaitable_wrapper(): # type: ignore # pyright: ignore[reportUnknownVariableType,reportUnknownLambdaType,reportUnknownMemberType]
- return await search_future # pyright: ignore[reportUnknownVariableType]
+ dummy_task = asyncio.create_task(asyncio.sleep(60))
+ cancellation_token.link_future(dummy_task)
- task = asyncio.create_task(awaitable_wrapper()) # type: ignore # pyright: ignore[reportUnknownVariableType]
- cancellation_token.link_future(task) # pyright: ignore[reportUnknownArgumentType]
- search_results = await task # pyright: ignore[reportUnknownVariableType]
+ def is_cancelled() -> bool:
+ return cancellation_token.is_cancelled()
else:
- search_results = await search_future # pyright: ignore[reportUnknownVariableType]
- async for doc in search_results: # type: ignore
- search_doc: Any = doc
- doc_dict: Dict[str, Any] = {}
+ def is_cancelled() -> bool:
+ return False
+
+ client = await self._get_client()
+ search_results: SearchResultsIterable = await client.search(**search_kwargs) # type: ignore[arg-type]
+
+ results: List[SearchResult] = []
+ async for doc in search_results:
+ if is_cancelled():
+ raise asyncio.CancelledError("Operation was cancelled")
try:
- if hasattr(search_doc, "items") and callable(search_doc.items):
- dict_like_doc = cast(Dict[str, Any], search_doc)
- for key, value in dict_like_doc.items():
- doc_dict[str(key)] = value
- else:
- for key in [
- k
- for k in dir(search_doc)
- if not k.startswith("_") and not callable(getattr(search_doc, k, None))
- ]:
- doc_dict[key] = getattr(search_doc, key)
+ metadata: Dict[str, Any] = {}
+ content: Dict[str, Any] = {}
+
+ for key, value in doc.items():
+ if isinstance(key, str) and key.startswith(("@", "_")):
+ metadata[key] = value
+ else:
+ content[str(key)] = value
+
+ score = float(metadata.get("@search.score", 0.0))
+ results.append(SearchResult(score=score, content=content, metadata=metadata))
except Exception as e:
logger.warning(f"Error processing search document: {e}")
continue
- metadata: Dict[str, Any] = {}
- content: Dict[str, Any] = {}
- for key, value in doc_dict.items():
- key_str: str = str(key)
- if key_str.startswith("@") or key_str.startswith("_"):
- metadata[key_str] = value
- else:
- content[key_str] = value
-
- score: float = 0.0
- if "@search.score" in doc_dict:
- score = float(doc_dict["@search.score"])
-
- result = SearchResult(
- score=score,
- content=content,
- metadata=metadata,
- )
- results.append(result)
-
if self.search_config.enable_caching:
- cache_key = f"{text_query}_{self.search_config.top}"
self._cache[cache_key] = {"results": results, "timestamp": time.time()}
- return SearchResults(
- results=[SearchResult(score=r.score, content=r.content, metadata=r.metadata) for r in results]
- )
+ return SearchResults(results=results)
+
+ except asyncio.CancelledError:
+ raise
except Exception as e:
+ error_msg = str(e)
if isinstance(e, HttpResponseError):
if hasattr(e, "message") and e.message:
- if "401 unauthorized" in e.message.lower() or "access denied" in e.message.lower():
- raise ValueError(
- f"Authentication failed: {e.message}. Please check your API key and credentials."
- ) from e
- elif "500" in e.message:
- raise ValueError(f"Error from Azure AI Search: {e.message}") from e
- else:
- raise ValueError(f"Error from Azure AI Search: {e.message}") from e
-
- if hasattr(self, "_name") and self._name == "test_search":
- if (
- hasattr(self, "_credential")
- and isinstance(self._credential, AzureKeyCredential)
- and self._credential.key == "invalid-key"
- ):
- raise ValueError(
- "Authentication failed: 401 Unauthorized. Please check your API key and credentials."
- ) from e
- elif "invalid status" in str(e).lower():
- raise ValueError(
- "Error from Azure AI Search: 500 Internal Server Error: Something went wrong"
- ) from e
+ error_msg = e.message
- error_msg = str(e)
if "not found" in error_msg.lower():
- raise ValueError(
- f"Index '{self.search_config.index_name}' not found. Please check the index name and try again."
- ) from e
+ raise ValueError(f"Index '{self.search_config.index_name}' not found.") from e
elif "unauthorized" in error_msg.lower() or "401" in error_msg:
- raise ValueError(
- f"Authentication failed: {error_msg}. Please check your API key and credentials."
- ) from e
+ raise ValueError(f"Authentication failed: {error_msg}") from e
else:
raise ValueError(f"Error from Azure AI Search: {error_msg}") from e
- @abstractmethod
- async def _get_embedding(self, query: str) -> List[float]:
- """Generate embedding vector for the query text.
-
- This method must be implemented by subclasses to provide embeddings for vector search.
-
- Args:
- query (str): The text to generate embeddings for.
-
- Returns:
- List[float]: The embedding vector as a list of floats.
- """
- pass
-
- def _to_config(self) -> Any:
- """Get the tool configuration.
-
- Returns:
- Any: The search configuration object
- """
+ def _to_config(self) -> AzureAISearchConfig:
+ """Convert the current instance to a configuration object."""
return self.search_config
- def dump_component(self) -> ComponentModel:
- """Serialize the tool to a component model.
-
- Returns:
- ComponentModel: A serialized representation of the tool
- """
- config = self._to_config()
- return ComponentModel(
- provider="autogen_ext.tools.azure.BaseAzureAISearchTool",
- config=config.model_dump(exclude_none=True),
- )
-
- @classmethod
- def _from_config(cls, config: Any) -> "BaseAzureAISearchTool":
- """Create a tool instance from configuration.
-
- Args:
- config (Any): The configuration object containing tool settings
-
- Returns:
- BaseAzureAISearchTool: An initialized instance of the search tool
- """
- query_type_str = getattr(config, "query_type", "keyword")
-
- query_type_mapping = {
- "keyword": "keyword",
- "simple": "fulltext",
- "fulltext": "fulltext",
- "vector": "vector",
- "semantic": "semantic",
- }
-
- query_type = cast(
- Literal["keyword", "fulltext", "vector", "semantic"], query_type_mapping.get(query_type_str, "vector")
- )
-
- openai_client_attr = getattr(config, "openai_client", None)
- if openai_client_attr is None:
- raise ValueError("openai_client must be provided in config")
-
- embedding_model_attr = getattr(config, "embedding_model", "")
- if not embedding_model_attr:
- raise ValueError("embedding_model must be specified in config")
-
- # If query_type="semantic", you must provide a valid semantic_config_name.
- # If query_type is anything else, semantic_config_name is ignored.
- return cls(
- name=getattr(config, "name", ""),
- endpoint=getattr(config, "endpoint", ""),
- index_name=getattr(config, "index_name", ""),
- credential=getattr(config, "credential", {}),
- description=getattr(config, "description", None),
- api_version=getattr(config, "api_version", "2023-11-01"),
- query_type=query_type,
- search_fields=getattr(config, "search_fields", None),
- select_fields=getattr(config, "select_fields", None),
- vector_fields=getattr(config, "vector_fields", None),
- top=getattr(config, "top", None),
- filter=getattr(config, "filter", None),
- semantic_config_name=getattr(config, "semantic_config_name", None),
- enable_caching=getattr(config, "enable_caching", False),
- cache_ttl_seconds=getattr(config, "cache_ttl_seconds", 300),
- )
-
- @overload
- @classmethod
- def load_component(
- cls, model: Union[ComponentModel, Dict[str, Any]], expected: None = None
- ) -> "BaseAzureAISearchTool": ...
-
- @overload
- @classmethod
- def load_component(
- cls, model: Union[ComponentModel, Dict[str, Any]], expected: Type[ExpectedType]
- ) -> ExpectedType: ...
-
- @classmethod
- def load_component(
- cls,
- model: Union[ComponentModel, Dict[str, Any]],
- expected: Optional[Type[ExpectedType]] = None,
- ) -> Union["BaseAzureAISearchTool", ExpectedType]:
- """Load the tool from a component model.
-
- Args:
- model (Union[ComponentModel, Dict[str, Any]]): The component configuration.
- expected (Optional[Type[ExpectedType]]): Optional component class for deserialization.
-
- Returns:
- Union[BaseAzureAISearchTool, ExpectedType]: An instance of the tool.
-
- Raises:
- ValueError: If the component configuration is invalid.
- """
- if expected is not None and not issubclass(expected, BaseAzureAISearchTool):
- raise TypeError(f"Cannot create instance of {expected} from AzureAISearchConfig")
-
- target_class = expected if expected is not None else cls
- assert hasattr(target_class, "_from_config"), f"{target_class} has no _from_config method"
-
- if isinstance(model, ComponentModel) and hasattr(model, "config"):
- config_dict = model.config
- elif isinstance(model, dict):
- config_dict = model
- else:
- raise ValueError(f"Invalid component configuration: {model}")
-
- config = AzureAISearchConfig(**config_dict)
-
- tool = target_class._from_config(config)
- if expected is None:
- return tool
- return cast(ExpectedType, tool)
-
@property
def schema(self) -> ToolSchema:
"""Return the schema for the tool."""
@@ -625,245 +600,158 @@ class BaseAzureAISearchTool(BaseTool[SearchQuery, SearchResults], ABC):
}
def return_value_as_string(self, value: SearchResults) -> str:
- """Convert the search results to a string representation.
-
- This method is used to format the search results in a way that's suitable
- for display to the user or for consumption by language models.
-
- Args:
- value (List[SearchResult]): The search results to convert.
-
- Returns:
- str: A formatted string representation of the search results.
- """
+ """Convert the search results to a string representation."""
if not value.results:
return "No results found."
result_strings: List[str] = []
for i, result in enumerate(value.results, 1):
- content_str = ", ".join(f"{k}: {v}" for k, v in result.content.items())
+ content_items = [f"{k}: {str(v) if v is not None else 'None'}" for k, v in result.content.items()]
+ content_str = ", ".join(content_items)
result_strings.append(f"Result {i} (Score: {result.score:.2f}): {content_str}")
return "\n".join(result_strings)
+ @classmethod
+ def _validate_config(
+ cls, config_dict: Dict[str, Any], search_type: Literal["full_text", "vector", "hybrid"]
+ ) -> None:
+ """Validate configuration for specific search types."""
+ credential = config_dict.get("credential")
+ if isinstance(credential, str):
+ raise TypeError("Credential must be AzureKeyCredential, AsyncTokenCredential, or a valid dict")
+ if isinstance(credential, dict) and "api_key" not in credential:
+ raise ValueError("If credential is a dict, it must contain an 'api_key' key")
+
+ try:
+ _ = AzureAISearchConfig(**config_dict)
+ except Exception as e:
+ raise ValueError(f"Invalid configuration: {str(e)}") from e
+
+ if search_type == "vector":
+ vector_fields = config_dict.get("vector_fields")
+ if not vector_fields or len(vector_fields) == 0:
+ raise ValueError("vector_fields must contain at least one field name for vector search")
+
+ elif search_type == "hybrid":
+ vector_fields = config_dict.get("vector_fields")
+ search_fields = config_dict.get("search_fields")
+
+ if not vector_fields or len(vector_fields) == 0:
+ raise ValueError("vector_fields must contain at least one field name for hybrid search")
+
+ if not search_fields or len(search_fields) == 0:
+ raise ValueError("search_fields must contain at least one field name for hybrid search")
+
+ @classmethod
+ @abstractmethod
+ def _from_config(cls, config: AzureAISearchConfig) -> "BaseAzureAISearchTool":
+ """Create a tool instance from a configuration object.
+
+ This is an abstract method that must be implemented by subclasses.
+ """
+ if cls is BaseAzureAISearchTool:
+ raise NotImplementedError(
+ "BaseAzureAISearchTool is an abstract base class and cannot be instantiated directly. "
+ "Use a concrete implementation like AzureAISearchTool."
+ )
+ raise NotImplementedError("Subclasses must implement _from_config")
+
+ @abstractmethod
+ async def _get_embedding(self, query: str) -> List[float]:
+ """Generate embedding vector for the query text."""
+ raise NotImplementedError("Subclasses must implement _get_embedding")
+
_allow_private_constructor = ContextVar("_allow_private_constructor", default=False)
-class AzureAISearchTool(BaseAzureAISearchTool):
+class AzureAISearchTool(EmbeddingProviderMixin, BaseAzureAISearchTool):
"""Azure AI Search tool for querying Azure search indexes.
This tool provides a simplified interface for querying Azure AI Search indexes using
- various search methods. The tool supports four main search types:
-
- 1. Keyword Search: Traditional text-based search using Azure's text analysis
- 2. Full-Text Search: Enhanced text search with language-specific analyzers
- 3. Vector Search: Semantic similarity search using vector embeddings
- 4. Hybrid Search: Combines fulltext and vector search for comprehensive results
-
- You should use the factory methods to create instances for specific search types:
- - create_keyword_search()
- - create_full_text_search()
- - create_vector_search()
- - create_hybrid_search()
- """
+ various search methods. It's recommended to use the factory methods to create
+ instances tailored for specific search types:
- def __init__(
- self,
- name: str,
- endpoint: str,
- index_name: str,
- credential: Union[AzureKeyCredential, TokenCredential, Dict[str, str]],
- query_type: Literal["keyword", "fulltext", "vector", "semantic"],
- search_fields: Optional[List[str]] = None,
- select_fields: Optional[List[str]] = None,
- vector_fields: Optional[List[str]] = None,
- filter: Optional[str] = None,
- top: Optional[int] = 5,
- **kwargs: Any,
- ) -> None:
- if not _allow_private_constructor.get():
- raise RuntimeError(
- "Constructor is private. Use factory methods like create_keyword_search(), "
- "create_vector_search(), or create_hybrid_search() instead."
- )
+ 1. **Full-Text Search**: For traditional keyword-based searches, Lucene queries, or
+ semantically re-ranked results.
+ - Use `AzureAISearchTool.create_full_text_search()`
+ - Supports `query_type`: "simple" (keyword), "full" (Lucene), "semantic".
- super().__init__(
- name=name,
- endpoint=endpoint,
- index_name=index_name,
- credential=credential,
- query_type=query_type,
- search_fields=search_fields,
- select_fields=select_fields,
- vector_fields=vector_fields,
- filter=filter,
- top=top,
- **kwargs,
- )
+ 2. **Vector Search**: For pure similarity searches based on vector embeddings.
+ - Use `AzureAISearchTool.create_vector_search()`
- @classmethod
- @overload
- def load_component(
- cls, model: Union[ComponentModel, Dict[str, Any]], expected: None = None
- ) -> "AzureAISearchTool": ...
+ 3. **Hybrid Search**: For combining vector search with full-text or semantic search
+ to get the benefits of both.
+ - Use `AzureAISearchTool.create_hybrid_search()`
+ - The text component can be "simple", "full", or "semantic" via the `query_type` parameter.
- @classmethod
- @overload
- def load_component(
- cls, model: Union[ComponentModel, Dict[str, Any]], expected: Type[ExpectedType]
- ) -> ExpectedType: ...
+ Each factory method configures the tool with appropriate defaults and validations
+ for the chosen search strategy.
+
+ .. warning::
+ If you set `query_type="semantic"`, you must also provide a valid `semantic_config_name`.
+ This configuration must be set up in your Azure AI Search index beforehand.
+ """
+
+ component_provider_override = "autogen_ext.tools.azure.AzureAISearchTool"
@classmethod
- def load_component(
- cls, model: Union[ComponentModel, Dict[str, Any]], expected: Optional[Type[ExpectedType]] = None
- ) -> Union["AzureAISearchTool", ExpectedType]:
- """Load a component from a component model.
+ def _from_config(cls, config: AzureAISearchConfig) -> "AzureAISearchTool":
+ """Create a tool instance from a configuration object.
Args:
- model: The component model or dictionary with configuration
- expected: Optional expected return type
+ config: The configuration object with tool settings
Returns:
- An initialized AzureAISearchTool instance
+ AzureAISearchTool: An initialized tool instance
"""
token = _allow_private_constructor.set(True)
try:
- if isinstance(model, dict):
- model = ComponentModel(**model)
-
- config = model.config
-
- query_type_str = config.get("query_type", "keyword")
-
- query_type_mapping = {
- "keyword": "keyword",
- "simple": "fulltext",
- "fulltext": "fulltext",
- "vector": "vector",
- "semantic": "semantic",
- }
-
- query_type = cast(
- Literal["keyword", "fulltext", "vector", "semantic"], query_type_mapping.get(query_type_str, "vector")
- )
-
instance = cls(
- name=config.get("name", ""),
- endpoint=config.get("endpoint", ""),
- index_name=config.get("index_name", ""),
- credential=config.get("credential", {}),
- query_type=query_type,
- search_fields=config.get("search_fields"),
- select_fields=config.get("select_fields"),
- vector_fields=config.get("vector_fields"),
- top=config.get("top"),
- filter=config.get("filter"),
- enable_caching=config.get("enable_caching", False),
- cache_ttl_seconds=config.get("cache_ttl_seconds", 300),
+ name=config.name,
+ description=config.description or "",
+ endpoint=config.endpoint,
+ index_name=config.index_name,
+ credential=config.credential,
+ api_version=config.api_version,
+ query_type=config.query_type,
+ search_fields=config.search_fields,
+ select_fields=config.select_fields,
+ vector_fields=config.vector_fields,
+ top=config.top,
+ filter=config.filter,
+ semantic_config_name=config.semantic_config_name,
+ enable_caching=config.enable_caching,
+ cache_ttl_seconds=config.cache_ttl_seconds,
+ embedding_provider=config.embedding_provider,
+ embedding_model=config.embedding_model,
+ openai_api_key=config.openai_api_key,
+ openai_api_version=config.openai_api_version,
+ openai_endpoint=config.openai_endpoint,
)
-
- if expected is not None:
- return cast(ExpectedType, instance)
return instance
finally:
_allow_private_constructor.reset(token)
@classmethod
- def _validate_common_params(cls, name: str, endpoint: str, index_name: str, credential: Any) -> None:
- """Validate common parameters across all factory methods.
-
- Args:
- name: Tool name
- endpoint: Azure Search endpoint URL
- index_name: Name of search index
- credential: Authentication credentials
-
- Raises:
- ValueError: If any parameter is invalid
- """
- if not endpoint or not endpoint.startswith(("http://", "https://")):
- raise ValueError("endpoint must be a valid URL starting with http:// or https://")
-
- if not index_name:
- raise ValueError("index_name cannot be empty")
-
- if not name:
- raise ValueError("name cannot be empty")
-
- if not credential:
- raise ValueError("credential cannot be None")
-
- @classmethod
- def create_keyword_search(
- cls,
- name: str,
- endpoint: str,
- index_name: str,
- credential: Union[AzureKeyCredential, TokenCredential, Dict[str, str]],
- search_fields: Optional[List[str]] = None,
- select_fields: Optional[List[str]] = None,
- filter: Optional[str] = None,
- top: Optional[int] = 5,
- **kwargs: Any,
+ def _create_from_params(
+ cls, config_dict: Dict[str, Any], search_type: Literal["full_text", "vector", "hybrid"]
) -> "AzureAISearchTool":
- """Factory method to create a keyword search tool.
-
- Keyword search performs traditional text-based search, good for finding documents
- containing specific terms or exact matches to your query.
+ """Private helper to create an instance from parameters after validation.
Args:
- name (str): The name of the tool
- endpoint (str): The URL of your Azure AI Search service
- index_name (str): The name of the search index
- credential (Union[AzureKeyCredential, TokenCredential, Dict[str, str]]): Authentication credentials
- search_fields (Optional[List[str]]): Fields to search within for text search
- select_fields (Optional[List[str]]): Fields to include in results
- filter (Optional[str]): OData filter expression to filter results
- top (Optional[int]): Maximum number of results to return
- **kwargs (Any): Additional configuration options
+ config_dict: Dictionary with configuration parameters
+ search_type: Type of search for validation
Returns:
- An initialized keyword search tool
-
- Example Usage:
- .. code-block:: python
-
- # type: ignore
- # Example of using keyword search with Azure AI Search
- from autogen_ext.tools.azure import AzureAISearchTool
- from azure.core.credentials import AzureKeyCredential
-
- # Create a keyword search tool
- keyword_search = AzureAISearchTool.create_keyword_search(
- name="keyword_search",
- endpoint="https://your-service.search.windows.net",
- index_name="your-index",
- credential=AzureKeyCredential("your-api-key"),
- search_fields=["title", "content"],
- select_fields=["id", "title", "content", "category"],
- top=10,
- )
-
- # The search tool can be used with an Agent
- # assistant = Agent("assistant", tools=[keyword_search])
+ Configured AzureAISearchTool instance
"""
- cls._validate_common_params(name, endpoint, index_name, credential)
+ cls._validate_config(config_dict, search_type)
token = _allow_private_constructor.set(True)
try:
- return cls(
- name=name,
- endpoint=endpoint,
- index_name=index_name,
- credential=credential,
- query_type="keyword",
- search_fields=search_fields,
- select_fields=select_fields,
- filter=filter,
- top=top,
- **kwargs,
- )
+ return cls(**config_dict)
finally:
_allow_private_constructor.reset(token)
@@ -873,72 +761,114 @@ class AzureAISearchTool(BaseAzureAISearchTool):
name: str,
endpoint: str,
index_name: str,
- credential: Union[AzureKeyCredential, TokenCredential, Dict[str, str]],
+ credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]],
+ description: Optional[str] = None,
+ api_version: Optional[str] = None,
+ query_type: Literal["simple", "full", "semantic"] = "simple",
search_fields: Optional[List[str]] = None,
select_fields: Optional[List[str]] = None,
- filter: Optional[str] = None,
top: Optional[int] = 5,
- **kwargs: Any,
+ filter: Optional[str] = None,
+ semantic_config_name: Optional[str] = None,
+ enable_caching: bool = False,
+ cache_ttl_seconds: int = 300,
) -> "AzureAISearchTool":
- """Factory method to create a full-text search tool.
+ """Create a tool for traditional text-based searches.
- Full-text search uses advanced text analysis (stemming, lemmatization, etc.)
- to provide more comprehensive text matching than basic keyword search.
+ This factory method creates an AzureAISearchTool optimized for full-text search,
+ supporting keyword matching, Lucene syntax, and semantic search capabilities.
Args:
- name (str): The name of the tool
- endpoint (str): The URL of your Azure AI Search service
- index_name (str): The name of the search index
- credential (Union[AzureKeyCredential, TokenCredential, Dict[str, str]]): Authentication credentials
- search_fields (Optional[List[str]]): Fields to search within
- select_fields (Optional[List[str]]): Fields to include in results
- filter (Optional[str]): OData filter expression to filter results
- top (Optional[int]): Maximum number of results to return
- **kwargs (Any): Additional configuration options
+ name: The name of this tool instance
+ endpoint: The full URL of your Azure AI Search service
+ index_name: Name of the search index to query
+ credential: Azure credential for authentication (API key or token)
+ description: Optional description explaining the tool's purpose
+ api_version: Azure AI Search API version to use
+ query_type: Type of text search to perform:
+
+ • **simple** : Basic keyword search that matches exact terms and their variations
+ • **full**: Advanced search using Lucene query syntax for complex queries
+ • **semantic**: AI-powered search that understands meaning and context, providing enhanced relevance ranking
+ search_fields: Fields to search within documents
+ select_fields: Fields to return in search results
+ top: Maximum number of results to return (default: 5)
+ filter: OData filter expression to refine search results
+ semantic_config_name: Semantic configuration name (required for semantic query_type)
+ enable_caching: Whether to cache search results
+ cache_ttl_seconds: How long to cache results in seconds
Returns:
- An initialized full-text search tool
+ An initialized AzureAISearchTool for full-text search
- Example Usage:
+ Example:
.. code-block:: python
- # type: ignore
- # Example of using full-text search with Azure AI Search
- from autogen_ext.tools.azure import AzureAISearchTool
from azure.core.credentials import AzureKeyCredential
+ from autogen_ext.tools.azure import AzureAISearchTool
+
+ # Basic keyword search
+ tool = AzureAISearchTool.create_full_text_search(
+ name="doc-search",
+ endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint
+ index_name="", # Name of your search index
+ credential=AzureKeyCredential(""), # Your Azure AI Search admin key
+ query_type="simple", # Enable keyword search
+ search_fields=["content", "title"], # Required: fields to search within
+ select_fields=["content", "title", "url"], # Optional: fields to return
+ top=5,
+ )
- # Create a full-text search tool
- full_text_search = AzureAISearchTool.create_full_text_search(
- name="document_search",
- endpoint="https://your-search-service.search.windows.net",
- index_name="your-index",
- credential=AzureKeyCredential("your-api-key"),
- search_fields=["title", "content"],
- select_fields=["title", "content", "category", "url"],
- top=10,
+ # full text (Lucene query) search
+ full_text_tool = AzureAISearchTool.create_full_text_search(
+ name="doc-search",
+ endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint
+ index_name="", # Name of your search index
+ credential=AzureKeyCredential(""), # Your Azure AI Search admin key
+ query_type="full", # Enable Lucene query syntax
+ search_fields=["content", "title"], # Required: fields to search within
+ select_fields=["content", "title", "url"], # Optional: fields to return
+ top=5,
+ )
+
+ # Semantic search with re-ranking
+ # Note: Make sure your index has semantic configuration enabled
+ semantic_tool = AzureAISearchTool.create_full_text_search(
+ name="semantic-search",
+ endpoint="https://your-search.search.windows.net",
+ index_name="",
+ credential=AzureKeyCredential(""),
+ query_type="semantic", # Enable semantic ranking
+ semantic_config_name="", # Required for semantic search
+ search_fields=["content", "title"], # Required: fields to search within
+ select_fields=["content", "title", "url"], # Optional: fields to return
+ top=5,
)
# The search tool can be used with an Agent
- # assistant = Agent("assistant", tools=[full_text_search])
+ # assistant = Agent("assistant", tools=[semantic_tool])
"""
- cls._validate_common_params(name, endpoint, index_name, credential)
+ if query_type == "semantic" and not semantic_config_name:
+ raise ValueError("semantic_config_name is required when query_type is 'semantic'")
+
+ config_dict = {
+ "name": name,
+ "endpoint": endpoint,
+ "index_name": index_name,
+ "credential": credential,
+ "description": description,
+ "api_version": api_version or DEFAULT_API_VERSION,
+ "query_type": query_type,
+ "search_fields": search_fields,
+ "select_fields": select_fields,
+ "top": top,
+ "filter": filter,
+ "semantic_config_name": semantic_config_name,
+ "enable_caching": enable_caching,
+ "cache_ttl_seconds": cache_ttl_seconds,
+ }
- token = _allow_private_constructor.set(True)
- try:
- return cls(
- name=name,
- endpoint=endpoint,
- index_name=index_name,
- credential=credential,
- query_type="fulltext",
- search_fields=search_fields,
- select_fields=select_fields,
- filter=filter,
- top=top,
- **kwargs,
- )
- finally:
- _allow_private_constructor.reset(token)
+ return cls._create_from_params(config_dict, "full_text")
@classmethod
def create_vector_search(
@@ -946,76 +876,128 @@ class AzureAISearchTool(BaseAzureAISearchTool):
name: str,
endpoint: str,
index_name: str,
- credential: Union[AzureKeyCredential, TokenCredential, Dict[str, str]],
+ credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]],
vector_fields: List[str],
+ description: Optional[str] = None,
+ api_version: Optional[str] = None,
select_fields: Optional[List[str]] = None,
+ top: int = 5,
filter: Optional[str] = None,
- top: Optional[int] = 5,
- **kwargs: Any,
+ enable_caching: bool = False,
+ cache_ttl_seconds: int = 300,
+ embedding_provider: Optional[str] = None,
+ embedding_model: Optional[str] = None,
+ openai_api_key: Optional[str] = None,
+ openai_api_version: Optional[str] = None,
+ openai_endpoint: Optional[str] = None,
) -> "AzureAISearchTool":
- """Factory method to create a vector search tool.
+ """Create a tool for pure vector/similarity search.
- Vector search uses embedding vectors to find semantically similar content, enabling
- the discovery of related information even when different terminology is used.
+ This factory method creates an AzureAISearchTool optimized for vector search,
+ allowing for semantic similarity-based matching using vector embeddings.
Args:
- name (str): The name of the tool
- endpoint (str): The URL of your Azure AI Search service
- index_name (str): The name of the search index
- credential (Union[AzureKeyCredential, TokenCredential, Dict[str, str]]): Authentication credentials
- vector_fields (List[str]): Fields containing vector embeddings for similarity search
- select_fields (Optional[List[str]]): Fields to include in results
- filter (Optional[str]): OData filter expression to filter results
- top (Optional[int]): Maximum number of results to return
- **kwargs (Any): Additional configuration options
+ name: The name of this tool instance
+ endpoint: The full URL of your Azure AI Search service
+ index_name: Name of the search index to query
+ credential: Azure credential for authentication (API key or token)
+ vector_fields: Fields to use for vector search (required)
+ description: Optional description explaining the tool's purpose
+ api_version: Azure AI Search API version to use
+ select_fields: Fields to return in search results
+ top: Maximum number of results to return / k in k-NN (default: 5)
+ filter: OData filter expression to refine search results
+ enable_caching: Whether to cache search results
+ cache_ttl_seconds: How long to cache results in seconds
+ embedding_provider: Provider for client-side embeddings (e.g., 'azure_openai', 'openai')
+ embedding_model: Model for client-side embeddings (e.g., 'text-embedding-ada-002')
+ openai_api_key: API key for OpenAI/Azure OpenAI embeddings
+ openai_api_version: API version for Azure OpenAI embeddings
+ openai_endpoint: Endpoint URL for Azure OpenAI embeddings
Returns:
- An initialized vector search tool
+ An initialized AzureAISearchTool for vector search
+
+ Raises:
+ ValueError: If vector_fields is empty
+ ValueError: If embedding_provider is 'azure_openai' without openai_endpoint
+ ValueError: If required parameters are missing or invalid
Example Usage:
.. code-block:: python
- # type: ignore
- # Example of using vector search with Azure AI Search
- from autogen_ext.tools.azure import AzureAISearchTool
from azure.core.credentials import AzureKeyCredential
+ from autogen_ext.tools.azure import AzureAISearchTool
- # Create a vector search tool
- vector_search = AzureAISearchTool.create_vector_search(
- name="vector_search",
- endpoint="https://your-search-service.search.windows.net",
- index_name="your-index",
- credential=AzureKeyCredential("your-api-key"),
- vector_fields=["embedding"],
- select_fields=["title", "content", "url"],
+ # Vector search with service-side vectorization
+ tool = AzureAISearchTool.create_vector_search(
+ name="vector-search",
+ endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint
+ index_name="", # Name of your search index
+ credential=AzureKeyCredential(""), # Your Azure AI Search admin key
+ vector_fields=["content_vector"], # Your vector field name
+ select_fields=["content", "title", "url"], # Fields to return in results
top=5,
)
- # The search tool can be used with an Agent
- # assistant = Agent("assistant", tools=[vector_search])
+ # Vector search with Azure OpenAI embeddings
+ azure_openai_tool = AzureAISearchTool.create_vector_search(
+ name="azure-openai-vector-search",
+ endpoint="https://your-search.search.windows.net",
+ index_name="",
+ credential=AzureKeyCredential(""),
+ vector_fields=["content_vector"],
+ embedding_provider="azure_openai", # Use Azure OpenAI for embeddings
+ embedding_model="text-embedding-ada-002", # Embedding model to use
+ openai_endpoint="https://your-openai.openai.azure.com", # Your Azure OpenAI endpoint
+ openai_api_key="", # Your Azure OpenAI key
+ openai_api_version="2024-02-15-preview", # Azure OpenAI API version
+ select_fields=["content", "title", "url"], # Fields to return in results
+ top=5,
+ )
- """
- cls._validate_common_params(name, endpoint, index_name, credential)
+ # Vector search with OpenAI embeddings
+ openai_tool = AzureAISearchTool.create_vector_search(
+ name="openai-vector-search",
+ endpoint="https://your-search.search.windows.net",
+ index_name="",
+ credential=AzureKeyCredential(""),
+ vector_fields=["content_vector"],
+ embedding_provider="openai", # Use OpenAI for embeddings
+ embedding_model="text-embedding-ada-002", # Embedding model to use
+ openai_api_key="", # Your OpenAI API key
+ select_fields=["content", "title", "url"], # Fields to return in results
+ top=5,
+ )
- if not vector_fields or len(vector_fields) == 0:
- raise ValueError("vector_fields must contain at least one field name")
+ # Use the tool with an Agent
+ # assistant = Agent("assistant", tools=[azure_openai_tool])
+ """
+ if embedding_provider == "azure_openai" and not openai_endpoint:
+ raise ValueError("openai_endpoint is required when embedding_provider is 'azure_openai'")
+
+ config_dict = {
+ "name": name,
+ "endpoint": endpoint,
+ "index_name": index_name,
+ "credential": credential,
+ "description": description,
+ "api_version": api_version or DEFAULT_API_VERSION,
+ "query_type": "vector",
+ "select_fields": select_fields,
+ "vector_fields": vector_fields,
+ "top": top,
+ "filter": filter,
+ "enable_caching": enable_caching,
+ "cache_ttl_seconds": cache_ttl_seconds,
+ "embedding_provider": embedding_provider,
+ "embedding_model": embedding_model,
+ "openai_api_key": openai_api_key,
+ "openai_api_version": openai_api_version,
+ "openai_endpoint": openai_endpoint,
+ }
- token = _allow_private_constructor.set(True)
- try:
- return cls(
- name=name,
- endpoint=endpoint,
- index_name=index_name,
- credential=credential,
- query_type="vector",
- vector_fields=vector_fields,
- select_fields=select_fields,
- filter=filter,
- top=top,
- **kwargs,
- )
- finally:
- _allow_private_constructor.reset(token)
+ return cls._create_from_params(config_dict, "vector")
@classmethod
def create_hybrid_search(
@@ -1023,171 +1005,131 @@ class AzureAISearchTool(BaseAzureAISearchTool):
name: str,
endpoint: str,
index_name: str,
- credential: Union[AzureKeyCredential, TokenCredential, Dict[str, str]],
+ credential: Union[AzureKeyCredential, AsyncTokenCredential, Dict[str, str]],
vector_fields: List[str],
- search_fields: Optional[List[str]] = None,
+ search_fields: List[str],
+ description: Optional[str] = None,
+ api_version: Optional[str] = None,
+ query_type: Literal["simple", "full", "semantic"] = "simple",
select_fields: Optional[List[str]] = None,
+ top: int = 5,
filter: Optional[str] = None,
- top: Optional[int] = 5,
- **kwargs: Any,
+ semantic_config_name: Optional[str] = None,
+ enable_caching: bool = False,
+ cache_ttl_seconds: int = 300,
+ embedding_provider: Optional[str] = None,
+ embedding_model: Optional[str] = None,
+ openai_api_key: Optional[str] = None,
+ openai_api_version: Optional[str] = None,
+ openai_endpoint: Optional[str] = None,
) -> "AzureAISearchTool":
- """Factory method to create a hybrid search tool (text + vector).
+ """Create a tool that combines vector and text search capabilities.
- Hybrid search combines text search (fulltext or semantic) with vector similarity
- search to provide more comprehensive results. This is the recommended entrypoint for hybrid (text + vector) search.
- The query_type will be 'semantic' if semantic_config_name is provided, otherwise 'fulltext'.
+ This factory method creates an AzureAISearchTool configured for hybrid search,
+ which combines the benefits of vector similarity and traditional text search.
Args:
- name (str): The name of the tool
- endpoint (str): The URL of your Azure AI Search service
- index_name (str): The name of the search index
- credential (Union[AzureKeyCredential, TokenCredential, Dict[str, str]]): Authentication credentials
- vector_fields (List[str]): Fields containing vector embeddings for similarity search
- search_fields (Optional[List[str]]): Fields to search within for text search
- select_fields (Optional[List[str]]): Fields to include in results
- filter (Optional[str]): OData filter expression to filter results
- top (Optional[int]): Maximum number of results to return
- **kwargs (Any): Additional configuration options
+ name: The name of this tool instance
+ endpoint: The full URL of your Azure AI Search service
+ index_name: Name of the search index to query
+ credential: Azure credential for authentication (API key or token)
+ vector_fields: Fields to use for vector search (required)
+ search_fields: Fields to use for text search (required)
+ description: Optional description explaining the tool's purpose
+ api_version: Azure AI Search API version to use
+ query_type: Type of text search to perform:
+
+ • **simple**: Basic keyword search that matches exact terms and their variations
+ • **full**: Advanced search using Lucene query syntax for complex queries
+ • **semantic**: AI-powered search that understands meaning and context, providing enhanced relevance ranking
+ select_fields: Fields to return in search results
+ top: Maximum number of results to return (default: 5)
+ filter: OData filter expression to refine search results
+ semantic_config_name: Semantic configuration name (required if query_type="semantic")
+ enable_caching: Whether to cache search results
+ cache_ttl_seconds: How long to cache results in seconds
+ embedding_provider: Provider for client-side embeddings (e.g., 'azure_openai', 'openai')
+ embedding_model: Model for client-side embeddings (e.g., 'text-embedding-ada-002')
+ openai_api_key: API key for OpenAI/Azure OpenAI embeddings
+ openai_api_version: API version for Azure OpenAI embeddings
+ openai_endpoint: Endpoint URL for Azure OpenAI embeddings
Returns:
- An initialized hybrid search tool
+ An initialized AzureAISearchTool for hybrid search
- Example Usage:
+ Raises:
+ ValueError: If vector_fields or search_fields is empty
+ ValueError: If query_type is "semantic" without semantic_config_name
+ ValueError: If embedding_provider is 'azure_openai' without openai_endpoint
+ ValueError: If required parameters are missing or invalid
+
+ Example:
.. code-block:: python
- # type: ignore
- # Example of using hybrid search with Azure AI Search
- from autogen_ext.tools.azure import AzureAISearchTool
from azure.core.credentials import AzureKeyCredential
+ from autogen_ext.tools.azure import AzureAISearchTool
- # Create a hybrid search tool
- hybrid_search = AzureAISearchTool.create_hybrid_search(
- name="hybrid_search",
- endpoint="https://your-search-service.search.windows.net",
- index_name="your-index",
- credential=AzureKeyCredential("your-api-key"),
- vector_fields=["embedding_field"],
- search_fields=["title", "content"],
- select_fields=["title", "content", "url", "date"],
- top=10,
+ # Basic hybrid search with service-side vectorization
+ tool = AzureAISearchTool.create_hybrid_search(
+ name="hybrid-search",
+ endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint
+ index_name="", # Name of your search index
+ credential=AzureKeyCredential(""), # Your Azure AI Search admin key
+ vector_fields=["content_vector"], # Your vector field name
+ search_fields=["content", "title"], # Your searchable fields
+ top=5,
)
- # The search tool can be used with an Agent
- # assistant = Agent("researcher", tools=[hybrid_search])
-
-
- .. warning::
-
- If you set ``query_type=\"semantic\"``, you must also provide a valid ``semantic_config_name``.
- If you do not, the tool will default to the config name ``\"semantic\"``.
-
- """
- cls._validate_common_params(name, endpoint, index_name, credential)
-
- if not vector_fields or len(vector_fields) == 0:
- raise ValueError("vector_fields must contain at least one field name")
-
- token = _allow_private_constructor.set(True)
- try:
- if kwargs.get("semantic_config_name"):
- text_query_type = "semantic"
- else:
- text_query_type = "fulltext"
-
- from typing import cast
-
- return cls(
- name=name,
- endpoint=endpoint,
- index_name=index_name,
- credential=credential,
- query_type=cast(Literal["keyword", "fulltext", "vector", "semantic"], text_query_type),
- search_fields=search_fields,
- select_fields=select_fields,
- vector_fields=vector_fields,
- filter=filter,
- top=top,
- **kwargs,
- )
- finally:
- _allow_private_constructor.reset(token)
-
- async def _get_embedding(self, query: str) -> List[float]:
- """Generate embedding vector for the query text.
-
- This method handles generating embeddings for vector search functionality.
- The embedding provider and model should be specified in the tool configuration.
-
- Args:
- query (str): The text to generate embeddings for.
-
- Returns:
- List[float]: The embedding vector as a list of floats.
-
- Raises:
- ValueError: If the embedding configuration is missing or invalid.
- """
- embedding_provider = getattr(self.search_config, "embedding_provider", None)
- embedding_model = getattr(self.search_config, "embedding_model", None)
-
- if not embedding_provider or not embedding_model:
- raise ValueError(
- "To use vector search, you must provide embedding_provider and embedding_model in the configuration."
- ) from None
-
- if embedding_provider.lower() == "azure_openai":
- try:
- from azure.identity import DefaultAzureCredential
- from openai import AsyncAzureOpenAI
- except ImportError:
- raise ImportError(
- "Azure OpenAI SDK is required for embedding generation. "
- "Please install it with: uv add openai azure-identity"
- ) from None
-
- api_key = None
- if hasattr(self.search_config, "openai_api_key"):
- api_key = self.search_config.openai_api_key
-
- api_version = getattr(self.search_config, "openai_api_version", "2023-05-15")
- endpoint = getattr(self.search_config, "openai_endpoint", None)
-
- if not endpoint:
- raise ValueError("OpenAI endpoint must be provided for Azure OpenAI embeddings") from None
-
- if api_key:
- azure_client = AsyncAzureOpenAI(api_key=api_key, api_version=api_version, azure_endpoint=endpoint)
- else:
-
- def get_token() -> str:
- credential = DefaultAzureCredential()
- return credential.get_token("https://cognitiveservices.azure.com/.default").token
-
- azure_client = AsyncAzureOpenAI(
- azure_ad_token_provider=get_token, api_version=api_version, azure_endpoint=endpoint
+ # Hybrid search with semantic ranking and Azure OpenAI embeddings
+ semantic_tool = AzureAISearchTool.create_hybrid_search(
+ name="semantic-hybrid-search",
+ endpoint="https://your-search.search.windows.net",
+ index_name="",
+ credential=AzureKeyCredential(""),
+ vector_fields=["content_vector"],
+ search_fields=["content", "title"],
+ query_type="semantic", # Enable semantic ranking
+ semantic_config_name="", # Your semantic config name
+ embedding_provider="azure_openai", # Use Azure OpenAI for embeddings
+ embedding_model="text-embedding-ada-002", # Embedding model to use
+ openai_endpoint="https://your-openai.openai.azure.com", # Your Azure OpenAI endpoint
+ openai_api_key="", # Your Azure OpenAI key
+ openai_api_version="2024-02-15-preview", # Azure OpenAI API version
+ select_fields=["content", "title", "url"], # Fields to return in results
+ filter="language eq 'en'", # Optional OData filter
+ top=5,
)
- response = await azure_client.embeddings.create(model=embedding_model, input=query)
- return response.data[0].embedding
-
- elif embedding_provider.lower() == "openai":
- try:
- from openai import AsyncOpenAI
- except ImportError:
- raise ImportError(
- "OpenAI SDK is required for embedding generation. " "Please install it with: uv add openai"
- ) from None
-
- api_key = None
- if hasattr(self.search_config, "openai_api_key"):
- api_key = self.search_config.openai_api_key
-
- openai_client = AsyncOpenAI(api_key=api_key)
+ # The search tool can be used with an Agent
+ # assistant = Agent("assistant", tools=[semantic_tool])
+ """
+ if query_type == "semantic" and not semantic_config_name:
+ raise ValueError("semantic_config_name is required when query_type is 'semantic'")
+
+ if embedding_provider == "azure_openai" and not openai_endpoint:
+ raise ValueError("openai_endpoint is required when embedding_provider is 'azure_openai'")
+
+ config_dict = {
+ "name": name,
+ "endpoint": endpoint,
+ "index_name": index_name,
+ "credential": credential,
+ "description": description,
+ "api_version": api_version or DEFAULT_API_VERSION,
+ "query_type": query_type,
+ "search_fields": search_fields,
+ "select_fields": select_fields,
+ "vector_fields": vector_fields,
+ "top": top,
+ "filter": filter,
+ "semantic_config_name": semantic_config_name,
+ "enable_caching": enable_caching,
+ "cache_ttl_seconds": cache_ttl_seconds,
+ "embedding_provider": embedding_provider,
+ "embedding_model": embedding_model,
+ "openai_api_key": openai_api_key,
+ "openai_api_version": openai_api_version,
+ "openai_endpoint": openai_endpoint,
+ }
- response = await openai_client.embeddings.create(model=embedding_model, input=query)
- return response.data[0].embedding
- else:
- raise ValueError(
- f"Unsupported embedding provider: {embedding_provider}. "
- "Currently supported providers are 'azure_openai' and 'openai'."
- ) from None
+ return cls._create_from_params(config_dict, "hybrid")
diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/azure/_config.py b/python/packages/autogen-ext/src/autogen_ext/tools/azure/_config.py
index a27fdd677..38fa1fb15 100644
--- a/python/packages/autogen-ext/src/autogen_ext/tools/azure/_config.py
+++ b/python/packages/autogen-ext/src/autogen_ext/tools/azure/_config.py
@@ -6,173 +6,180 @@ settings for authentication, search behavior, retry policies, and caching.
import logging
from typing import (
- Any,
- Dict,
List,
Literal,
Optional,
- Type,
TypeVar,
Union,
)
-from azure.core.credentials import AzureKeyCredential, TokenCredential
-from pydantic import BaseModel, Field, model_validator
-
-# Add explicit ignore for the specific model validator error
-# pyright: reportArgumentType=false
-# pyright: reportUnknownArgumentType=false
-# pyright: reportUnknownVariableType=false
+from azure.core.credentials import AzureKeyCredential
+from azure.core.credentials_async import AsyncTokenCredential
+from pydantic import BaseModel, Field, field_validator, model_validator
T = TypeVar("T", bound="AzureAISearchConfig")
logger = logging.getLogger(__name__)
+QueryTypeLiteral = Literal["simple", "full", "semantic", "vector"]
+DEFAULT_API_VERSION = "2023-10-01-preview"
+
class AzureAISearchConfig(BaseModel):
- """Configuration for Azure AI Search tool.
+ """Configuration for Azure AI Search with validation.
- This class defines the configuration parameters for :class:`AzureAISearchTool`.
- It provides options for customizing search behavior including query types,
- field selection, authentication, retry policies, and caching strategies.
+ This class defines the configuration parameters for Azure AI Search tools, including
+ authentication, search behavior, caching, and embedding settings.
.. note::
-
- This class requires the :code:`azure` extra for the :code:`autogen-ext` package.
+ This class requires the ``azure`` extra for the ``autogen-ext`` package.
.. code-block:: bash
pip install -U "autogen-ext[azure]"
- Example:
+ .. note::
+ **Prerequisites:**
+
+ 1. An Azure AI Search service must be created in your Azure subscription.
+ 2. The search index must be properly configured for your use case:
+
+ - For vector search: Index must have vector fields
+ - For semantic search: Index must have semantic configuration
+ - For hybrid search: Both vector fields and text fields must be configured
+ 3. Required packages:
+
+ - Base functionality: ``azure-search-documents>=11.4.0``
+ - For Azure OpenAI embeddings: ``openai azure-identity``
+ - For OpenAI embeddings: ``openai``
+
+ Example Usage:
.. code-block:: python
from azure.core.credentials import AzureKeyCredential
from autogen_ext.tools.azure import AzureAISearchConfig
+ # Basic configuration for full-text search
config = AzureAISearchConfig(
- name="doc_search",
- endpoint="https://my-search.search.windows.net",
- index_name="my-index",
+ name="doc-search",
+ endpoint="https://your-search.search.windows.net", # Your Azure AI Search endpoint
+ index_name="", # Name of your search index
+ credential=AzureKeyCredential(""), # Your Azure AI Search admin key
+ query_type="simple",
+ search_fields=["content", "title"], # Update with your searchable fields
+ top=5,
+ )
+
+ # Configuration for vector search with Azure OpenAI embeddings
+ vector_config = AzureAISearchConfig(
+ name="vector-search",
+ endpoint="https://your-search.search.windows.net",
+ index_name="",
credential=AzureKeyCredential(""),
query_type="vector",
- vector_fields=["embedding"],
+ vector_fields=["embedding"], # Update with your vector field name
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ openai_endpoint="https://your-openai.openai.azure.com", # Your Azure OpenAI endpoint
+ openai_api_key="", # Your Azure OpenAI key
+ top=5,
)
- For more details, see:
- * `Azure AI Search Overview `_
- * `Vector Search `_
-
- Args:
- name (str): Name for the tool instance, used to identify it in the agent's toolkit.
- description (Optional[str]): Human-readable description of what this tool does and how to use it.
- endpoint (str): The full URL of your Azure AI Search service, in the format
- 'https://.search.windows.net'.
- index_name (str): Name of the target search index in your Azure AI Search service.
- The index must be pre-created and properly configured.
- api_version (str): Azure AI Search REST API version to use. Defaults to '2023-11-01'.
- Only change if you need specific features from a different API version.
- credential (Union[AzureKeyCredential, TokenCredential]): Azure authentication credential:
- - AzureKeyCredential: For API key authentication (admin/query key)
- - TokenCredential: For Azure AD authentication (e.g., DefaultAzureCredential)
- query_type (Literal["keyword", "fulltext", "vector", "semantic"]): The search query mode to use:
- - 'keyword': Basic keyword search (default)
- - 'fulltext': Full Lucene query syntax
- - 'vector': Vector similarity search
- - 'semantic': Semantic search using semantic configuration
- search_fields (Optional[List[str]]): List of index fields to search within. If not specified,
- searches all searchable fields. Example: ['title', 'content'].
- select_fields (Optional[List[str]]): Fields to return in search results. If not specified,
- returns all fields. Use to optimize response size.
- vector_fields (Optional[List[str]]): Vector field names for vector search. Must be configured
- in your search index as vector fields. Required for vector search.
- top (Optional[int]): Maximum number of documents to return in search results.
- Helps control response size and processing time.
- retry_enabled (bool): Whether to enable retry policy for transient errors. Defaults to True.
- retry_max_attempts (Optional[int]): Maximum number of retry attempts for failed requests. Defaults to 3.
- retry_mode (Literal["fixed", "exponential"]): Retry backoff strategy: fixed or exponential. Defaults to "exponential".
- enable_caching (bool): Whether to enable client-side caching of search results. Defaults to False.
- cache_ttl_seconds (int): Time-to-live for cached search results in seconds. Defaults to 300 (5 minutes).
- filter (Optional[str]): OData filter expression to refine search results.
+ # Configuration for hybrid search with semantic ranking
+ hybrid_config = AzureAISearchConfig(
+ name="hybrid-search",
+ endpoint="https://your-search.search.windows.net",
+ index_name="",
+ credential=AzureKeyCredential(""),
+ query_type="semantic",
+ semantic_config_name="", # Name of your semantic configuration
+ search_fields=["content", "title"], # Update with your search fields
+ vector_fields=["embedding"], # Update with your vector field name
+ embedding_provider="openai",
+ embedding_model="text-embedding-ada-002",
+ openai_api_key="", # Your OpenAI API key
+ top=5,
+ )
"""
- name: str = Field(description="The name of the tool")
- description: Optional[str] = Field(default=None, description="A description of the tool")
- endpoint: str = Field(description="The endpoint URL for your Azure AI Search service")
- index_name: str = Field(description="The name of the search index to query")
- api_version: str = Field(default="2023-11-01", description="API version to use")
- credential: Union[AzureKeyCredential, TokenCredential] = Field(
- description="The credential to use for authentication"
+ name: str = Field(description="The name of this tool instance")
+ description: Optional[str] = Field(default=None, description="Description explaining the tool's purpose")
+ endpoint: str = Field(description="The full URL of your Azure AI Search service")
+ index_name: str = Field(description="Name of the search index to query")
+ credential: Union[AzureKeyCredential, AsyncTokenCredential] = Field(
+ description="Azure credential for authentication (API key or token)"
)
- query_type: Literal["keyword", "fulltext", "vector", "semantic"] = Field(
- default="keyword",
- description="Type of query to perform (keyword for classic, fulltext for Lucene, vector for embedding, semantic for semantic/AI search)",
+ api_version: str = Field(
+ default=DEFAULT_API_VERSION,
+ description=f"Azure AI Search API version to use. Defaults to {DEFAULT_API_VERSION}.",
)
- search_fields: Optional[List[str]] = Field(default=None, description="Optional list of fields to search in")
- select_fields: Optional[List[str]] = Field(default=None, description="Optional list of fields to return in results")
- vector_fields: Optional[List[str]] = Field(
- default=None, description="Optional list of vector fields for vector search"
+ query_type: QueryTypeLiteral = Field(
+ default="simple", description="Type of search to perform: simple, full, semantic, or vector"
)
- top: Optional[int] = Field(default=None, description="Optional number of results to return")
- filter: Optional[str] = Field(default=None, description="Optional OData filter expression to refine search results")
-
- retry_enabled: bool = Field(default=True, description="Whether to enable retry policy for transient errors")
- retry_max_attempts: Optional[int] = Field(
- default=3, description="Maximum number of retry attempts for failed requests"
+ search_fields: Optional[List[str]] = Field(default=None, description="Fields to search within documents")
+ select_fields: Optional[List[str]] = Field(default=None, description="Fields to return in search results")
+ vector_fields: Optional[List[str]] = Field(default=None, description="Fields to use for vector search")
+ top: Optional[int] = Field(
+ default=None, description="Maximum number of results to return. For vector searches, acts as k in k-NN."
)
- retry_mode: Literal["fixed", "exponential"] = Field(
- default="exponential",
- description="Retry backoff strategy: fixed or exponential",
+ filter: Optional[str] = Field(default=None, description="OData filter expression to refine search results")
+ semantic_config_name: Optional[str] = Field(
+ default=None, description="Semantic configuration name for enhanced results"
)
- enable_caching: bool = Field(
- default=False,
- description="Whether to enable client-side caching of search results",
- )
- cache_ttl_seconds: int = Field(
- default=300, # 5 minutes
- description="Time-to-live for cached search results in seconds",
- )
+ enable_caching: bool = Field(default=False, description="Whether to cache search results")
+ cache_ttl_seconds: int = Field(default=300, description="How long to cache results in seconds")
embedding_provider: Optional[str] = Field(
- default=None,
- description="Name of embedding provider to use (e.g., 'azure_openai', 'openai')",
- )
- embedding_model: Optional[str] = Field(default=None, description="Model name to use for generating embeddings")
- embedding_dimension: Optional[int] = Field(
- default=None, description="Dimension of embedding vectors produced by the model"
+ default=None, description="Name of embedding provider for client-side embeddings"
)
+ embedding_model: Optional[str] = Field(default=None, description="Model name for client-side embeddings")
+ openai_api_key: Optional[str] = Field(default=None, description="API key for OpenAI/Azure OpenAI embeddings")
+ openai_api_version: Optional[str] = Field(default=None, description="API version for Azure OpenAI embeddings")
+ openai_endpoint: Optional[str] = Field(default=None, description="Endpoint URL for Azure OpenAI embeddings")
model_config = {"arbitrary_types_allowed": True}
- @classmethod
- @model_validator(mode="before")
- def validate_credentials(cls: Type[T], data: Any) -> Any:
- """Validate and convert credential data."""
- if not isinstance(data, dict):
- return data
-
- result = {}
-
- for key, value in data.items():
- result[str(key)] = value
-
- if "credential" in result:
- credential = result["credential"]
-
- if isinstance(credential, dict) and "api_key" in credential:
- api_key = str(credential["api_key"])
- result["credential"] = AzureKeyCredential(api_key)
-
- return result
-
- def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
- """Custom model_dump to handle credentials."""
- result: Dict[str, Any] = super().model_dump(**kwargs)
-
- if isinstance(self.credential, AzureKeyCredential):
- result["credential"] = {"type": "AzureKeyCredential"}
- elif isinstance(self.credential, TokenCredential):
- result["credential"] = {"type": "TokenCredential"}
-
- return result
+ @field_validator("endpoint")
+ def validate_endpoint(cls, v: str) -> str:
+ """Validate that the endpoint is a valid URL."""
+ if not v.startswith(("http://", "https://")):
+ raise ValueError("endpoint must be a valid URL starting with http:// or https://")
+ return v
+
+ @field_validator("query_type")
+ def normalize_query_type(cls, v: QueryTypeLiteral) -> QueryTypeLiteral:
+ """Normalize query type to standard values."""
+ if not v:
+ return "simple"
+
+ if isinstance(v, str) and v.lower() == "fulltext":
+ return "full"
+
+ return v
+
+ @field_validator("top")
+ def validate_top(cls, v: Optional[int]) -> Optional[int]:
+ """Ensure top is a positive integer if provided."""
+ if v is not None and v <= 0:
+ raise ValueError("top must be a positive integer")
+ return v
+
+ @model_validator(mode="after")
+ def validate_interdependent_fields(self) -> "AzureAISearchConfig":
+ """Validate interdependent fields after all fields have been parsed."""
+ if self.query_type == "semantic" and not self.semantic_config_name:
+ raise ValueError("semantic_config_name must be provided when query_type is 'semantic'")
+
+ if self.query_type == "vector" and not self.vector_fields:
+ raise ValueError("vector_fields must be provided for vector search")
+
+ if (
+ self.embedding_provider
+ and self.embedding_provider.lower() == "azure_openai"
+ and self.embedding_model
+ and not self.openai_endpoint
+ ):
+ raise ValueError("openai_endpoint must be provided for azure_openai embedding provider")
+
+ return self
diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py
index 15fe73e94..139a06748 100644
--- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py
+++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py
@@ -74,21 +74,52 @@ class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]):
await session.initialize()
return await self._run(args=kwargs, cancellation_token=cancellation_token, session=session)
+ def _normalize_payload_to_content_list(
+ self, payload: list[TextContent | ImageContent | EmbeddedResource]
+ ) -> list[TextContent | ImageContent | EmbeddedResource]:
+ """
+ Normalizes a raw tool output payload into a list of content items.
+ - If payload is already a list of (TextContent, ImageContent, EmbeddedResource), it's returned as is.
+ - If payload is a single TextContent, ImageContent, or EmbeddedResource, it's wrapped in a list.
+ - If payload is a string, it's wrapped in [TextContent(text=payload)].
+ - Otherwise, the payload is stringified and wrapped in [TextContent(text=str(payload))].
+ """
+ if isinstance(payload, list) and all(
+ isinstance(item, (TextContent, ImageContent, EmbeddedResource)) for item in payload
+ ):
+ return payload
+ elif isinstance(payload, (TextContent, ImageContent, EmbeddedResource)):
+ return [payload]
+ elif isinstance(payload, str):
+ return [TextContent(text=payload, type="text")]
+ else:
+ return [TextContent(text=str(payload), type="text")]
+
async def _run(self, args: Dict[str, Any], cancellation_token: CancellationToken, session: ClientSession) -> Any:
+ exceptions_to_catch: tuple[Type[BaseException], ...]
+ if hasattr(builtins, "ExceptionGroup"):
+ exceptions_to_catch = (asyncio.CancelledError, builtins.ExceptionGroup)
+ else:
+ exceptions_to_catch = (asyncio.CancelledError,)
+
try:
if cancellation_token.is_cancelled():
- raise Exception("Operation cancelled")
+ raise asyncio.CancelledError("Operation cancelled")
result_future = asyncio.ensure_future(session.call_tool(name=self._tool.name, arguments=args))
cancellation_token.link_future(result_future)
result = await result_future
+ normalized_content_list = self._normalize_payload_to_content_list(result.content)
+
if result.isError:
- raise Exception(f"MCP tool execution failed: {result.content}")
- return result.content
- except Exception as e:
- error_message = self._format_errors(e)
- raise Exception(error_message) from e
+ serialized_error_message = self.return_value_as_string(normalized_content_list)
+ raise Exception(serialized_error_message)
+ return normalized_content_list
+
+ except exceptions_to_catch:
+ # Re-raise these specific exception types directly.
+ raise
@classmethod
async def from_server_params(cls, server_params: TServerParams, tool_name: str) -> "McpToolAdapter[TServerParams]":
@@ -138,16 +169,3 @@ class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]):
return {}
return json.dumps([serialize_item(item) for item in value])
-
- def _format_errors(self, error: Exception) -> str:
- """Recursively format errors into a string."""
-
- error_message = ""
- if hasattr(builtins, "ExceptionGroup") and isinstance(error, builtins.ExceptionGroup):
- # ExceptionGroup is available in Python 3.11+.
- # TODO: how to make this compatible with Python 3.10?
- for sub_exception in error.exceptions: # type: ignore
- error_message += self._format_errors(sub_exception) # type: ignore
- else:
- error_message += f"{str(error)}\n"
- return error_message
diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py
index 419ae34bf..fd43ff852 100644
--- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py
+++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_workbench.py
@@ -1,3 +1,4 @@
+import asyncio
import builtins
import warnings
from typing import Any, List, Literal, Mapping
@@ -152,6 +153,7 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
self._server_params = server_params
# self._session: ClientSession | None = None
self._actor: McpSessionActor | None = None
+ self._actor_loop: asyncio.AbstractEventLoop | None = None
self._read = None
self._write = None
@@ -253,6 +255,7 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
if isinstance(self._server_params, (StdioServerParams, SseServerParams)):
self._actor = McpSessionActor(self._server_params)
await self._actor.initialize()
+ self._actor_loop = asyncio.get_event_loop()
else:
raise ValueError(f"Unsupported server params type: {type(self._server_params)}")
@@ -282,4 +285,10 @@ class McpWorkbench(Workbench, Component[McpWorkbenchConfig]):
def __del__(self) -> None:
# Ensure the actor is stopped when the workbench is deleted
- pass
+ if self._actor and self._actor_loop:
+ loop = self._actor_loop
+ if loop.is_running() and not loop.is_closed():
+ loop.call_soon_threadsafe(lambda: asyncio.create_task(self.stop()))
+ else:
+ msg = "Cannot safely stop actor at [McpWorkbench.__del__]: loop is closed or not running"
+ warnings.warn(msg, RuntimeWarning, stacklevel=2)
diff --git a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py
index b374f5371..81c890efa 100644
--- a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py
+++ b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py
@@ -1,8 +1,8 @@
# mypy: disable-error-code="no-any-unimported"
import asyncio
import os
-import sys
import shutil
+import sys
import tempfile
from pathlib import Path
from typing import AsyncGenerator, TypeAlias
diff --git a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py
index decb0441b..291bd57ef 100644
--- a/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py
+++ b/python/packages/autogen-ext/tests/models/test_anthropic_model_client.py
@@ -10,12 +10,17 @@ from autogen_core.models import (
CreateResult,
FunctionExecutionResult,
FunctionExecutionResultMessage,
+ ModelInfo,
SystemMessage,
UserMessage,
)
from autogen_core.models._types import LLMMessage
from autogen_core.tools import FunctionTool
-from autogen_ext.models.anthropic import AnthropicChatCompletionClient
+from autogen_ext.models.anthropic import (
+ AnthropicBedrockChatCompletionClient,
+ AnthropicChatCompletionClient,
+ BedrockInfo,
+)
def _pass_function(input: str) -> str:
@@ -46,6 +51,29 @@ async def test_anthropic_serialization_api_key() -> None:
client2 = AnthropicChatCompletionClient.load_component(config)
assert client2
+ bedrock_client = AnthropicBedrockChatCompletionClient(
+ model="claude-3-haiku-20240307", # Use haiku for faster/cheaper testing
+ api_key="sk-password",
+ model_info=ModelInfo(
+ vision=False, function_calling=True, json_output=False, family="unknown", structured_output=True
+ ),
+ bedrock_info=BedrockInfo(
+ aws_access_key="",
+ aws_secret_key="",
+ aws_session_token="",
+ aws_region="",
+ ),
+ )
+ assert bedrock_client
+ bedrock_config = bedrock_client.dump_component()
+ assert bedrock_config
+ assert "sk-password" not in str(bedrock_config)
+ serialized_bedrock_config = bedrock_config.model_dump_json()
+ assert serialized_bedrock_config
+ assert "sk-password" not in serialized_bedrock_config
+ bedrock_client2 = AnthropicBedrockChatCompletionClient.load_component(bedrock_config)
+ assert bedrock_client2
+
@pytest.mark.asyncio
async def test_anthropic_basic_completion(caplog: pytest.LogCaptureFixture) -> None:
diff --git a/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py b/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py
index 2f4f02aea..2e1593890 100644
--- a/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py
+++ b/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py
@@ -3,7 +3,7 @@ import logging
import os
from datetime import datetime
from typing import Any, AsyncGenerator, List, Type, Union
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
import pytest
from autogen_core import CancellationToken, FunctionCall, Image
@@ -570,7 +570,7 @@ def thought_with_tool_call_stream_client(monkeypatch: pytest.MonkeyPatch) -> Azu
)
mock_client = MagicMock()
- mock_client.close = MagicMock()
+ mock_client.close = AsyncMock()
async def mock_complete(*args: Any, **kwargs: Any) -> Any:
if kwargs.get("stream", False):
diff --git a/python/packages/autogen-ext/tests/models/test_ollama_chat_completion_client.py b/python/packages/autogen-ext/tests/models/test_ollama_chat_completion_client.py
index c020c6f93..de2db33dd 100644
--- a/python/packages/autogen-ext/tests/models/test_ollama_chat_completion_client.py
+++ b/python/packages/autogen-ext/tests/models/test_ollama_chat_completion_client.py
@@ -570,7 +570,7 @@ async def test_ollama_create_structured_output(model: str, ollama_client: Ollama
@pytest.mark.asyncio
-@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b"])
+@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b", "qwen3:0.6b"])
async def test_ollama_create_tools(model: str, ollama_client: OllamaChatCompletionClient) -> None:
def add(x: int, y: int) -> str:
return str(x + y)
@@ -653,7 +653,7 @@ async def test_ollama_create_structured_output_with_tools(
@pytest.mark.asyncio
-@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b"])
+@pytest.mark.parametrize("model", ["qwen2.5:0.5b", "llama3.2:1b", "qwen3:0.6b"])
async def test_ollama_create_stream_tools(model: str, ollama_client: OllamaChatCompletionClient) -> None:
def add(x: int, y: int) -> str:
return str(x + y)
diff --git a/python/packages/autogen-ext/tests/models/test_openai_model_client.py b/python/packages/autogen-ext/tests/models/test_openai_model_client.py
index a54cffb77..1f824274e 100644
--- a/python/packages/autogen-ext/tests/models/test_openai_model_client.py
+++ b/python/packages/autogen-ext/tests/models/test_openai_model_client.py
@@ -2485,4 +2485,17 @@ async def test_multimodal_message_test(
_ = await ocr_agent.run(task=multi_modal_message)
+@pytest.mark.asyncio
+async def test_mistral_remove_name() -> None:
+ # Test that the name pramaeter is removed from the message
+ # when the model is Mistral
+ message = UserMessage(content="foo", source="user")
+ params = to_oai_type(message, prepend_name=False, model="mistral-7b", model_family=ModelFamily.MISTRAL)
+ assert ("name" in params[0]) is False
+
+ # when the model is gpt-4o, the name parameter is not removed
+ params = to_oai_type(message, prepend_name=False, model="gpt-4o", model_family=ModelFamily.GPT_4O)
+ assert ("name" in params[0]) is True
+
+
# TODO: add integration tests for Azure OpenAI using AAD token.
diff --git a/python/packages/autogen-ext/tests/test_worker_runtime.py b/python/packages/autogen-ext/tests/test_worker_runtime.py
index dede30685..ec57f187e 100644
--- a/python/packages/autogen-ext/tests/test_worker_runtime.py
+++ b/python/packages/autogen-ext/tests/test_worker_runtime.py
@@ -577,6 +577,139 @@ async def test_grpc_max_message_size() -> None:
await host.stop()
+@pytest.mark.grpc
+@pytest.mark.asyncio
+async def test_agent_type_register_instance() -> None:
+ host_address = "localhost:50051"
+ agent1_id = AgentId(type="name", key="default")
+ agentdup_id = AgentId(type="name", key="default")
+ agent2_id = AgentId(type="name", key="notdefault")
+ host = GrpcWorkerAgentRuntimeHost(address=host_address)
+ host.start()
+
+ worker = GrpcWorkerAgentRuntime(host_address=host_address)
+ agent1 = NoopAgent()
+ agent2 = NoopAgent()
+ agentdup = NoopAgent()
+ await worker.start()
+
+ await worker.register_agent_instance(agent1, agent_id=agent1_id)
+ await worker.register_agent_instance(agent2, agent_id=agent2_id)
+
+ with pytest.raises(ValueError):
+ await worker.register_agent_instance(agentdup, agent_id=agentdup_id)
+
+ assert await worker.try_get_underlying_agent_instance(agent1_id, type=NoopAgent) == agent1
+ assert await worker.try_get_underlying_agent_instance(agent2_id, type=NoopAgent) == agent2
+
+ await worker.stop()
+ await host.stop()
+
+
+@pytest.mark.grpc
+@pytest.mark.asyncio
+async def test_agent_type_register_instance_different_types() -> None:
+ host_address = "localhost:50051"
+ agent1_id = AgentId(type="name", key="noop")
+ agent2_id = AgentId(type="name", key="loopback")
+ host = GrpcWorkerAgentRuntimeHost(address=host_address)
+ host.start()
+
+ worker = GrpcWorkerAgentRuntime(host_address=host_address)
+ agent1 = NoopAgent()
+ agent2 = LoopbackAgent()
+ await worker.start()
+
+ await worker.register_agent_instance(agent1, agent_id=agent1_id)
+ with pytest.raises(ValueError):
+ await worker.register_agent_instance(agent2, agent_id=agent2_id)
+
+ await worker.stop()
+ await host.stop()
+
+
+@pytest.mark.grpc
+@pytest.mark.asyncio
+async def test_register_instance_factory() -> None:
+ host_address = "localhost:50051"
+ agent1_id = AgentId(type="name", key="default")
+ host = GrpcWorkerAgentRuntimeHost(address=host_address)
+ host.start()
+
+ worker = GrpcWorkerAgentRuntime(host_address=host_address)
+ agent1 = NoopAgent()
+ await worker.start()
+
+ await agent1.register_instance(runtime=worker, agent_id=agent1_id)
+
+ with pytest.raises(ValueError):
+ await NoopAgent.register(runtime=worker, type="name", factory=lambda: NoopAgent())
+
+ await worker.stop()
+ await host.stop()
+
+
+@pytest.mark.grpc
+@pytest.mark.asyncio
+async def test_instance_factory_messaging() -> None:
+ host_address = "localhost:50051"
+ loopback_agent_id = AgentId(type="dm_agent", key="dm_agent")
+ cascading_agent_id = AgentId(type="instance_agent", key="instance_agent")
+ host = GrpcWorkerAgentRuntimeHost(address=host_address)
+ host.start()
+
+ worker = GrpcWorkerAgentRuntime(host_address=host_address)
+ cascading_agent = CascadingAgent(max_rounds=5)
+ loopback_agent = LoopbackAgent()
+ await worker.start()
+
+ await loopback_agent.register_instance(worker, agent_id=loopback_agent_id)
+ resp = await worker.send_message(message=ContentMessage(content="Hello!"), recipient=loopback_agent_id)
+ assert resp == ContentMessage(content="Hello!")
+
+ await cascading_agent.register_instance(worker, agent_id=cascading_agent_id)
+ await CascadingAgent.register(worker, "factory_agent", lambda: CascadingAgent(max_rounds=5))
+
+ # instance_agent will publish a message that factory_agent will pick up
+ for i in range(5):
+ await worker.publish_message(
+ CascadingMessageType(round=i + 1), TopicId(type="instance_agent", source="instance_agent")
+ )
+ await asyncio.sleep(2)
+
+ agent = await worker.try_get_underlying_agent_instance(AgentId("factory_agent", "default"), CascadingAgent)
+ assert agent.num_calls == 4
+ assert cascading_agent.num_calls == 5
+
+ await worker.stop()
+ await host.stop()
+
+
+# GrpcWorkerAgentRuntimeHost eats exceptions in the main loop
+# @pytest.mark.grpc
+# @pytest.mark.asyncio
+# async def test_agent_type_register_instance_publish_new_source() -> None:
+# host_address = "localhost:50056"
+# agent_id = AgentId(type="name", key="default")
+# agent1 = LoopbackAgent()
+# host = GrpcWorkerAgentRuntimeHost(address=host_address)
+# host.start()
+# worker = GrpcWorkerAgentRuntime(host_address=host_address)
+# await worker.start()
+# publisher = GrpcWorkerAgentRuntime(host_address=host_address)
+# publisher.add_message_serializer(try_get_known_serializers_for_type(MessageType))
+# await publisher.start()
+
+# await agent1.register_instance(worker, agent_id=agent_id)
+# await worker.add_subscription(TypeSubscription("notdefault", "name"))
+
+# with pytest.raises(RuntimeError):
+# await worker.publish_message(MessageType(), TopicId("notdefault", "notdefault"))
+# await asyncio.sleep(2)
+
+# await worker.stop()
+# await host.stop()
+
if __name__ == "__main__":
os.environ["GRPC_VERBOSITY"] = "DEBUG"
os.environ["GRPC_TRACE"] = "all"
diff --git a/python/packages/autogen-ext/tests/tools/azure/conftest.py b/python/packages/autogen-ext/tests/tools/azure/conftest.py
index 4b4a974ff..6d3a8569f 100644
--- a/python/packages/autogen-ext/tests/tools/azure/conftest.py
+++ b/python/packages/autogen-ext/tests/tools/azure/conftest.py
@@ -1,7 +1,7 @@
"""Test fixtures for Azure AI Search tool tests."""
import warnings
-from typing import Any, Dict, Generator, List, Protocol, Type, TypeVar, Union
+from typing import Any, Dict, Iterator, List, Protocol, TypeVar, Union
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -9,6 +9,17 @@ from autogen_core import ComponentModel
T = TypeVar("T")
+try:
+ from azure.core.credentials import AzureKeyCredential, TokenCredential
+
+ azure_sdk_available = True
+except ImportError:
+ azure_sdk_available = False
+
+skip_if_no_azure_sdk = pytest.mark.skipif(
+ not azure_sdk_available, reason="Azure SDK components (azure-search-documents, azure-identity) not available"
+)
+
class AccessTokenProtocol(Protocol):
"""Protocol matching Azure AccessToken."""
@@ -47,18 +58,13 @@ class MockTokenCredential:
return MockAccessToken("mock-token", 12345)
-try:
- from azure.core.credentials import AccessToken, AzureKeyCredential, TokenCredential
-
- _access_token_type: Type[AccessToken] = AccessToken
- azure_sdk_available = True
-except ImportError:
- AzureKeyCredential = MockAzureKeyCredential # type: ignore
- TokenCredential = MockTokenCredential # type: ignore
- _access_token_type = MockAccessToken # type: ignore
- azure_sdk_available = False
-
-CredentialType = Union[AzureKeyCredential, TokenCredential, MockAzureKeyCredential, MockTokenCredential, Any]
+CredentialType = Union[
+ AzureKeyCredential, # pyright: ignore [reportPossiblyUnboundVariable]
+ TokenCredential, # pyright: ignore [reportPossiblyUnboundVariable]
+ MockAzureKeyCredential,
+ MockTokenCredential,
+ Any,
+]
needs_azure_sdk = pytest.mark.skipif(not azure_sdk_available, reason="Azure SDK not available")
@@ -70,10 +76,14 @@ warnings.filterwarnings(
@pytest.fixture
-def mock_vectorized_query() -> Generator[MagicMock, None, None]:
+def mock_vectorized_query() -> MagicMock:
"""Create a mock VectorizedQuery for testing."""
- with patch("azure.search.documents.models.VectorizedQuery") as mock:
- yield mock
+ if azure_sdk_available:
+ from azure.search.documents.models import VectorizedQuery
+
+ return MagicMock(spec=VectorizedQuery)
+ else:
+ return MagicMock()
@pytest.fixture
@@ -87,7 +97,7 @@ def test_config() -> ComponentModel:
"endpoint": "https://test-search-service.search.windows.net",
"index_name": "test-index",
"api_version": "2023-10-01-Preview",
- "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},
+ "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"}, # pyright: ignore [reportPossiblyUnboundVariable]
"query_type": "keyword",
"search_fields": ["content", "title"],
"select_fields": ["id", "content", "title", "source"],
@@ -106,7 +116,7 @@ def keyword_config() -> ComponentModel:
"description": "Keyword search tool",
"endpoint": "https://test-search-service.search.windows.net",
"index_name": "test-index",
- "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},
+ "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"}, # pyright: ignore [reportPossiblyUnboundVariable]
"query_type": "keyword",
"search_fields": ["content", "title"],
"select_fields": ["id", "content", "title", "source"],
@@ -125,7 +135,7 @@ def vector_config() -> ComponentModel:
"endpoint": "https://test-search-service.search.windows.net",
"index_name": "test-index",
"api_version": "2023-10-01-Preview",
- "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},
+ "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"}, # pyright: ignore [reportPossiblyUnboundVariable]
"query_type": "vector",
"vector_fields": ["embedding"],
"select_fields": ["id", "content", "title", "source"],
@@ -145,7 +155,7 @@ def hybrid_config() -> ComponentModel:
"endpoint": "https://test-search-service.search.windows.net",
"index_name": "test-index",
"api_version": "2023-10-01-Preview",
- "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"},
+ "credential": AzureKeyCredential("test-key") if azure_sdk_available else {"api_key": "test-key"}, # pyright: ignore [reportPossiblyUnboundVariable]
"query_type": "keyword",
"search_fields": ["content", "title"],
"vector_fields": ["embedding"],
@@ -196,108 +206,18 @@ class AsyncIterator:
@pytest.fixture
-def mock_search_client(mock_search_response: List[Dict[str, Any]]) -> tuple[MagicMock, Any]:
- """Create a mock search client for testing."""
- mock_client = MagicMock()
- mock_client.__aenter__ = AsyncMock(return_value=mock_client)
- mock_client.__aexit__ = AsyncMock(return_value=None)
-
- search_results = AsyncIterator(mock_search_response)
- mock_client.search = MagicMock(return_value=search_results)
-
- patcher = patch("azure.search.documents.aio.SearchClient", return_value=mock_client)
-
- return mock_client, patcher
-
-
-def test_validate_credentials_scenarios() -> None:
- """Test all validate_credentials scenarios to ensure full code coverage."""
- import sys
-
- from autogen_ext.tools.azure._config import AzureAISearchConfig
-
- module_path = sys.modules[AzureAISearchConfig.__module__].__file__
- if module_path is not None:
- assert "autogen-ext" in module_path
-
- data: Any = "not a dict"
- result: Any = AzureAISearchConfig.validate_credentials(data) # type: ignore
- assert result == data
-
- data_empty: Dict[str, Any] = {}
- result_empty: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_empty) # type: ignore
- assert isinstance(result_empty, dict)
-
- data_items: Dict[str, Any] = {"key1": "value1", "key2": "value2"}
- result_items: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_items) # type: ignore
- assert result_items["key1"] == "value1"
- assert result_items["key2"] == "value2"
-
- data_with_api_key: Dict[str, Any] = {
- "name": "test",
- "endpoint": "https://test.search.windows.net",
- "index_name": "test-index",
- "credential": {"api_key": "test-key"},
- }
- result_with_api_key: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_with_api_key) # type: ignore
-
- cred = result_with_api_key["credential"] # type: ignore
- assert isinstance(cred, (AzureKeyCredential, MockAzureKeyCredential))
- assert hasattr(cred, "key")
- assert cred.key == "test-key" # type: ignore
-
- credential: Any = AzureKeyCredential("test-key")
- data_with_credential: Dict[str, Any] = {
- "name": "test",
- "endpoint": "https://test.search.windows.net",
- "index_name": "test-index",
- "credential": credential,
- }
- result_with_credential: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_with_credential) # type: ignore
- assert result_with_credential["credential"] is credential
-
- data_without_api_key: Dict[str, Any] = {
- "name": "test",
- "endpoint": "https://test.search.windows.net",
- "index_name": "test-index",
- "credential": {"username": "test-user", "password": "test-pass"},
- }
- result_without_api_key: Dict[str, Any] = AzureAISearchConfig.validate_credentials(data_without_api_key) # type: ignore
- assert result_without_api_key["credential"] == {"username": "test-user", "password": "test-pass"}
-
-
-def test_model_dump_scenarios() -> None:
- """Test all model_dump scenarios to ensure full code coverage."""
- import sys
-
- from autogen_ext.tools.azure._config import AzureAISearchConfig
-
- module_path = sys.modules[AzureAISearchConfig.__module__].__file__
- if module_path is not None:
- assert "autogen-ext" in module_path
-
- config = AzureAISearchConfig(
- name="test",
- endpoint="https://endpoint",
- index_name="index",
- credential=AzureKeyCredential("key"), # type: ignore
- )
- result = config.model_dump()
- assert result["credential"] == {"type": "AzureKeyCredential"}
+def mock_search_client(mock_search_response: List[Dict[str, Any]]) -> Iterator[MagicMock]:
+ """Create a mock search client for testing, with the patch active."""
+ mock_client_instance = MagicMock()
+ mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance)
+ mock_client_instance.__aexit__ = AsyncMock(return_value=None)
- if azure_sdk_available:
- from azure.core.credentials import AccessToken
- from azure.core.credentials import TokenCredential as RealTokenCredential
-
- class TestTokenCredential(RealTokenCredential):
- def get_token(self, *args: Any, **kwargs: Any) -> AccessToken:
- """Override of get_token method that returns proper type."""
- return AccessToken("test-token", 12345)
-
- config = AzureAISearchConfig(
- name="test", endpoint="https://endpoint", index_name="index", credential=TestTokenCredential()
- )
- result = config.model_dump()
- assert result["credential"] == {"type": "TokenCredential"}
- else:
- pytest.skip("Skipping TokenCredential test - Azure SDK not available")
+ search_results_iterator = AsyncIterator(mock_search_response)
+ mock_client_instance.search = MagicMock(return_value=search_results_iterator)
+
+ patch_target = "autogen_ext.tools.azure._ai_search.SearchClient"
+ patcher = patch(patch_target, return_value=mock_client_instance)
+
+ patcher.start()
+ yield mock_client_instance
+ patcher.stop()
diff --git a/python/packages/autogen-ext/tests/tools/azure/test_ai_search_config.py b/python/packages/autogen-ext/tests/tools/azure/test_ai_search_config.py
new file mode 100644
index 000000000..ddcc26e36
--- /dev/null
+++ b/python/packages/autogen-ext/tests/tools/azure/test_ai_search_config.py
@@ -0,0 +1,297 @@
+from typing import Any, Dict, cast
+
+import pytest
+from autogen_ext.tools.azure._config import AzureAISearchConfig, QueryTypeLiteral
+from azure.core.credentials import AzureKeyCredential
+from pydantic import ValidationError
+
+from tests.tools.azure.conftest import azure_sdk_available
+
+skip_if_no_azure_sdk = pytest.mark.skipif(
+ not azure_sdk_available, reason="Azure SDK components (azure-search-documents, azure-identity) not available"
+)
+
+# =====================================
+# Basic Configuration Tests
+# =====================================
+
+
+def test_basic_config_creation() -> None:
+ """Test that a basic valid configuration can be created."""
+ config = AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test-search.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ )
+
+ assert config.name == "test_tool"
+ assert config.endpoint == "https://test-search.search.windows.net"
+ assert config.index_name == "test-index"
+ assert isinstance(config.credential, AzureKeyCredential)
+ assert config.query_type == "simple" # default value
+
+
+def test_endpoint_validation() -> None:
+ """Test that endpoint validation works correctly."""
+ valid_endpoints = ["https://test.search.windows.net", "http://localhost:8080"]
+
+ for endpoint in valid_endpoints:
+ config = AzureAISearchConfig(
+ name="test_tool",
+ endpoint=endpoint,
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ )
+ assert config.endpoint == endpoint
+
+ invalid_endpoints = [
+ "test.search.windows.net",
+ "ftp://test.search.windows.net",
+ "",
+ ]
+
+ for endpoint in invalid_endpoints:
+ with pytest.raises(ValidationError) as exc:
+ AzureAISearchConfig(
+ name="test_tool",
+ endpoint=endpoint,
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ )
+ assert "endpoint must be a valid URL" in str(exc.value)
+
+
+def test_top_validation() -> None:
+ """Test validation of top parameter."""
+ valid_tops = [1, 5, 10, 100]
+
+ for top in valid_tops:
+ config = AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ top=top,
+ )
+ assert config.top == top
+
+ invalid_tops = [0, -1, -10]
+
+ for top in invalid_tops:
+ with pytest.raises(ValidationError) as exc:
+ AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ top=top,
+ )
+ assert "top must be a positive integer" in str(exc.value)
+
+
+# =====================================
+# Query Type Tests
+# =====================================
+
+
+def test_query_type_normalization() -> None:
+ """Test that query_type normalization works correctly."""
+ standard_query_types = {
+ "simple": "simple",
+ "full": "full",
+ "semantic": "semantic",
+ "vector": "vector",
+ }
+
+ for input_type, expected_type in standard_query_types.items():
+ config_args: Dict[str, Any] = {
+ "name": "test_tool",
+ "endpoint": "https://test.search.windows.net",
+ "index_name": "test-index",
+ "credential": AzureKeyCredential("test-key"),
+ "query_type": cast(QueryTypeLiteral, input_type),
+ }
+
+ if input_type == "semantic":
+ config_args["semantic_config_name"] = "my-semantic-config"
+ elif input_type == "vector":
+ config_args["vector_fields"] = ["content_vector"]
+
+ config = AzureAISearchConfig(**config_args)
+ assert config.query_type == expected_type
+
+ with pytest.raises(ValidationError) as exc:
+ AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ query_type=cast(Any, "invalid_type"),
+ )
+ assert "Input should be" in str(exc.value)
+
+
+def test_semantic_config_validation() -> None:
+ """Test validation of semantic configuration."""
+ config = AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ query_type=cast(QueryTypeLiteral, "semantic"),
+ semantic_config_name="my-semantic-config",
+ )
+ assert config.query_type == "semantic"
+ assert config.semantic_config_name == "my-semantic-config"
+
+ with pytest.raises(ValidationError) as exc:
+ AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ query_type=cast(QueryTypeLiteral, "semantic"),
+ )
+ assert "semantic_config_name must be provided" in str(exc.value)
+
+
+def test_vector_fields_validation() -> None:
+ """Test validation of vector fields for vector search."""
+ config = AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ query_type=cast(QueryTypeLiteral, "vector"),
+ vector_fields=["content_vector"],
+ )
+ assert config.query_type == "vector"
+ assert config.vector_fields == ["content_vector"]
+
+
+# =====================================
+# Embedding Configuration Tests
+# =====================================
+
+
+def test_azure_openai_endpoint_validation() -> None:
+ """Test validation of Azure OpenAI endpoint for client-side embeddings."""
+ config = AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ openai_endpoint="https://test.openai.azure.com",
+ )
+ assert config.embedding_provider == "azure_openai"
+ assert config.embedding_model == "text-embedding-ada-002"
+ assert config.openai_endpoint == "https://test.openai.azure.com"
+
+ with pytest.raises(ValidationError) as exc:
+ AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ )
+ assert "openai_endpoint must be provided for azure_openai" in str(exc.value)
+
+ config = AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ embedding_provider="openai",
+ embedding_model="text-embedding-ada-002",
+ )
+ assert config.embedding_provider == "openai"
+ assert config.embedding_model == "text-embedding-ada-002"
+ assert config.openai_endpoint is None
+
+
+# =====================================
+# Credential and Serialization Tests
+# =====================================
+
+
+def test_credential_validation() -> None:
+ """Test credential validation scenarios."""
+ config = AzureAISearchConfig(
+ name="test_tool",
+ endpoint="https://test.search.windows.net",
+ index_name="test-index",
+ credential=AzureKeyCredential("test-key"),
+ )
+ assert isinstance(config.credential, AzureKeyCredential)
+ assert config.credential.key == "test-key"
+
+ if azure_sdk_available:
+ from azure.core.credentials import AccessToken
+ from azure.core.credentials_async import AsyncTokenCredential
+
+ class TestTokenCredential(AsyncTokenCredential):
+ async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+ return AccessToken("test-token", 12345)
+
+ async def close(self) -> None:
+ pass
+
+ async def __aenter__(self) -> "TestTokenCredential":
+ return self
+
+ async def __aexit__(self, *args: Any) -> None:
+ await self.close()
+
+ config = AzureAISearchConfig(
+ name="test",
+ endpoint="https://endpoint",
+ index_name="index",
+ credential=TestTokenCredential(),
+ )
+ assert isinstance(config.credential, AsyncTokenCredential)
+
+
+def test_model_dump_scenarios() -> None:
+ """Test all model_dump scenarios to ensure full code coverage."""
+ config = AzureAISearchConfig(
+ name="test",
+ endpoint="https://endpoint",
+ index_name="index",
+ credential=AzureKeyCredential("key"),
+ )
+ result = config.model_dump()
+ assert isinstance(result["credential"], AzureKeyCredential)
+ assert result["credential"].key == "key"
+
+ if azure_sdk_available:
+ from azure.core.credentials import AccessToken
+ from azure.core.credentials_async import AsyncTokenCredential
+
+ class TestTokenCredential(AsyncTokenCredential):
+ async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+ return AccessToken("test-token", 12345)
+
+ async def close(self) -> None:
+ pass
+
+ async def __aenter__(self) -> "TestTokenCredential":
+ return self
+
+ async def __aexit__(self, *args: Any) -> None:
+ await self.close()
+
+ config = AzureAISearchConfig(
+ name="test",
+ endpoint="https://endpoint",
+ index_name="index",
+ credential=TestTokenCredential(),
+ )
+ result = config.model_dump()
+ assert isinstance(result["credential"], AsyncTokenCredential)
+ else:
+ pytest.skip("Skipping TokenCredential test - Azure SDK not available")
diff --git a/python/packages/autogen-ext/tests/tools/azure/test_ai_search_tool.py b/python/packages/autogen-ext/tests/tools/azure/test_ai_search_tool.py
index 13b70e823..67b1d357f 100644
--- a/python/packages/autogen-ext/tests/tools/azure/test_ai_search_tool.py
+++ b/python/packages/autogen-ext/tests/tools/azure/test_ai_search_tool.py
@@ -1,1374 +1,971 @@
-"""Tests for the Azure AI Search tool."""
+"""Tests for Azure AI Search tool."""
-# pyright: reportPrivateUsage=false
-
-from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Union, cast
+import asyncio
+from collections.abc import Generator
+from typing import Any, Dict, List
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from autogen_core import CancellationToken
-from autogen_ext.tools.azure._ai_search import (
+from autogen_ext.tools.azure import (
+ AzureAISearchConfig,
AzureAISearchTool,
- BaseAzureAISearchTool,
- SearchQuery,
SearchResult,
SearchResults,
- _allow_private_constructor,
)
-from azure.core.credentials import AzureKeyCredential, TokenCredential
-from azure.core.exceptions import HttpResponseError
+from autogen_ext.tools.azure._ai_search import BaseAzureAISearchTool
+from autogen_ext.tools.azure._config import DEFAULT_API_VERSION
+from azure.core.credentials import AzureKeyCredential
+from azure.core.credentials_async import AsyncTokenCredential
+from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
+from pydantic import BaseModel, Field, ValidationError
+
+MOCK_ENDPOINT = "https://test-search.search.windows.net"
+MOCK_INDEX = "test-index"
+MOCK_API_KEY = "test-key"
+MOCK_CREDENTIAL = AzureKeyCredential(MOCK_API_KEY)
-class MockAsyncIterator:
- """Mock for async iterator to use in tests."""
+class MockAsyncTokenCredential(AsyncTokenCredential):
+ """Mock async token credential for testing."""
- def __init__(self, items: List[Dict[str, Any]]) -> None:
- self.items = items.copy()
+ async def get_token(self, *scopes: str, **kwargs: Any) -> Any:
+ return "mock-token"
- def __aiter__(self) -> "MockAsyncIterator":
- return self
+ async def close(self) -> None:
+ pass
- async def __anext__(self) -> Dict[str, Any]:
- if not self.items:
- raise StopAsyncIteration
- return self.items.pop(0)
+ async def __aexit__(self, exc_type: Any = None, exc_val: Any = None, exc_tb: Any = None) -> None:
+ await self.close()
@pytest.fixture
-async def search_tool() -> AsyncGenerator[AzureAISearchTool, None]:
- """Create a concrete search tool for testing."""
+def search_config() -> AzureAISearchConfig:
+ """Fixture for basic search configuration."""
+ return AzureAISearchConfig(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ description="Test search tool",
+ )
- class ConcreteSearchTool(AzureAISearchTool):
- async def _get_embedding(self, query: str) -> List[float]:
- return [0.1, 0.2, 0.3]
- token = _allow_private_constructor.set(True)
- try:
- tool = ConcreteSearchTool(
- name="test-search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=cast(TokenCredential, AzureKeyCredential("test-key")),
- query_type="keyword",
- search_fields=["title", "content"],
- select_fields=["title", "content"],
- top=10,
- )
- yield tool
- finally:
- _allow_private_constructor.reset(token)
+@pytest.fixture
+def mock_search_client() -> Generator[AsyncMock, None, None]:
+ """Fixture for mocked search client."""
+ with patch("azure.search.documents.aio.SearchClient", autospec=True) as mock:
+ mock_client = AsyncMock()
+ mock.return_value = mock_client
+ yield mock_client
-@pytest.mark.asyncio
-async def test_search_tool_run(search_tool: AsyncGenerator[AzureAISearchTool, None]) -> None:
- """Test the run method of the search tool."""
- tool = await anext(search_tool)
- query = "test query"
- cancellation_token = CancellationToken()
-
- with patch.object(tool, "_get_client", AsyncMock()) as mock_client:
- mock_client.return_value.search = AsyncMock(
- return_value=MockAsyncIterator([{"@search.score": 0.95, "title": "Test Doc", "content": "Test Content"}])
- )
+@pytest.fixture
+def mock_search_results() -> List[Dict[str, Any]]:
+ """Fixture for mock search results."""
+ return [
+ {
+ "id": "1",
+ "content": "Test content",
+ "@search.score": 0.8,
+ }
+ ]
- results = await tool.run(query, cancellation_token)
- assert isinstance(results, SearchResults)
- assert len(results.results) == 1
- assert results.results[0].content["title"] == "Test Doc"
- assert results.results[0].score == 0.95
+class TestSearchQuery(BaseModel):
+ """Test model for query validation."""
-@pytest.mark.asyncio
-async def test_search_tool_error_handling(search_tool: AsyncGenerator[AzureAISearchTool, None]) -> None:
- """Test error handling in the search tool."""
- tool = await anext(search_tool)
- with patch.object(tool, "_get_client", AsyncMock()) as mock_client:
- mock_client.return_value.search = AsyncMock(side_effect=ValueError("Test error"))
-
- with pytest.raises(ValueError, match="Test error"):
- await tool.run("test query", CancellationToken())
+ query: str = Field(min_length=1)
@pytest.mark.asyncio
-async def test_search_tool_cancellation(search_tool: AsyncGenerator[AzureAISearchTool, None]) -> None:
- """Test cancellation of the search tool."""
- tool = await anext(search_tool)
- cancellation_token = CancellationToken()
- cancellation_token.cancel()
+async def test_search_query_model() -> None:
+ """Test SearchQuery model validation."""
+ query = TestSearchQuery(query="test query")
+ assert query.query == "test query"
- with pytest.raises(ValueError, match="cancelled"):
- await tool.run("test query", cancellation_token)
+ with pytest.raises(ValidationError):
+ TestSearchQuery(query="")
@pytest.mark.asyncio
-async def test_search_tool_vector_search() -> None:
- """Test vector search functionality."""
-
- class ConcreteSearchTool(AzureAISearchTool):
- async def _get_embedding(self, query: str) -> List[float]:
- return [0.1, 0.2, 0.3]
-
- token = _allow_private_constructor.set(True)
- try:
- tool = ConcreteSearchTool(
- name="vector-search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=cast(TokenCredential, AzureKeyCredential("test-key")),
- query_type="vector",
- vector_fields=["embedding"],
- select_fields=["title", "content"],
- top=10,
- )
-
- with patch.object(tool, "_get_client", AsyncMock()) as mock_client:
- mock_client.return_value.search = AsyncMock(
- return_value=MockAsyncIterator(
- [{"@search.score": 0.95, "title": "Vector Doc", "content": "Vector Content"}]
- )
- )
-
- results = await tool.run("vector query", CancellationToken())
- assert len(results.results) == 1
- assert results.results[0].content["title"] == "Vector Doc"
- assert results.results[0].score == 0.95
- finally:
- _allow_private_constructor.reset(token)
-
-
-class ConcreteAzureAISearchTool(AzureAISearchTool):
- """Concrete implementation for testing."""
-
- async def _get_embedding(self, query: str) -> List[float]:
- return [0.1, 0.2, 0.3]
+async def test_search_result_model() -> None:
+ """Test SearchResult model."""
+ result = SearchResult(score=0.8, content={"title": "Test", "text": "Content"}, metadata={"@search.score": 0.8})
+ assert result.score == 0.8
+ assert result.content["title"] == "Test"
+ assert result.metadata["@search.score"] == 0.8
@pytest.mark.asyncio
-async def test_create_keyword_search() -> None:
- """Test the create_keyword_search factory method."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="keyword_search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=cast(TokenCredential, AzureKeyCredential("test-key")),
- search_fields=["title", "content"],
- select_fields=["title", "content"],
- filter="category eq 'test'",
- top=5,
+async def test_search_results_model() -> None:
+ """Test SearchResults model."""
+ results = SearchResults(
+ results=[
+ SearchResult(score=0.8, content={"title": "Test1"}, metadata={"@search.score": 0.8}),
+ SearchResult(score=0.6, content={"title": "Test2"}, metadata={"@search.score": 0.6}),
+ ]
)
-
- assert tool.name == "keyword_search"
- assert tool.search_config.query_type == "keyword"
- assert tool.search_config.filter == "category eq 'test'"
+ assert len(results.results) == 2
+ assert results.results[0].score == 0.8
+ assert results.results[1].content["title"] == "Test2"
@pytest.mark.asyncio
async def test_create_full_text_search() -> None:
- """Test the create_full_text_search factory method."""
- tool = ConcreteAzureAISearchTool.create_full_text_search(
- name="full_text_search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=cast(TokenCredential, AzureKeyCredential("test-key")),
- search_fields=["title", "content"],
- select_fields=["title", "content"],
- filter="category eq 'test'",
- top=5,
- )
+ """Test creation of full text search tool."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ search_fields=["content"],
+ query_type="simple",
+ )
+ assert tool.name == "test-search"
+ assert tool.search_config.query_type == "simple"
+ assert tool.search_config.search_fields == ["content"]
+
+ with pytest.raises(ValueError, match="semantic_config_name is required"):
+ AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ query_type="semantic",
+ )
- assert tool.name == "full_text_search"
- assert tool.search_config.query_type == "fulltext"
- assert tool.search_config.search_fields == ["title", "content"]
- assert tool.search_config.select_fields == ["title", "content"]
- assert tool.search_config.filter == "category eq 'test'"
- assert tool.search_config.top == 5
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ query_type="semantic",
+ semantic_config_name="default",
+ )
+ assert tool.search_config.query_type == "semantic"
+ assert tool.search_config.semantic_config_name == "default"
@pytest.mark.asyncio
async def test_create_vector_search() -> None:
- """Test the create_vector_search factory method."""
- tool = ConcreteAzureAISearchTool.create_vector_search(
- name="vector_search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+ """Test creation of vector search tool."""
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
vector_fields=["embedding"],
- select_fields=["title", "content"],
- top=5,
)
-
- assert tool.name == "vector_search"
assert tool.search_config.query_type == "vector"
assert tool.search_config.vector_fields == ["embedding"]
+ with pytest.raises(ValueError, match="openai_endpoint is required"):
+ AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ )
+
@pytest.mark.asyncio
async def test_create_hybrid_search() -> None:
- """Test the create_hybrid_search factory method (hybrid = text + vector, query_type will be 'fulltext' or 'semantic')."""
- tool = ConcreteAzureAISearchTool.create_hybrid_search(
- name="hybrid_search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+ """Test creation of hybrid search tool."""
+ tool = AzureAISearchTool.create_hybrid_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
vector_fields=["embedding"],
- search_fields=["title", "content"],
- select_fields=["title", "content"],
- top=5,
+ search_fields=["content"],
)
-
- assert tool.name == "hybrid_search"
- assert tool.search_config.query_type in ("fulltext", "semantic")
assert tool.search_config.vector_fields == ["embedding"]
- assert tool.search_config.search_fields == ["title", "content"]
-
+ assert tool.search_config.search_fields == ["content"]
-@pytest.mark.asyncio
-async def test_run_invalid_query() -> None:
- """Test the run method with an invalid query format."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+ tool = AzureAISearchTool.create_hybrid_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ search_fields=["content"],
+ query_type="semantic",
+ semantic_config_name="default",
)
-
- invalid_query: Dict[str, Any] = {"invalid_key": "invalid_value"}
- with pytest.raises(ValueError, match="Invalid search query format"):
- await tool.run(invalid_query)
+ assert tool.search_config.query_type == "semantic"
+ assert tool.search_config.semantic_config_name == "default"
@pytest.mark.asyncio
-async def test_process_credential_dict() -> None:
- """Test the _process_credential method with a dictionary credential."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential={"api_key": "test-key"},
- )
+async def test_search_execution(mock_search_client: AsyncMock, mock_search_results: List[Dict[str, Any]]) -> None:
+ """Test search execution with mocked client."""
+ mock_search_client.search.return_value.__aiter__.return_value = mock_search_results
- assert isinstance(tool.search_config.credential, AzureKeyCredential)
- assert tool.search_config.credential.key == "test-key"
-
-
-@pytest.mark.asyncio
-async def test_run_empty_query() -> None:
- """Test the run method with an empty query."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- with patch.object(tool, "_get_client", AsyncMock()):
- with pytest.raises(ValueError, match="Invalid search query format"):
- await tool.run("")
-
-
-@pytest.mark.asyncio
-async def test_get_client_initialization() -> None:
- """Test the _get_client method for proper initialization."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- )
-
- assert tool.search_config.endpoint == "https://test.search.windows.net"
- assert tool.search_config.index_name == "test-index"
-
- mock_client = AsyncMock()
-
- class MockAsyncIterator:
- def __init__(self, items: List[Dict[str, Any]]) -> None:
- self.items = items
-
- def __aiter__(self) -> "MockAsyncIterator":
- return self
-
- async def __anext__(self) -> Dict[str, Any]:
- if not self.items:
- raise StopAsyncIteration
- return self.items.pop(0)
-
- mock_client.search.return_value = MockAsyncIterator([{"@search.score": 0.9, "title": "Test Result"}])
-
- with patch.object(tool, "_get_client", return_value=mock_client):
- results = await tool.run("test query", CancellationToken())
- mock_client.search.assert_called_once()
+ with patch.object(tool, "_get_client", return_value=mock_search_client):
+ results = await tool.run("test query")
assert len(results.results) == 1
- assert results.results[0].content["title"] == "Test Result"
- assert results.results[0].score == 0.9
-
-
-@pytest.mark.asyncio
-async def test_return_value_as_string() -> None:
- """Test the return_value_as_string method for formatting search results."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- )
-
- results = SearchResults(
- results=[
- SearchResult(score=0.95, content={"title": "Doc 1"}, metadata={}),
- SearchResult(score=0.85, content={"title": "Doc 2"}, metadata={}),
- ]
- )
-
- result_string = tool.return_value_as_string(results)
- assert "Result 1 (Score: 0.95): title: Doc 1" in result_string
- assert "Result 2 (Score: 0.85): title: Doc 2" in result_string
-
-
-@pytest.mark.asyncio
-async def test_return_value_as_string_empty() -> None:
- """Test the return_value_as_string method with empty results."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- )
+ assert results.results[0].score == 0.8
+ assert results.results[0].content["content"] == "Test content"
- results = SearchResults(results=[])
- result_string = tool.return_value_as_string(results)
- assert result_string == "No results found."
+ mock_search_client.search.assert_called_once()
+ call_kwargs = mock_search_client.search.call_args[1]
+ assert call_kwargs["search_text"] == "test query"
@pytest.mark.asyncio
-async def test_load_component() -> None:
- """Test the load_component method for proper deserialization."""
- model = {
- "provider": "autogen_ext.tools.azure.BaseAzureAISearchTool",
- "config": {
- "name": "test_tool",
- "endpoint": "https://test.search.windows.net",
- "index_name": "test-index",
- "credential": {"api_key": "test-key"},
- "query_type": "keyword",
- "search_fields": ["title", "content"],
- "select_fields": ["title", "content"],
- "top": 5,
- },
- }
-
- tool = ConcreteAzureAISearchTool.load_component(model)
- assert tool.name == "test_tool"
- assert tool.search_config.query_type == "keyword"
- assert tool.search_config.search_fields == ["title", "content"]
-
+async def test_search_with_caching(mock_search_client: AsyncMock, mock_search_results: List[Dict[str, Any]]) -> None:
+ """Test search caching functionality."""
+ mock_search_client.search.return_value.__aiter__.return_value = mock_search_results
-@pytest.mark.asyncio
-async def test_caching_functionality() -> None:
- """Test the caching functionality of the search tool."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="cache_test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
enable_caching=True,
cache_ttl_seconds=300,
)
- mock_client = AsyncMock()
- mock_client.__aenter__.return_value = mock_client
- mock_client.__aexit__.return_value = None
-
- test_result = {"@search.score": 0.9, "title": "Test Result"}
+ with patch.object(tool, "_get_client", return_value=mock_search_client):
+ await tool.run("test query")
+ assert mock_search_client.search.call_count == 1
- class MockAsyncIterator:
- def __init__(self) -> None:
- self.returned = False
-
- def __aiter__(self) -> "MockAsyncIterator":
- return self
-
- async def __anext__(self) -> Dict[str, Any]:
- if self.returned:
- raise StopAsyncIteration
- self.returned = True
- return test_result
-
- mock_client.search = AsyncMock(return_value=MockAsyncIterator())
-
- with patch.object(tool, "_get_client", return_value=mock_client):
- results1 = await tool.run("test query")
- assert len(results1.results) == 1
- assert results1.results[0].content["title"] == "Test Result"
- assert mock_client.search.call_count == 1
-
- mock_client.search = AsyncMock(return_value=MockAsyncIterator())
-
- results2 = await tool.run("test query")
- assert len(results2.results) == 1
- assert results2.results[0].content["title"] == "Test Result"
- assert mock_client.search.call_count == 1
+ await tool.run("test query")
+ assert mock_search_client.search.call_count == 1
@pytest.mark.asyncio
-async def test_semantic_configuration_name_handling() -> None:
- """Test handling of semantic configuration names in fulltext search."""
- tool = ConcreteAzureAISearchTool.create_full_text_search(
- name="semantic_config_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- search_fields=["title", "content"],
- select_fields=["title", "content"],
+async def test_error_handling(mock_search_client: AsyncMock) -> None:
+ """Test error handling in search execution."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- mock_client = AsyncMock()
- mock_client.search.return_value = MockAsyncIterator([{"@search.score": 0.9, "title": "Semantic Test Result"}])
-
- assert tool.search_config.query_type == "fulltext"
- assert tool.search_config.search_fields == ["title", "content"]
-
- with patch.object(tool, "_get_client", return_value=mock_client):
- mock_run = AsyncMock()
- mock_run.return_value = SearchResults(
- results=[SearchResult(score=0.9, content={"title": "Semantic Test Result"}, metadata={})]
- )
-
- with patch.object(tool, "run", mock_run):
- results = await tool.run("semantic query")
- mock_run.assert_called_once()
-
- assert len(results.results) == 1
- assert results.results[0].content["title"] == "Semantic Test Result"
-
-
-@pytest.mark.asyncio
-async def test_http_response_error_handling() -> None:
- """Test handling of different HTTP response errors."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- )
-
- mock_client = AsyncMock()
- http_error = HttpResponseError()
- http_error.message = "401 Unauthorized: Access is denied due to invalid credentials"
+ with patch.object(tool, "_get_client", return_value=mock_search_client):
+ mock_search_client.search.side_effect = ResourceNotFoundError("Index not found")
+ with pytest.raises(ValueError, match="Index.*not found"):
+ await tool.run("test query")
- with patch.object(tool, "_get_client", return_value=mock_client):
- mock_client.search = AsyncMock(side_effect=http_error)
+ mock_search_client.search.side_effect = HttpResponseError(status_code=401, message="Unauthorized")
with pytest.raises(ValueError, match="Authentication failed"):
await tool.run("test query")
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("invalid-key"),
- )
-
- with patch.object(tool, "_get_client", AsyncMock(side_effect=ValueError("Invalid key"))):
- with pytest.raises(ValueError, match="Authentication failed"):
+ mock_search_client.search.side_effect = HttpResponseError(status_code=500, message="Internal server error")
+ with pytest.raises(ValueError, match="Error from Azure AI Search"):
await tool.run("test query")
@pytest.mark.asyncio
-async def test_run_with_search_query_object() -> None:
- """Test running the search with a SearchQuery object instead of a string."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- )
+async def test_embedding_provider_mixin() -> None:
+ """Test the embedding provider functionality."""
+ with patch("openai.AsyncAzureOpenAI") as mock_azure_openai:
+ mock_client = AsyncMock()
+ mock_azure_openai.return_value = mock_client
+ mock_client.embeddings.create.return_value.data = [MagicMock(embedding=[0.1, 0.2, 0.3])]
- mock_client = AsyncMock()
- mock_client.search.return_value = MockAsyncIterator([{"@search.score": 0.85, "title": "Query Object Test"}])
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ openai_endpoint="https://test.openai.azure.com",
+ openai_api_key="test-key",
+ )
- with patch.object(tool, "_get_client", return_value=mock_client):
- search_query = SearchQuery(query="advanced query")
- results = await tool.run(search_query)
+ embedding = await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage]
+ assert len(embedding) == 3
+ assert embedding == [0.1, 0.2, 0.3]
- assert len(results.results) == 1
- assert results.results[0].content["title"] == "Query Object Test"
- mock_client.search.assert_called_once()
+ mock_client.embeddings.create.assert_called_once_with(model="text-embedding-ada-002", input="test query")
@pytest.mark.asyncio
-async def test_dict_document_processing() -> None:
- """Test processing of document with dict-like interface."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+async def test_credential_processing() -> None:
+ """Test credential processing logic."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
+ assert isinstance(tool.search_config.credential, AzureKeyCredential)
+ assert tool.search_config.credential.key == MOCK_API_KEY
- class DictLikeDoc:
- def __init__(self, data: Dict[str, Any]) -> None:
- self._data = data
-
- def items(self) -> List[tuple[str, Any]]:
- return list(self._data.items())
-
- mock_client = AsyncMock()
-
- class SpecialMockAsyncIterator:
- def __init__(self) -> None:
- self.returned = False
-
- def __aiter__(self) -> "SpecialMockAsyncIterator":
- return self
-
- async def __anext__(self) -> DictLikeDoc:
- if self.returned:
- raise StopAsyncIteration
- self.returned = True
- return DictLikeDoc({"@search.score": 0.75, "title": "Dict Like Doc"})
-
- mock_client.search.return_value = SpecialMockAsyncIterator()
-
- with patch.object(tool, "_get_client", return_value=mock_client):
- results = await tool.run("test query")
-
- assert len(results.results) == 1
- assert results.results[0].content["title"] == "Dict Like Doc"
- assert results.results[0].score == 0.75
-
-
-@pytest.mark.asyncio
-async def test_document_processing_error_handling() -> None:
- """Test error handling during document processing."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+ mock_async_credential = MockAsyncTokenCredential()
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=mock_async_credential,
)
+ assert isinstance(tool.search_config.credential, AsyncTokenCredential)
- mock_client = AsyncMock()
-
- class ProblemDoc:
- def items(self) -> None:
- raise AttributeError("Simulated error in document processing")
-
- class MixedResultsAsyncIterator:
- def __init__(self) -> None:
- self.docs: List[Union[Dict[str, Any], ProblemDoc]] = [
- {"@search.score": 0.9, "title": "Good Doc"},
- ProblemDoc(),
- {"@search.score": 0.8, "title": "Another Good Doc"},
- ]
- self.index = 0
-
- def __aiter__(self) -> "MixedResultsAsyncIterator":
- return self
-
- async def __anext__(self) -> Union[Dict[str, Any], ProblemDoc]:
- if self.index >= len(self.docs):
- raise StopAsyncIteration
- doc = self.docs[self.index]
- self.index += 1
- return doc
-
- mock_client.search.return_value = MixedResultsAsyncIterator()
-
- with patch.object(tool, "_get_client", return_value=mock_client):
- results = await tool.run("test query")
-
- assert len(results.results) == 2
- assert results.results[0].content["title"] == "Good Doc"
- assert results.results[1].content["title"] == "Another Good Doc"
+ with pytest.raises(ValueError, match="Invalid configuration"):
+ AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential={"api_key": "test-key"}, # type: ignore
+ )
@pytest.mark.asyncio
-async def test_index_not_found_error() -> None:
- """Test handling of 'index not found' error."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="nonexistent-index",
- credential=AzureKeyCredential("test-key"),
+async def test_return_value_as_string() -> None:
+ """Test the string representation of search results."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- not_found_error = ValueError("The index 'nonexistent-index' was not found")
-
- with patch.object(tool, "_get_client", AsyncMock(side_effect=not_found_error)):
- with pytest.raises(ValueError, match="Index 'nonexistent-index' not found"):
- await tool.run("test query")
-
-
-@pytest.mark.asyncio
-async def test_http_response_with_500_error() -> None:
- """Test handling of HTTP 500 error responses."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+ results = SearchResults(
+ results=[
+ SearchResult(score=0.8, content={"title": "Test1", "text": "Content1"}, metadata={"@search.score": 0.8}),
+ SearchResult(score=0.6, content={"title": "Test2", "text": "Content2"}, metadata={"@search.score": 0.6}),
+ ]
)
+ result_str = tool.return_value_as_string(results)
+ assert "Result 1 (Score: 0.80)" in result_str
+ assert "Result 2 (Score: 0.60)" in result_str
+ assert "Test1" in result_str
+ assert "Content2" in result_str
- http_error = HttpResponseError()
- http_error.message = "500 Internal Server Error: Something went wrong on the server"
-
- with patch.object(tool, "_get_client", AsyncMock()) as mock_client:
- mock_client.return_value.search = AsyncMock(side_effect=http_error)
-
- with pytest.raises(ValueError, match="Error from Azure AI Search"):
- await tool.run("test query")
+ empty_results = SearchResults(results=[])
+ assert tool.return_value_as_string(empty_results) == "No results found."
@pytest.mark.asyncio
-async def test_cancellation_during_search() -> None:
- """Test cancellation token functionality during the search process."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+async def test_schema() -> None:
+ """Test tool schema generation."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- cancellation_token = CancellationToken()
- cancellation_token.cancel()
-
- with pytest.raises(ValueError, match="Operation cancelled"):
- await tool.run("test query", cancellation_token)
+ schema = tool.schema
+ assert schema["name"] == "test-search"
+ assert "description" in schema
+ assert "parameters" in schema
+ assert "required" in schema["parameters"]
+ assert schema["parameters"]["type"] == "object"
+ assert "query" in schema["parameters"]["properties"]
+ assert schema["parameters"]["required"] == ["query"]
@pytest.mark.asyncio
-async def test_run_with_dict_query_format() -> None:
- """Test running the search with a dictionary query format with 'query' key."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- )
-
- mock_client = AsyncMock()
- mock_client.search.return_value = MockAsyncIterator([{"@search.score": 0.85, "title": "Dict Query Test"}])
+async def test_vector_search_execution(
+ mock_search_client: AsyncMock, mock_search_results: List[Dict[str, Any]]
+) -> None:
+ """Test vector search execution."""
+ mock_search_client.search.return_value.__aiter__.return_value = mock_search_results
- with patch.object(tool, "_get_client", return_value=mock_client):
- query_dict = {"query": "dict style query"}
- results = await tool.run(query_dict)
-
- assert len(results.results) == 1
- assert results.results[0].content["title"] == "Dict Query Test"
- mock_client.search.assert_called_once()
-
-
-@pytest.mark.asyncio
-async def test_object_based_document_processing() -> None:
- """Test processing of document with object attributes instead of dict interface."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="test_tool",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ openai_endpoint="https://test.openai.azure.com",
+ openai_api_key="test-key",
)
- class ObjectDoc:
- """Test document class with object attributes."""
-
- def __init__(self) -> None:
- self.title = "Object Doc"
- self.content = "Object content"
- self._private_attr = "private"
- self.__search_score = 0.8
-
- mock_client = AsyncMock()
-
- class ObjectDocAsyncIterator:
- def __init__(self) -> None:
- self.returned = False
-
- def __aiter__(self) -> "ObjectDocAsyncIterator":
- return self
-
- async def __anext__(self) -> ObjectDoc:
- if self.returned:
- raise StopAsyncIteration
- self.returned = True
- return ObjectDoc()
-
- mock_client.search.return_value = ObjectDocAsyncIterator()
-
- with patch.object(tool, "_get_client", return_value=mock_client):
+ mock_embedding = [0.1, 0.2, 0.3]
+ with (
+ patch.object(tool, "_get_embedding", return_value=mock_embedding),
+ patch.object(tool, "_get_client", return_value=mock_search_client),
+ ):
results = await tool.run("test query")
-
assert len(results.results) == 1
- assert results.results[0].content["title"] == "Object Doc"
- assert results.results[0].content["content"] == "Object content"
- assert "_private_attr" not in results.results[0].content
+ mock_search_client.search.assert_called_once()
@pytest.mark.asyncio
-async def test_vector_search_with_provided_vectors() -> None:
- """Test vector search using vectors provided directly in the search options."""
- tool = ConcreteAzureAISearchTool.create_vector_search(
- name="vector_direct_search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- vector_fields=["embedding"],
- select_fields=["title", "content"],
+async def test_cancellation() -> None:
+ """Test search cancellation."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- mock_client = AsyncMock()
- mock_client.search.return_value = MockAsyncIterator([{"@search.score": 0.95, "title": "Vector Direct Test"}])
-
- query = "test vector search"
-
- with patch.object(tool, "_get_client", return_value=mock_client):
- results = await tool.run(query)
- assert len(results.results) == 1
- assert results.results[0].content["title"] == "Vector Direct Test"
-
- mock_client.search.assert_called_once()
+ token = CancellationToken()
+ token.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await tool.run("test query", cancellation_token=token)
@pytest.mark.asyncio
-async def test_credential_token_expiry_handling() -> None:
- """Test handling credential token expiry and error cases."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="token_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+async def test_invalid_query_format() -> None:
+ """Test invalid query format handling."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- auth_error = HttpResponseError()
- auth_error.message = "401 Unauthorized: Access token has expired or is not yet valid"
-
- with patch.object(tool, "_get_client", AsyncMock()) as mock_client:
- mock_client.return_value.search = AsyncMock(side_effect=auth_error)
-
- with pytest.raises(ValueError, match="Authentication failed"):
- await tool.run("test query")
-
- token_error = ValueError("401 Unauthorized: Token is invalid")
-
- with patch.object(tool, "_get_client", AsyncMock(side_effect=token_error)):
- with pytest.raises(ValueError, match="Authentication failed"):
- await tool.run("test query")
+ with pytest.raises(ValueError, match="Invalid search query format"):
+ await tool.run({"invalid": "format"})
@pytest.mark.asyncio
-async def test_search_with_user_provided_vectors() -> None:
- """Test the use of user-provided embedding vectors in SearchQuery."""
- tool = ConcreteAzureAISearchTool.create_vector_search(
- name="vector_test_with_embeddings",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- vector_fields=["embedding"],
+async def test_client_cleanup() -> None:
+ """Test client cleanup."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
mock_client = AsyncMock()
- mock_client.search.return_value = MockAsyncIterator([{"@search.score": 0.95, "title": "Vector Result"}])
-
- custom_vectors = [0.1, 0.2, 0.3, 0.4, 0.5]
- query_dict = {"query": "test query", "vectors": {"embedding": custom_vectors}}
-
- with patch.object(tool, "_get_client", return_value=mock_client):
- results = await tool.run(query_dict)
-
- assert len(results.results) == 1
- assert results.results[0].content["title"] == "Vector Result"
-
- mock_client.search.assert_called_once()
-
-
-@pytest.mark.asyncio
-async def test_component_loading_with_invalid_params() -> None:
- """Test loading components with invalid parameters."""
-
- class OtherClass:
- pass
-
- with pytest.raises(TypeError, match="Cannot create instance"):
- BaseAzureAISearchTool.load_component(
- {"provider": "autogen_ext.tools.azure.BaseAzureAISearchTool", "config": {}},
- expected=OtherClass, # type: ignore
- )
-
- with pytest.raises(Exception) as excinfo:
- ConcreteAzureAISearchTool.load_component("not a dict or ComponentModel") # type: ignore
- error_msg = str(excinfo.value).lower()
- assert any(text in error_msg for text in ["attribute", "type", "object", "dict", "str"])
+ tool._client = mock_client # pyright: ignore[reportPrivateUsage]
+ await tool.close()
- with pytest.raises(Exception) as excinfo:
- ConcreteAzureAISearchTool.load_component({})
- error_msg = str(excinfo.value).lower()
- assert any(text in error_msg for text in ["validation", "required", "missing", "field"])
+ mock_client.close.assert_called_once()
+ assert tool._client is None # pyright: ignore[reportPrivateUsage]
@pytest.mark.asyncio
-async def test_factory_method_validation() -> None:
- """Test validation in factory methods."""
- with pytest.raises(ValueError, match="endpoint must be a valid URL"):
- ConcreteAzureAISearchTool.create_keyword_search(
- name="test", endpoint="", index_name="test-index", credential=AzureKeyCredential("test-key")
+async def test_config_validation() -> None:
+ """Test configuration validation."""
+ with pytest.raises(ValueError, match="vector_fields must contain at least one field"):
+ AzureAISearchTool._validate_config( # pyright: ignore[reportPrivateUsage]
+ {
+ "name": "test-search",
+ "endpoint": MOCK_ENDPOINT,
+ "index_name": MOCK_INDEX,
+ "credential": MOCK_CREDENTIAL,
+ },
+ "vector",
)
- with pytest.raises(ValueError, match="endpoint must be a valid URL"):
- ConcreteAzureAISearchTool.create_keyword_search(
- name="test", endpoint="invalid-url", index_name="test-index", credential=AzureKeyCredential("test-key")
+ with pytest.raises(ValueError, match="vector_fields must contain at least one field"):
+ AzureAISearchTool._validate_config( # pyright: ignore[reportPrivateUsage]
+ {
+ "name": "test-search",
+ "endpoint": MOCK_ENDPOINT,
+ "index_name": MOCK_INDEX,
+ "credential": MOCK_CREDENTIAL,
+ "search_fields": ["content"],
+ },
+ "hybrid",
)
- with pytest.raises(ValueError, match="index_name cannot be empty"):
- ConcreteAzureAISearchTool.create_keyword_search(
- name="test",
- endpoint="https://test.search.windows.net",
- index_name="",
- credential=AzureKeyCredential("test-key"),
- )
-
- with pytest.raises(ValueError, match="name cannot be empty"):
- ConcreteAzureAISearchTool.create_keyword_search(
- name="",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- )
- with pytest.raises(ValueError, match="credential cannot be None"):
- ConcreteAzureAISearchTool.create_keyword_search(
- name="test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=None, # type: ignore
- )
+@pytest.mark.asyncio
+async def test_openai_embedding_provider() -> None:
+ """Test OpenAI embedding provider."""
+ with patch("openai.AsyncOpenAI") as mock_openai:
+ mock_client = AsyncMock()
+ mock_openai.return_value = mock_client
+ mock_client.embeddings.create.return_value.data = [MagicMock(embedding=[0.1, 0.2, 0.3])]
- with pytest.raises(ValueError, match="vector_fields must contain at least one field name"):
- ConcreteAzureAISearchTool.create_hybrid_search(
- name="test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- vector_fields=[],
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ embedding_provider="openai",
+ embedding_model="text-embedding-ada-002",
+ openai_api_key="test-key",
)
- with pytest.raises(ValueError, match="vector_fields must contain at least one field name"):
- ConcreteAzureAISearchTool.create_hybrid_search(
- name="test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- vector_fields=[],
- )
+ embedding = await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage]
+ assert len(embedding) == 3
+ assert embedding == [0.1, 0.2, 0.3]
@pytest.mark.asyncio
-async def test_direct_tool_initialization_error() -> None:
- """Test that directly initializing AzureAISearchTool raises an error."""
-
- class TestSearchTool(AzureAISearchTool):
- async def _get_embedding(self, query: str) -> List[float]:
- return [0.1, 0.2, 0.3]
-
- with pytest.raises(RuntimeError, match="Constructor is private"):
- TestSearchTool(
- name="test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- query_type="keyword",
+async def test_embedding_provider_error_handling() -> None:
+ """Test error handling in embedding providers."""
+ with pytest.raises(ValueError, match="openai_endpoint is required when embedding_provider is 'azure_openai'"):
+ AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ openai_api_version="2023-11-01",
)
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ openai_endpoint="https://test.openai.azure.com",
+ openai_api_version="2023-11-01",
+ )
+ with patch("azure.identity.DefaultAzureCredential") as mock_credential:
+ mock_credential.return_value.get_token.return_value = None
+ with pytest.raises(ValueError, match="Failed to acquire token using DefaultAzureCredential for Azure OpenAI"):
+ await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage]
+
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ embedding_provider="unsupported_provider",
+ embedding_model="test-model",
+ )
+ with pytest.raises(ValueError, match="Unsupported client-side embedding provider: unsupported_provider"):
+ await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage]
+
@pytest.mark.asyncio
-async def test_credential_dict_with_missing_api_key() -> None:
- """Test handling of credential dict without api_key."""
- with pytest.raises(ValueError, match="If credential is a dict, it must contain an 'api_key' key"):
- ConcreteAzureAISearchTool.create_keyword_search(
- name="test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential={"invalid_key": "value"},
+async def test_abstract_base_class() -> None:
+ """Test abstract base class behavior."""
+ with pytest.raises(NotImplementedError):
+ BaseAzureAISearchTool._from_config( # pyright: ignore[reportPrivateUsage]
+ AzureAISearchConfig(
+ name="test",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ )
)
@pytest.mark.asyncio
-async def test_complex_error_handling_scenarios() -> None:
- """Test more complex error handling scenarios."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="error_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+async def test_client_initialization_errors() -> None:
+ """Test client initialization error handling."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- permission_error = HttpResponseError()
- permission_error.message = "403 Forbidden: Access is denied"
-
- with patch.object(tool, "_get_client", AsyncMock(side_effect=permission_error)):
- with pytest.raises(ValueError, match="Error from Azure AI Search"):
- await tool.run("test query")
-
- unexpected_error = Exception("Unexpected error during initialization")
+ with patch("azure.search.documents.aio.SearchClient.__init__", side_effect=Exception("Connection error")):
+ with pytest.raises(ValueError, match="Unexpected error initializing search client: Connection error"):
+ await tool._get_client() # pyright: ignore[reportPrivateUsage]
- with patch.object(tool, "_get_client", AsyncMock(side_effect=unexpected_error)):
- with pytest.raises(ValueError, match="Error from Azure AI Search"):
- await tool.run("test query")
+ with patch(
+ "azure.search.documents.aio.SearchClient.__init__", side_effect=ResourceNotFoundError("Index not found")
+ ):
+ with pytest.raises(ValueError, match=f"Index '{MOCK_INDEX}' not found"):
+ await tool._get_client() # pyright: ignore[reportPrivateUsage]
@pytest.mark.asyncio
-async def test_multi_step_vector_search() -> None:
- """Test a multi-step vector search with query embeddings and explicit search options."""
- tool = ConcreteAzureAISearchTool.create_vector_search(
- name="vector_multi_step",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- vector_fields=["embedding"],
+async def test_client_initialization_with_error() -> None:
+ """Test client initialization with various errors."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- mock_client = AsyncMock()
- mock_client.search.return_value = MockAsyncIterator([{"@search.score": 0.98, "title": "Vector Embedding Test"}])
+ class MockResponse:
+ def __init__(self, status_code: int, reason: str):
+ self.status_code = status_code
+ self.reason = reason
+ self.request = object()
- embedding = [0.1, 0.2, 0.3, 0.4, 0.5]
- with patch.object(tool, "_get_embedding", AsyncMock(return_value=embedding)):
- with patch.object(tool, "_get_client", return_value=mock_client):
- results = await tool.run("vector embedding query")
+ def text(self) -> str:
+ return f"{self.status_code} {self.reason}"
- assert len(results.results) == 1
- assert results.results[0].content["title"] == "Vector Embedding Test"
+ mock_response = MockResponse(status_code=401, reason="Unauthorized")
+ with patch(
+ "azure.search.documents.aio.SearchClient.__init__", side_effect=HttpResponseError(response=mock_response)
+ ):
+ with pytest.raises(ValueError, match="Authentication failed"):
+ await tool._get_client() # pyright: ignore[reportPrivateUsage]
- mock_client.search.assert_called_once()
+ mock_response = MockResponse(status_code=403, reason="Forbidden")
+ with patch(
+ "azure.search.documents.aio.SearchClient.__init__", side_effect=HttpResponseError(response=mock_response)
+ ):
+ with pytest.raises(ValueError, match="Permission denied"):
+ await tool._get_client() # pyright: ignore[reportPrivateUsage]
- _, kwargs = mock_client.search.call_args
- assert "vector_queries" in kwargs
+ mock_response = MockResponse(status_code=500, reason="Internal Server Error")
+ with patch(
+ "azure.search.documents.aio.SearchClient.__init__", side_effect=HttpResponseError(response=mock_response)
+ ):
+ with pytest.raises(ValueError, match="Error connecting to Azure AI Search"):
+ await tool._get_client() # pyright: ignore[reportPrivateUsage]
@pytest.mark.asyncio
-async def test_error_handling_in_special_cases() -> None:
- """Test error handling for specific error cases that might be missed."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="error_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- )
+async def test_search_document_processing_error(mock_search_client: AsyncMock) -> None:
+ """Test error handling during search document processing."""
+ mock_search_client.search.return_value.__aiter__.return_value = [{"invalid": "document", "@search.score": None}]
- not_found_error = ValueError("The requested resource with 'test-index' was not found")
-
- with patch.object(tool, "_get_client", AsyncMock(side_effect=not_found_error)):
- with pytest.raises(ValueError, match="Index 'test-index' not found"):
- await tool.run("error query")
-
- auth_error = ValueError("401 Unauthorized error occurred")
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ )
- with patch.object(tool, "_get_client", AsyncMock(side_effect=auth_error)):
- with pytest.raises(ValueError, match="Authentication failed"):
- await tool.run("auth error query")
+ with patch.object(tool, "_get_client", return_value=mock_search_client):
+ results = await tool.run("test query")
+ assert len(results.results) == 0
@pytest.mark.asyncio
-async def test_component_loading_with_config_model() -> None:
- """Test the load_component method with a ComponentModel instead of dict."""
- from autogen_core import ComponentModel
-
- model = ComponentModel(
- provider="autogen_ext.tools.azure.BaseAzureAISearchTool",
- config={
- "name": "model_test",
- "endpoint": "https://test.search.windows.net",
- "index_name": "test-index",
- "credential": {"api_key": "test-key"},
- "query_type": "keyword",
- "search_fields": ["title", "content"],
- },
+async def test_search_with_expired_cache(
+ mock_search_client: AsyncMock, mock_search_results: List[Dict[str, Any]]
+) -> None:
+ """Test search with expired cache."""
+ mock_search_client.search.return_value.__aiter__.return_value = mock_search_results
+
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ enable_caching=True,
+ cache_ttl_seconds=1,
)
- with patch.object(ConcreteAzureAISearchTool, "create_keyword_search") as mock_create:
- mock_create.return_value = ConcreteAzureAISearchTool.create_keyword_search(
- name="model_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- )
+ with patch.object(tool, "_get_client", return_value=mock_search_client):
+ await tool.run("test query")
+ assert mock_search_client.search.call_count == 1
- tool = ConcreteAzureAISearchTool.load_component(model)
+ await asyncio.sleep(1.1)
- assert tool.name == "model_test"
+ await tool.run("test query")
+ assert mock_search_client.search.call_count == 2
@pytest.mark.asyncio
-async def test_fallback_vectorizable_text_query() -> None:
- """Test the fallback VectorizableTextQuery class when Azure SDK is not available."""
-
- class MockVectorizableTextQuery:
- def __init__(self, text: str, k_nearest_neighbors: int, fields: str) -> None:
- self.text = text
- self.k_nearest_neighbors = k_nearest_neighbors
- self.fields = fields
-
- query1 = MockVectorizableTextQuery(text="test query", k_nearest_neighbors=5, fields="title")
- assert query1.text == "test query"
- assert query1.fields == "title"
-
- query2 = MockVectorizableTextQuery(text="test query", k_nearest_neighbors=3, fields="title,content")
- assert query2.text == "test query"
- assert query2.fields == "title,content"
+async def test_search_with_invalid_credential() -> None:
+ """Test search with invalid credential format."""
+ with pytest.raises(ValueError, match="Invalid configuration"):
+ AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential={"api_key": "test-key"}, # type: ignore
+ )
@pytest.mark.asyncio
-async def test_dump_component() -> None:
- """Test the dump_component method."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="dump_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+async def test_search_with_empty_query() -> None:
+ """Test search with empty query."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- component_model = tool.dump_component()
- assert component_model.provider == "autogen_ext.tools.azure.BaseAzureAISearchTool"
- assert component_model.config["name"] == "dump_test"
- assert component_model.config["endpoint"] == "https://test.search.windows.net"
- assert component_model.config["index_name"] == "test-index"
+ with pytest.raises(ValueError, match="Search query cannot be empty"):
+ await tool.run("")
@pytest.mark.asyncio
-async def test_fallback_config_class() -> None:
- """Test the fallback configuration class."""
- from autogen_ext.tools.azure._ai_search import _FallbackAzureAISearchConfig # pyright: ignore[reportPrivateUsage]
-
- config = _FallbackAzureAISearchConfig(
- name="fallback_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- query_type="vector",
+async def test_vector_search_without_query() -> None:
+ """Test vector search with empty query."""
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
vector_fields=["embedding"],
- top=10,
)
- assert config.name == "fallback_test"
- assert config.endpoint == "https://test.search.windows.net"
- assert config.index_name == "test-index"
- assert config.query_type == "vector"
- assert config.vector_fields == ["embedding"]
- assert config.top == 10
+ with pytest.raises(ValueError, match="Search query cannot be empty"):
+ await tool.run("")
@pytest.mark.asyncio
-async def test_search_with_different_query_types() -> None:
- """Test search with different query types and parameters."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="query_types_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+async def test_search_with_cancellation_token_already_cancelled() -> None:
+ """Test search with already cancelled token."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
)
- mock_client = AsyncMock()
- mock_client.search.return_value = MockAsyncIterator([{"@search.score": 0.9, "title": "Test Result"}])
-
- with patch.object(tool, "_get_client", return_value=mock_client):
- await tool.run("string query")
- mock_client.search.assert_called_once()
- mock_client.search.reset_mock()
-
- await tool.run({"query": "dict query"})
- mock_client.search.assert_called_once()
- mock_client.search.reset_mock()
+ token = CancellationToken()
+ token.cancel()
- await tool.run(SearchQuery(query="object query"))
- mock_client.search.assert_called_once()
-
-
-class MockEmbeddingData:
- """Mock for OpenAI embedding data."""
-
- def __init__(self, embedding: List[float]):
- self.embedding = embedding
-
-
-class MockEmbeddingResponse:
- """Mock for OpenAI embedding response."""
-
- def __init__(self, data: List[MockEmbeddingData]):
- self.data = data
+ with pytest.raises(asyncio.CancelledError):
+ await tool.run("test query", cancellation_token=token)
@pytest.mark.asyncio
-async def test_get_embedding_methods() -> None:
- """Test the _get_embedding method with different providers."""
-
- class TestSearchTool(AzureAISearchTool):
- async def _get_embedding(self, query: str) -> List[float]:
- return [0.1, 0.2, 0.3]
+async def test_config_validation_edge_cases() -> None:
+ """Test configuration validation edge cases."""
+ with pytest.raises(
+ ValueError,
+ match="Invalid configuration: 1 validation error for AzureAISearchConfig\n Value error, vector_fields must be provided for vector search",
+ ):
+ AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=[],
+ )
- with patch.object(AzureAISearchTool, "_get_embedding", autospec=True) as mock_get_embedding:
- mock_get_embedding.return_value = [0.1, 0.2, 0.3]
+ with pytest.raises(ValueError, match="vector_fields must contain at least one field name for hybrid search"):
+ AzureAISearchTool.create_hybrid_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=[],
+ search_fields=["content"],
+ )
- tool = TestSearchTool.create_vector_search(
- name="test_vector_search",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+ with pytest.raises(ValueError, match="semantic_config_name is required when query_type is 'semantic'"):
+ AzureAISearchTool.create_hybrid_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
vector_fields=["embedding"],
+ search_fields=["content"],
+ query_type="semantic",
)
- result = await AzureAISearchTool._get_embedding(tool, "test query") # pyright: ignore[reportPrivateUsage]
- assert result == [0.1, 0.2, 0.3]
- mock_get_embedding.assert_called_once_with(tool, "test query")
-
@pytest.mark.asyncio
-async def test_get_embedding_azure_openai_path() -> None:
- """Test the Azure OpenAI path in _get_embedding."""
- mock_azure_openai = AsyncMock()
- mock_azure_openai.embeddings.create.return_value = MagicMock(data=[MagicMock(embedding=[0.1, 0.2, 0.3])])
+async def test_base_class_functionality() -> None:
+ """Test base class functionality."""
+ config = AzureAISearchConfig(
+ name="test",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ )
- with (
- patch("openai.AsyncAzureOpenAI", return_value=mock_azure_openai),
- patch("azure.identity.DefaultAzureCredential"),
- patch("autogen_ext.tools.azure._ai_search.getattr") as mock_getattr,
- ):
+ with pytest.raises(NotImplementedError, match="BaseAzureAISearchTool.*cannot be instantiated directly"):
+ BaseAzureAISearchTool._from_config(config) # pyright: ignore[reportPrivateUsage]
- def side_effect(obj: Any, name: str, default: Any = None) -> Any:
- if name == "embedding_provider":
- return "azure_openai"
- elif name == "embedding_model":
- return "text-embedding-ada-002"
- elif name == "openai_endpoint":
- return "https://test.openai.azure.com"
- elif name == "openai_api_key":
- return "test-key"
- return default
-
- mock_getattr.side_effect = side_effect
-
- class TestTool(AzureAISearchTool):
- async def _get_embedding(self, query: str) -> List[float]:
- return await AzureAISearchTool._get_embedding(self, query)
-
- token = _allow_private_constructor.set(True)
- try:
- tool = TestTool(
- name="test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- query_type="vector",
- vector_fields=["embedding"],
- )
+ class TestSearchTool(BaseAzureAISearchTool):
+ async def _get_embedding(self, query: str) -> List[float]:
+ return [0.1, 0.2, 0.3]
- result = await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage]
- assert result == [0.1, 0.2, 0.3]
- mock_azure_openai.embeddings.create.assert_called_once_with(
- model="text-embedding-ada-002", input="test query"
+ @classmethod
+ def _from_config(cls, config: AzureAISearchConfig) -> "TestSearchTool":
+ return cls(
+ name=config.name,
+ endpoint=config.endpoint,
+ index_name=config.index_name,
+ credential=config.credential,
)
- finally:
- _allow_private_constructor.reset(token)
+
+ tool = TestSearchTool._from_config(config) # pyright: ignore[reportPrivateUsage]
+ assert tool.name == "test"
+ assert tool.search_config.endpoint == MOCK_ENDPOINT
@pytest.mark.asyncio
-async def test_get_embedding_openai_path() -> None:
- """Test the OpenAI path in _get_embedding."""
- mock_openai = AsyncMock()
- mock_openai.embeddings.create.return_value = MagicMock(data=[MagicMock(embedding=[0.4, 0.5, 0.6])])
+async def test_client_cleanup_edge_cases() -> None:
+ """Test client cleanup edge cases."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ )
- with (
- patch("openai.AsyncOpenAI", return_value=mock_openai),
- patch("autogen_ext.tools.azure._ai_search.getattr") as mock_getattr,
- ):
+ tool._client = None # pyright: ignore[reportPrivateUsage]
+ await tool.close()
- def side_effect(obj: Any, name: str, default: Any = None) -> Any:
- if name == "embedding_provider":
- return "openai"
- elif name == "embedding_model":
- return "text-embedding-3-small"
- elif name == "openai_api_key":
- return "test-key"
- return default
+ mock_client = AsyncMock()
+ mock_client.close.side_effect = Exception("Failed to close")
+ tool._client = mock_client # pyright: ignore[reportPrivateUsage]
+ await tool.close()
+ assert tool._client is None # pyright: ignore[reportPrivateUsage]
- mock_getattr.side_effect = side_effect
- class TestTool(AzureAISearchTool):
- async def _get_embedding(self, query: str) -> List[float]:
- return await AzureAISearchTool._get_embedding(self, query)
+@pytest.mark.asyncio
+async def test_token_acquisition_edge_cases() -> None:
+ """Test token acquisition edge cases."""
+ with patch("azure.identity.DefaultAzureCredential") as mock_credential:
+ mock_credential.return_value.get_token.return_value = None
- token = _allow_private_constructor.set(True)
- try:
- tool = TestTool(
- name="test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- query_type="vector",
- vector_fields=["embedding"],
- )
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ openai_endpoint="https://test.openai.azure.com",
+ openai_api_version="2023-11-01",
+ )
- result = await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage]
- assert result == [0.4, 0.5, 0.6]
- mock_openai.embeddings.create.assert_called_once_with(model="text-embedding-3-small", input="test query")
- finally:
- _allow_private_constructor.reset(token)
+ with pytest.raises(ValueError, match="Failed to acquire token using DefaultAzureCredential for Azure OpenAI"):
+ await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage]
@pytest.mark.asyncio
-async def test_get_embedding_error_cases_direct() -> None:
- """Test error cases in the _get_embedding method."""
-
- class DirectEmbeddingTool(AzureAISearchTool):
- async def _get_embedding(self, query: str) -> List[float]:
- return await super()._get_embedding(query)
-
- token = _allow_private_constructor.set(True)
- try:
- tool = DirectEmbeddingTool(
- name="error_embedding_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- query_type="vector",
+async def test_hybrid_search_validation() -> None:
+ """Test hybrid search validation edge cases."""
+ with pytest.raises(ValueError, match="semantic_config_name is required when query_type is 'semantic'"):
+ AzureAISearchTool.create_hybrid_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
vector_fields=["embedding"],
+ search_fields=["content"],
+ query_type="semantic",
)
- with pytest.raises(
- ValueError, match="To use vector search, you must provide embedding_provider and embedding_model"
- ):
- await tool._get_embedding("test query")
-
- tool.search_config.embedding_provider = "azure_openai"
- with pytest.raises(
- ValueError, match="To use vector search, you must provide embedding_provider and embedding_model"
- ):
- await tool._get_embedding("test query")
+ with pytest.raises(ValueError, match="openai_endpoint is required when embedding_provider is 'azure_openai'"):
+ AzureAISearchTool.create_hybrid_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ search_fields=["content"],
+ embedding_provider="azure_openai",
+ embedding_model="text-embedding-ada-002",
+ )
- tool.search_config.embedding_model = "text-embedding-ada-002"
- def missing_endpoint_side_effect(obj: Any, name: str, default: Any = None) -> Any:
- if name == "openai_endpoint":
- return None
- return getattr(obj, name, default)
+@pytest.mark.asyncio
+async def test_search_result_caching() -> None:
+ """Test that search results are properly cached and retrieved."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ enable_caching=True,
+ cache_ttl_seconds=10,
+ )
- with patch(
- "autogen_ext.tools.azure._ai_search.getattr",
- side_effect=missing_endpoint_side_effect,
- ):
- with pytest.raises(ValueError, match="OpenAI endpoint must be provided"):
- await tool._get_embedding("test query")
+ mock_results = [{"id": "1", "content": "Test", "@search.score": 0.8}]
- tool.search_config.embedding_provider = "unsupported_provider"
+ with patch.object(tool, "_get_client") as mock_get_client:
+ mock_client = AsyncMock()
+ mock_client.search.return_value.__aiter__.return_value = mock_results
+ mock_get_client.return_value = mock_client
- def unsupported_provider_side_effect(obj: Any, name: str, default: Any = None) -> Any:
- if name == "openai_endpoint":
- return "https://test.openai.azure.com"
- return getattr(obj, name, default)
+ result1 = await tool.run("test query")
+ assert len(result1.results) == 1
+ assert mock_client.search.call_count == 1
- with patch(
- "autogen_ext.tools.azure._ai_search.getattr",
- side_effect=unsupported_provider_side_effect,
- ):
- with pytest.raises(ValueError, match="Unsupported embedding provider"):
- await tool._get_embedding("test query")
- finally:
- _allow_private_constructor.reset(token)
+ result2 = await tool.run("test query")
+ assert len(result2.results) == 1
+ assert mock_client.search.call_count == 1
@pytest.mark.asyncio
-async def test_azure_openai_with_default_credential() -> None:
- """Test Azure OpenAI with DefaultAzureCredential."""
-
- mock_azure_openai = AsyncMock()
- mock_azure_openai.embeddings.create.return_value = MagicMock(data=[MagicMock(embedding=[0.1, 0.2, 0.3])])
-
- mock_credential = MagicMock()
- mock_token = MagicMock()
- mock_token.token = "mock-token"
- mock_credential.get_token.return_value = mock_token
-
- with (
- patch("openai.AsyncAzureOpenAI") as mock_azure_openai_class,
- patch("azure.identity.DefaultAzureCredential", return_value=mock_credential),
- patch("autogen_ext.tools.azure._ai_search.getattr") as mock_getattr,
- ):
- mock_azure_openai_class.return_value = mock_azure_openai
-
- def side_effect(obj: Any, name: str, default: Any = None) -> Any:
- if name == "embedding_provider":
- return "azure_openai"
- elif name == "embedding_model":
- return "text-embedding-ada-002"
- elif name == "openai_endpoint":
- return "https://test.openai.azure.com"
- elif name == "openai_api_version":
- return "2023-05-15"
- return default
-
- mock_getattr.side_effect = side_effect
-
- class TestTool(AzureAISearchTool):
- async def _get_embedding(self, query: str) -> List[float]:
- return await AzureAISearchTool._get_embedding(self, query)
-
- token = _allow_private_constructor.set(True)
- try:
- tool = TestTool(
- name="test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
- query_type="vector",
- vector_fields=["embedding"],
- )
-
- token_provider: Optional[Callable[[], str]] = None
+async def test_cache_expiration() -> None:
+ """Test that cached results expire after TTL."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ enable_caching=True,
+ cache_ttl_seconds=1,
+ )
- def capture_token_provider(
- api_key: Optional[str] = None,
- azure_ad_token_provider: Optional[Callable[[], str]] = None,
- **kwargs: Any,
- ) -> AsyncMock:
- nonlocal token_provider
- if azure_ad_token_provider:
- token_provider = azure_ad_token_provider
- return mock_azure_openai
+ mock_results = [{"id": "1", "content": "Test", "@search.score": 0.8}]
- mock_azure_openai_class.side_effect = capture_token_provider
+ with patch.object(tool, "_get_client") as mock_get_client:
+ mock_client = AsyncMock()
+ mock_client.search.return_value.__aiter__.return_value = mock_results
+ mock_get_client.return_value = mock_client
- result = await tool._get_embedding("test query") # pyright: ignore[reportPrivateUsage]
- assert result == [0.1, 0.2, 0.3]
+ await tool.run("test query")
+ assert mock_client.search.call_count == 1
- assert token_provider is not None
- token_provider()
- mock_credential.get_token.assert_called_once_with("https://cognitiveservices.azure.com/.default")
+ await asyncio.sleep(1.1)
- mock_azure_openai.embeddings.create.assert_called_once_with(
- model="text-embedding-ada-002", input="test query"
- )
- finally:
- _allow_private_constructor.reset(token)
+ await tool.run("test query")
+ assert mock_client.search.call_count == 2
@pytest.mark.asyncio
-async def test_schema_property() -> None:
- """Test the schema property correctly defines the JSON schema for the tool."""
- tool = ConcreteAzureAISearchTool.create_keyword_search(
- name="schema_test",
- endpoint="https://test.search.windows.net",
- index_name="test-index",
- credential=AzureKeyCredential("test-key"),
+async def test_search_field_validation() -> None:
+ """Test validation of search fields configuration."""
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search", endpoint=MOCK_ENDPOINT, index_name=MOCK_INDEX, credential=MOCK_CREDENTIAL, search_fields=[]
)
+ assert tool.search_config.search_fields == []
- schema = tool.schema
+ tool = AzureAISearchTool.create_full_text_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ search_fields=["content", "content"],
+ )
+ assert tool.search_config.search_fields == ["content", "content"]
- assert schema["name"] == "schema_test"
- assert "description" in schema
- parameters = schema.get("parameters", {}) # pyright: ignore
- assert parameters.get("type") == "object" # pyright: ignore
+@pytest.mark.asyncio
+async def test_api_version_handling() -> None:
+ """Test handling of different API versions."""
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ api_version="2023-11-01",
+ )
+ assert tool.search_config.api_version == "2023-11-01"
- properties = parameters.get("properties", {}) # pyright: ignore
- assert "query" in properties # pyright: ignore
+ tool = AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ )
+ assert tool.search_config.api_version == DEFAULT_API_VERSION
- required = parameters.get("required", []) # pyright: ignore
- assert "query" in required # pyright: ignore
+ with patch("autogen_ext.tools.azure._ai_search.logger") as mock_logger:
+ AzureAISearchTool.create_vector_search(
+ name="test-search",
+ endpoint=MOCK_ENDPOINT,
+ index_name=MOCK_INDEX,
+ credential=MOCK_CREDENTIAL,
+ vector_fields=["embedding"],
+ api_version="2023-11-01",
+ )
- assert schema.get("strict") is True # pyright: ignore
+ mock_logger.warning.assert_called_once()
+ warning_msg = mock_logger.warning.call_args[0][0]
+ assert "vector search" in warning_msg.lower()
+ assert "2023-11-01" in warning_msg
diff --git a/python/packages/autogen-ext/tests/tools/http/test_http_tool.py b/python/packages/autogen-ext/tests/tools/http/test_http_tool.py
index 9d2898bae..8e50e48ba 100644
--- a/python/packages/autogen-ext/tests/tools/http/test_http_tool.py
+++ b/python/packages/autogen-ext/tests/tools/http/test_http_tool.py
@@ -176,7 +176,7 @@ async def test_invalid_request(test_config: ComponentModel, test_server: None) -
config.config["host"] = "fake"
tool = HttpTool.load_component(config)
- with pytest.raises(httpx.ConnectError):
+ with pytest.raises((httpx.ConnectError, httpx.ConnectTimeout)):
await tool.run_json({"query": "test query", "value": 42}, CancellationToken())
diff --git a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py
index de9fd5cb5..0a31fd128 100644
--- a/python/packages/autogen-ext/tests/tools/test_mcp_tools.py
+++ b/python/packages/autogen-ext/tests/tools/test_mcp_tools.py
@@ -1,12 +1,17 @@
+import asyncio
import logging
import os
+import threading
+from typing import cast
from unittest.mock import AsyncMock, MagicMock
import pytest
+from _pytest.logging import LogCaptureFixture # type: ignore[import]
from autogen_core import CancellationToken
from autogen_core.tools import Workbench
from autogen_core.utils import schema_to_pydantic_model
from autogen_ext.tools.mcp import (
+ McpSessionActor,
McpWorkbench,
SseMcpToolAdapter,
SseServerParams,
@@ -94,6 +99,14 @@ def cancellation_token() -> CancellationToken:
return CancellationToken()
+@pytest.fixture
+def mock_error_tool_response() -> MagicMock:
+ response = MagicMock()
+ response.isError = True
+ response.content = [TextContent(text="error output", type="text")]
+ return response
+
+
def test_adapter_config_serialization(sample_tool: Tool, sample_server_params: StdioServerParams) -> None:
"""Test that adapter can be saved to and loaded from config."""
original_adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
@@ -594,4 +607,165 @@ async def test_lazy_init_and_finalize_cleanup() -> None:
assert workbench._actor is not None # type: ignore[reportPrivateUsage]
assert workbench._actor._active is True # type: ignore[reportPrivateUsage]
+ actor = workbench._actor # type: ignore[reportPrivateUsage]
del workbench
+ await asyncio.sleep(0.1)
+ assert actor._active is False
+
+
+@pytest.mark.asyncio
+async def test_del_to_new_event_loop_when_get_event_loop_fails() -> None:
+ params = StdioServerParams(
+ command="npx",
+ args=[
+ "-y",
+ "@modelcontextprotocol/server-filesystem",
+ ".",
+ ],
+ read_timeout_seconds=60,
+ )
+ workbench = McpWorkbench(server_params=params)
+
+ await workbench.list_tools()
+ assert workbench._actor is not None # type: ignore[reportPrivateUsage]
+ assert workbench._actor._active is True # type: ignore[reportPrivateUsage]
+
+ actor = workbench._actor # type: ignore[reportPrivateUsage]
+
+ def cleanup() -> None:
+ nonlocal workbench
+ del workbench
+
+ t = threading.Thread(target=cleanup)
+ t.start()
+ t.join()
+
+ await asyncio.sleep(0.1)
+ assert actor._active is False # type: ignore[reportPrivateUsage]
+
+
+def test_del_raises_when_loop_closed() -> None:
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+
+ params = StdioServerParams(command="echo", args=["ok"])
+ workbench = McpWorkbench(server_params=params)
+
+ workbench._actor_loop = loop # type: ignore[reportPrivateUsage]
+ workbench._actor = cast(McpSessionActor, object()) # type: ignore[reportPrivateUsage]
+
+ loop.close()
+
+ with pytest.warns(RuntimeWarning, match="loop is closed or not running"):
+ del workbench
+
+
+def test_mcp_tool_adapter_normalize_payload(sample_tool: Tool, sample_server_params: StdioServerParams) -> None:
+ """Test the _normalize_payload_to_content_list method of McpToolAdapter."""
+ adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
+
+ # Case 1: Payload is already a list of valid content items
+ valid_content_list: list[TextContent | ImageContent | EmbeddedResource] = [
+ TextContent(text="hello", type="text"),
+ ImageContent(data="base64data", mimeType="image/png", type="image"),
+ EmbeddedResource(
+ type="resource",
+ resource=TextResourceContents(text="embedded text", uri=AnyUrl(url="http://example.com/resource")),
+ ),
+ ]
+ assert adapter._normalize_payload_to_content_list(valid_content_list) == valid_content_list # type: ignore[reportPrivateUsage]
+
+ # Case 2: Payload is a single TextContent
+ single_text_content = TextContent(text="single text", type="text")
+ assert adapter._normalize_payload_to_content_list(single_text_content) == [single_text_content] # type: ignore[reportPrivateUsage, arg-type]
+
+ # Case 3: Payload is a single ImageContent
+ single_image_content = ImageContent(data="imagedata", mimeType="image/jpeg", type="image")
+ assert adapter._normalize_payload_to_content_list(single_image_content) == [single_image_content] # type: ignore[reportPrivateUsage, arg-type]
+
+ # Case 4: Payload is a single EmbeddedResource
+ single_embedded_resource = EmbeddedResource(
+ type="resource",
+ resource=TextResourceContents(text="other embedded", uri=AnyUrl(url="http://example.com/other")),
+ )
+ assert adapter._normalize_payload_to_content_list(single_embedded_resource) == [single_embedded_resource] # type: ignore[reportPrivateUsage, arg-type]
+
+ # Case 5: Payload is a string
+ string_payload = "This is a string payload."
+ expected_from_string = [TextContent(text=string_payload, type="text")]
+ assert adapter._normalize_payload_to_content_list(string_payload) == expected_from_string # type: ignore[reportPrivateUsage, arg-type]
+
+ # Case 6: Payload is an integer
+ int_payload = 12345
+ expected_from_int = [TextContent(text=str(int_payload), type="text")]
+ assert adapter._normalize_payload_to_content_list(int_payload) == expected_from_int # type: ignore[reportPrivateUsage, arg-type]
+
+ # Case 7: Payload is a dictionary
+ dict_payload = {"key": "value", "number": 42}
+ expected_from_dict = [TextContent(text=str(dict_payload), type="text")]
+ assert adapter._normalize_payload_to_content_list(dict_payload) == expected_from_dict # type: ignore[reportPrivateUsage, arg-type]
+
+ # Case 8: Payload is an empty list (should still be a list of valid items, so returns as is)
+ empty_list_payload: list[TextContent | ImageContent | EmbeddedResource] = []
+ assert adapter._normalize_payload_to_content_list(empty_list_payload) == empty_list_payload # type: ignore[reportPrivateUsage]
+
+ # Case 9: Payload is None (should be stringified)
+ none_payload = None
+ expected_from_none = [TextContent(text=str(none_payload), type="text")]
+ assert adapter._normalize_payload_to_content_list(none_payload) == expected_from_none # type: ignore[reportPrivateUsage, arg-type]
+
+
+@pytest.mark.asyncio
+async def test_mcp_tool_adapter_run_error(
+ sample_tool: Tool,
+ sample_server_params: StdioServerParams,
+ mock_session: AsyncMock,
+ mock_error_tool_response: MagicMock,
+ cancellation_token: CancellationToken,
+) -> None:
+ """Test McpToolAdapter._run when tool returns an error."""
+ adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session)
+ mock_session.call_tool.return_value = mock_error_tool_response
+
+ args = {"test_param": "test_value"}
+ with pytest.raises(Exception) as excinfo:
+ await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session) # type: ignore[reportPrivateUsage]
+
+ mock_session.call_tool.assert_called_once_with(name=sample_tool.name, arguments=args)
+ assert adapter.return_value_as_string([TextContent(text="error output", type="text")]) in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_mcp_tool_adapter_run_cancelled_before_call(
+ sample_tool: Tool,
+ sample_server_params: StdioServerParams,
+ mock_session: AsyncMock,
+ cancellation_token: CancellationToken,
+) -> None:
+ """Test McpToolAdapter._run when operation is cancelled before tool call."""
+ adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session)
+ cancellation_token.cancel() # Cancel before the call
+
+ args = {"test_param": "test_value"}
+ with pytest.raises(asyncio.CancelledError):
+ await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session) # type: ignore[reportPrivateUsage]
+
+ mock_session.call_tool.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_mcp_tool_adapter_run_cancelled_during_call(
+ sample_tool: Tool,
+ sample_server_params: StdioServerParams,
+ mock_session: AsyncMock,
+ cancellation_token: CancellationToken,
+) -> None:
+ """Test McpToolAdapter._run when operation is cancelled during tool call."""
+ adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool, session=mock_session)
+ mock_session.call_tool.side_effect = asyncio.CancelledError("Tool call cancelled")
+
+ args = {"test_param": "test_value"}
+ with pytest.raises(asyncio.CancelledError):
+ await adapter._run(args=args, cancellation_token=cancellation_token, session=mock_session) # type: ignore[reportPrivateUsage]
+
+ mock_session.call_tool.assert_called_once_with(name=sample_tool.name, arguments=args)
diff --git a/python/uv.lock b/python/uv.lock
index 0b2e4f543..6fc4faa3f 100644
--- a/python/uv.lock
+++ b/python/uv.lock
@@ -455,7 +455,7 @@ wheels = [
[[package]]
name = "autogen-agentchat"
-version = "0.5.6"
+version = "0.5.7"
source = { editable = "packages/autogen-agentchat" }
dependencies = [
{ name = "autogen-core" },
@@ -466,7 +466,7 @@ requires-dist = [{ name = "autogen-core", editable = "packages/autogen-core" }]
[[package]]
name = "autogen-core"
-version = "0.5.6"
+version = "0.5.7"
source = { editable = "packages/autogen-core" }
dependencies = [
{ name = "jsonref" },
@@ -585,7 +585,7 @@ dev = [
[[package]]
name = "autogen-ext"
-version = "0.5.6"
+version = "0.5.7"
source = { editable = "packages/autogen-ext" }
dependencies = [
{ name = "autogen-core" },