agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: [BREAKING] cleanup of thread API and serialization (#893)

Browse Source 
* cleanup of threads and serialization

* fix for sliding window

* fix redis test

* updated from comments

* updated context provider and threads

* updated lock

* add asyncio default

* fix redis tests

* fix tests

* fix tests

* renamed to invoking

* fixed tests

* fix for instructions
This commit is contained in:
Eduard van Valkenburg
2025-09-29 18:22:34 +02:00
committed by GitHub
Unverified
parent bf5931932e
commit 10d10364a9
52 changed files with 1642 additions and 1411 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/agent_framework/_agents.py
+227 -147
View File
@@ -4,18 +4,19 @@ import inspect
import sys
from collections.abc import AsyncIterable, Awaitable, Callable, MutableMapping, Sequence
from contextlib import AbstractAsyncContextManager, AsyncExitStack
from typing import Any, ClassVar, Literal, Protocol, TypeVar, runtime_checkable
from copy import copy
from itertools import chain
from typing import Any, ClassVar, Literal, Protocol, TypeVar, cast, runtime_checkable
from uuid import uuid4
from pydantic import BaseModel, Field, PrivateAttr, create_model
from pydantic import BaseModel, Field, create_model
from ._clients import BaseChatClient, ChatClientProtocol
from ._logging import get_logger
from ._mcp import MCPTool
from ._memory import AggregateContextProvider, Context, ContextProvider
from ._middleware import Middleware, use_agent_middleware
from ._pydantic import AFBaseModel
from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages
from ._threads import AgentThread, ChatMessageStoreProtocol
from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, AIFunction, ToolProtocol
from ._types import (
AgentRunResponse,
@@ -30,6 +31,10 @@ from ._types import (
from .exceptions import AgentExecutionException
from .observability import use_agent_observability
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore[import] # pragma: no cover
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
else:
@@ -122,7 +127,7 @@ class AgentProtocol(Protocol):
"""
...
def get_new_thread(self) -> AgentThread:
def get_new_thread(self, **kwargs: Any) -> AgentThread:
"""Creates a new conversation thread for the agent."""
...
@@ -130,7 +135,7 @@ class AgentProtocol(Protocol):
# region BaseAgent
class BaseAgent(AFBaseModel):
class BaseAgent:
"""Base class for all Agent Framework agents.
Attributes:
@@ -143,18 +148,55 @@ class BaseAgent(AFBaseModel):
middleware: List of middleware to intercept agent and function invocations.
"""
id: str = Field(default_factory=lambda: str(uuid4()))
name: str | None = None
description: str | None = None
context_providers: AggregateContextProvider | None = None
middleware: Middleware | list[Middleware] | None = None
def __init__(
self,
id: str | None = None,
name: str | None = None,
description: str | None = None,
context_providers: ContextProvider | Sequence[ContextProvider] | None = None,
middleware: Middleware | Sequence[Middleware] | None = None,
**kwargs: Any,
) -> None:
"""Base class for all Agent Framework agents.
Args:
id: The unique identifier of the agent If no id is provided,
a new UUID will be generated.
name: The name of the agent, can be None.
description: The description of the agent.
display_name: The display name of the agent, which is either the name or id.
context_providers: The collection of multiple context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
kwargs: will be stored in `additional_properties`
"""
if id is None:
id = str(uuid4())
self.id = id
self.name = name
self.description = description
self.context_provider = self._prepare_context_providers(context_providers)
if middleware is None or isinstance(middleware, Sequence):
self.middleware: list[Middleware] | None = cast(list[Middleware], middleware) if middleware else None
else:
self.middleware = [middleware]
self.additional_properties = kwargs
async def _notify_thread_of_new_messages(
self, thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage]
self,
thread: AgentThread,
input_messages: ChatMessage | Sequence[ChatMessage],
response_messages: ChatMessage | Sequence[ChatMessage],
) -> None:
"""Notify the thread of new messages."""
if isinstance(new_messages, ChatMessage) or len(new_messages) > 0:
await thread_on_new_messages(thread, new_messages)
"""Notify the thread of new messages.
This also calls the invoked method of a potential context provider on the thread.
"""
if isinstance(input_messages, ChatMessage) or len(input_messages) > 0:
await thread.on_new_messages(input_messages)
if isinstance(response_messages, ChatMessage) or len(response_messages) > 0:
await thread.on_new_messages(response_messages)
if thread.context_provider:
await thread.context_provider.invoked(input_messages, response_messages)
@property
def display_name(self) -> str:
@@ -164,14 +206,14 @@ class BaseAgent(AFBaseModel):
"""
return self.name or self.id
def get_new_thread(self) -> AgentThread:
def get_new_thread(self, **kwargs: Any) -> AgentThread:
"""Returns AgentThread instance that is compatible with the agent."""
return AgentThread()
return AgentThread(**kwargs, context_provider=self.context_provider)
async def deserialize_thread(self, serialized_thread: Any, **kwargs: Any) -> AgentThread:
"""Deserializes the thread."""
thread: AgentThread = self.get_new_thread()
await deserialize_thread_state(thread, serialized_thread, **kwargs)
await thread.deserialize(serialized_thread, **kwargs)
return thread
def as_tool(
@@ -236,7 +278,12 @@ class BaseAgent(AFBaseModel):
# Create final text from accumulated updates
return AgentRunResponse.from_agent_run_response_updates(response_updates).text
return AIFunction(name=tool_name, description=tool_description, func=agent_wrapper, input_model=input_model)
return AIFunction(
name=tool_name,
description=tool_description,
func=agent_wrapper,
input_model=input_model,
)
def _normalize_messages(
self,
@@ -253,6 +300,18 @@ class BaseAgent(AFBaseModel):
return [ChatMessage(role=Role.USER, text=msg) if isinstance(msg, str) else msg for msg in messages]
def _prepare_context_providers(
self,
context_providers: ContextProvider | Sequence[ContextProvider] | None = None,
) -> AggregateContextProvider | None:
if not context_providers:
return None
if isinstance(context_providers, AggregateContextProvider):
return context_providers
return AggregateContextProvider(context_providers)
# region ChatAgent
@@ -263,12 +322,6 @@ class ChatAgent(BaseAgent):
"""A Chat Client Agent."""
AGENT_SYSTEM_NAME: ClassVar[str] = "microsoft.agent_framework"
chat_client: ChatClientProtocol
instructions: str | None = None
chat_options: ChatOptions
chat_message_store_factory: Callable[[], ChatMessageStore] | None = None
_local_mcp_tools: list[MCPTool] = PrivateAttr(default_factory=list) # type: ignore[reportUnknownVariableType]
_async_exit_stack: AsyncExitStack = PrivateAttr(default_factory=AsyncExitStack)
def __init__(
self,
@@ -278,6 +331,9 @@ class ChatAgent(BaseAgent):
id: str | None = None,
name: str | None = None,
description: str | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
middleware: Middleware | list[Middleware] | None = None,
frequency_penalty: float | None = None,
logit_bias: dict[str | int, float] | None = None,
max_tokens: int | None = None,
@@ -297,10 +353,7 @@ class ChatAgent(BaseAgent):
| None = None,
top_p: float | None = None,
user: str | None = None,
additional_properties: dict[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
middleware: Middleware | list[Middleware] | None = None,
request_kwargs: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
"""Create a ChatAgent.
@@ -317,6 +370,10 @@ class ChatAgent(BaseAgent):
id: The unique identifier for the agent, will be created automatically if not provided.
name: The name of the agent.
description: A brief description of the agent's purpose.
chat_message_store_factory: factory function to create an instance of ChatMessageStoreProtocol.
If not provided, the default in-memory store will be used.
context_providers: The collection of multiple context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
frequency_penalty: the frequency penalty to use.
logit_bias: the logit bias to use.
max_tokens: The maximum number of tokens to generate.
@@ -332,64 +389,54 @@ class ChatAgent(BaseAgent):
tools: the tools to use for the request.
top_p: the nucleus sampling probability to use.
user: the user to associate with the request.
additional_properties: additional properties to include in the request.
chat_message_store_factory: factory function to create an instance of ChatMessageStore. If not provided,
the default in-memory store will be used.
context_providers: The collection of multiple context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
kwargs: any additional keyword arguments.
Unused, can be used by subclasses of this Agent.
request_kwargs: a dictionary of other values that will be passed through
to the chat_client `get_response` and `get_streaming_response` methods.
kwargs: any additional keyword arguments. Will be stored as `additional_properties`
"""
if not hasattr(chat_client, FUNCTION_INVOKING_CHAT_CLIENT_MARKER) and isinstance(chat_client, BaseChatClient):
logger.warning(
"The provided chat client does not support function invoking, this might limit agent capabilities."
)
kwargs.update(additional_properties or {})
aggregate_context_providers = self._prepare_context_providers(context_providers)
super().__init__(
id=id,
name=name,
description=description,
context_providers=context_providers,
middleware=middleware,
**kwargs,
)
self.chat_client = chat_client
self.chat_message_store_factory = chat_message_store_factory
# We ignore the MCP Servers here and store them separately,
# we add their functions to the tools list at runtime
normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
local_mcp_tools = [tool for tool in normalized_tools if isinstance(tool, MCPTool)]
final_tools = [tool for tool in normalized_tools if not isinstance(tool, MCPTool)]
args: dict[str, Any] = {
"chat_client": chat_client,
"chat_message_store_factory": chat_message_store_factory,
"context_providers": aggregate_context_providers,
"middleware": middleware,
"chat_options": ChatOptions(
ai_model_id=model,
frequency_penalty=frequency_penalty,
logit_bias=logit_bias,
max_tokens=max_tokens,
metadata=metadata,
presence_penalty=presence_penalty,
response_format=response_format,
seed=seed,
stop=stop,
store=store,
temperature=temperature,
tool_choice=tool_choice,
tools=final_tools, # type: ignore[reportArgumentType]
top_p=top_p,
user=user,
additional_properties=kwargs,
),
}
if instructions is not None:
args["instructions"] = instructions
if name is not None:
args["name"] = name
if description is not None:
args["description"] = description
if id is not None:
args["id"] = id
super().__init__(**args)
normalized_tools: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = ( # type:ignore[reportUnknownVariableType]
[] if tools is None else tools if isinstance(tools, list) else [tools]
)
self._local_mcp_tools = [tool for tool in normalized_tools if isinstance(tool, MCPTool)]
agent_tools = [tool for tool in normalized_tools if not isinstance(tool, MCPTool)]
self.chat_options = ChatOptions(
ai_model_id=model,
frequency_penalty=frequency_penalty,
instructions=instructions,
logit_bias=logit_bias,
max_tokens=max_tokens,
metadata=metadata,
presence_penalty=presence_penalty,
response_format=response_format,
seed=seed,
stop=stop,
store=store,
temperature=temperature,
tool_choice=tool_choice,
tools=agent_tools, # type: ignore[reportArgumentType]
top_p=top_p,
user=user,
additional_properties=request_kwargs or {}, # type: ignore
)
self._async_exit_stack = AsyncExitStack()
self._update_agent_name()
self._local_mcp_tools = local_mcp_tools # type: ignore[assignment]
async def __aenter__(self) -> "Self":
"""Async context manager entry.
@@ -399,16 +446,17 @@ class ChatAgent(BaseAgent):
This list might be extended in the future.
"""
context_managers = [self.chat_client, *self._local_mcp_tools]
if self.context_providers:
context_managers.append(self.context_providers)
for context_manager in context_managers:
for context_manager in chain([self.chat_client], self._local_mcp_tools):
if isinstance(context_manager, AbstractAsyncContextManager):
await self._async_exit_stack.enter_async_context(context_manager)
return self
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: Any,
) -> None:
"""Async context manager exit.
Close the async exit stack to ensure all context managers are exited properly.
@@ -443,11 +491,9 @@ class ChatAgent(BaseAgent):
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
tools: ToolProtocol
| list[ToolProtocol]
| Callable[..., Any]
| list[Callable[..., Any]]
| MutableMapping[str, Any]
| list[MutableMapping[str, Any]]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
@@ -485,28 +531,32 @@ class ChatAgent(BaseAgent):
will only be passed to functions that are called.
"""
input_messages = self._normalize_messages(messages)
context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
thread, thread_messages = await self._prepare_thread_and_messages(
thread=thread, context=context, input_messages=input_messages
thread, run_chat_options, thread_messages = await self._prepare_thread_and_messages(
thread=thread, input_messages=input_messages
)
normalized_tools: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = ( # type:ignore[reportUnknownVariableType]
[] if tools is None else tools if isinstance(tools, list) else [tools]
)
agent_name = self._get_agent_name()
# Resolve final tool list (runtime provided tools + local MCP server tools)
final_tools: list[ToolProtocol | Callable[..., Any] | dict[str, Any]] = []
# Normalize tools argument to a list without mutating the original parameter
normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
for tool in normalized_tools:
if isinstance(tool, MCPTool):
if not tool.is_connected:
await self._async_exit_stack.enter_async_context(tool)
final_tools.extend(tool.functions) # type: ignore
else:
final_tools.append(tool) # type: ignore
for mcp_server in self._local_mcp_tools:
if not mcp_server.is_connected:
await self._async_exit_stack.enter_async_context(mcp_server)
final_tools.extend(mcp_server.functions)
response = await self.chat_client.get_response(
messages=thread_messages,
chat_options=self.chat_options
chat_options=run_chat_options
& ChatOptions(
ai_model_id=model,
conversation_id=thread.service_thread_id,
@@ -529,7 +579,7 @@ class ChatAgent(BaseAgent):
**kwargs,
)
self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
await self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
# Ensure that the author name is set for each message in the response.
for message in response.messages:
@@ -538,13 +588,7 @@ class ChatAgent(BaseAgent):
# Only notify the thread of new messages if the chatResponse was successful
# to avoid inconsistent messages state in the thread.
await self._notify_thread_of_new_messages(thread, input_messages)
await self._notify_thread_of_new_messages(thread, response.messages)
if self.context_providers:
await self.context_providers.thread_created(response.conversation_id)
await self.context_providers.messages_adding(thread.service_thread_id, input_messages + response.messages)
await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
return AgentRunResponse(
messages=response.messages,
response_id=response.response_id,
@@ -614,29 +658,34 @@ class ChatAgent(BaseAgent):
"""
input_messages = self._normalize_messages(messages)
context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
thread, thread_messages = await self._prepare_thread_and_messages(
thread=thread, context=context, input_messages=input_messages
thread, run_chat_options, thread_messages = await self._prepare_thread_and_messages(
thread=thread, input_messages=input_messages
)
agent_name = self._get_agent_name()
response_updates: list[ChatResponseUpdate] = []
# Resolve final tool list (runtime provided tools + local MCP server tools)
final_tools: list[ToolProtocol | MutableMapping[str, Any] | Callable[..., Any]] = []
normalized_tools: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = ( # type: ignore[reportUnknownVariableType]
[] if tools is None else tools if isinstance(tools, list) else [tools]
)
# Normalize tools argument to a list without mutating the original parameter
normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
for tool in normalized_tools:
if isinstance(tool, MCPTool):
if not tool.is_connected:
await self._async_exit_stack.enter_async_context(tool)
final_tools.extend(tool.functions) # type: ignore
else:
final_tools.append(tool)
for mcp_server in self._local_mcp_tools:
if not mcp_server.is_connected:
await self._async_exit_stack.enter_async_context(mcp_server)
final_tools.extend(mcp_server.functions)
async for update in self.chat_client.get_streaming_response(
messages=thread_messages,
chat_options=self.chat_options
chat_options=run_chat_options
& ChatOptions(
conversation_id=thread.service_thread_id,
frequency_penalty=frequency_penalty,
@@ -675,27 +724,46 @@ class ChatAgent(BaseAgent):
)
response = ChatResponse.from_chat_response_updates(response_updates)
await self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
@override
def get_new_thread(
self,
*,
service_thread_id: str | None = None,
**kwargs: Any,
) -> AgentThread:
"""Get a new conversation thread for the agent.
# Only notify the thread of new messages if the chatResponse was successful
# to avoid inconsistent messages state in the thread.
await self._notify_thread_of_new_messages(thread, input_messages)
await self._notify_thread_of_new_messages(thread, response.messages)
If you supply a service_thread_id, the thread will be marked as service managed.
if self.context_providers:
await self.context_providers.thread_created(response.conversation_id)
await self.context_providers.messages_adding(thread.service_thread_id, input_messages + response.messages)
If you don't supply a service_thread_id but have a chat_message_store_factory configured on the agent,
that factory will be used to create a message store for the thread and the thread will be
managed locally.
def get_new_thread(self) -> AgentThread:
message_store: ChatMessageStore | None = None
When neither is present, the thread will be created without a service ID or message store,
this will be updated based on usage, when you run the agent with this thread.
If you run with store=True, the response will respond with a thread_id and that will be set.
Otherwise a messages store is created from the default factory.
if self.chat_message_store_factory:
message_store = self.chat_message_store_factory()
Args:
service_thread_id: Optional service managed thread ID.
kwargs: not used at present.
"""
if service_thread_id is not None:
return AgentThread(
service_thread_id=service_thread_id,
context_provider=self.context_provider,
)
if self.chat_message_store_factory is not None:
return AgentThread(
message_store=self.chat_message_store_factory(),
context_provider=self.context_provider,
)
return AgentThread(context_provider=self.context_provider)
return AgentThread() if message_store is None else AgentThread(message_store=message_store)
def _update_thread_with_type_and_conversation_id(
async def _update_thread_with_type_and_conversation_id(
self, thread: AgentThread, response_conversation_id: str | None
) -> None:
"""Update thread with storage type and conversation ID.
@@ -719,6 +787,8 @@ class ChatAgent(BaseAgent):
# If we got a conversation id back from the chat client, it means that the service
# supports server side thread storage so we should update the thread with the new id.
thread.service_thread_id = response_conversation_id
if thread.context_provider:
await thread.context_provider.thread_created(thread.service_thread_id)
elif thread.message_store is None and self.chat_message_store_factory is not None:
# If the service doesn't use service side thread storage (i.e. we got no id back from invocation), and
# the thread has no message_store yet, and we have a custom messages store, we should update the thread
@@ -729,14 +799,14 @@ class ChatAgent(BaseAgent):
self,
*,
thread: AgentThread | None,
context: Context | None,
input_messages: list[ChatMessage] | None = None,
) -> tuple[AgentThread, list[ChatMessage]]:
) -> tuple[AgentThread, ChatOptions, list[ChatMessage]]:
"""Prepare the messages for agent execution.
Also updates the chat_options of the agent, with
Args:
thread: The conversation thread.
context: Context to include in messages.
input_messages: Messages to process.
Returns:
@@ -745,32 +815,42 @@ class ChatAgent(BaseAgent):
Raises:
AgentExecutionException: If the thread is not of the expected type.
"""
chat_options = copy(self.chat_options) if self.chat_options else ChatOptions()
thread = thread or self.get_new_thread()
messages: list[ChatMessage] = []
if self.instructions:
messages.append(ChatMessage(role=Role.SYSTEM, text=self.instructions))
if context and context.contents:
messages.append(ChatMessage(role=Role.SYSTEM, contents=context.contents))
if thread.service_thread_id and thread.context_provider:
await thread.context_provider.thread_created(thread.service_thread_id)
thread_messages: list[ChatMessage] = []
if thread.message_store:
messages.extend(await thread.message_store.list_messages() or [])
messages.extend(input_messages or [])
return thread, messages
thread_messages.extend(await thread.message_store.list_messages() or [])
context: Context | None = None
if self.context_provider:
async with self.context_provider:
context = await self.context_provider.invoking(input_messages or [])
if context:
if context.messages:
thread_messages.extend(context.messages)
if context.tools:
if chat_options.tools is not None:
chat_options.tools.extend(context.tools)
else:
chat_options.tools = list(context.tools)
if context.instructions:
chat_options.instructions = (
context.instructions
if not chat_options.instructions
else f"{chat_options.instructions}\n{context.instructions}"
)
thread_messages.extend(input_messages or [])
if (
thread.service_thread_id
and chat_options.conversation_id
and thread.service_thread_id != chat_options.conversation_id
):
raise AgentExecutionException(
"The conversation_id set on the agent is different from the one set on the thread, "
"only one ID can be used for a run."
)
return thread, chat_options, thread_messages
def _get_agent_name(self) -> str:
return self.name or "UnnamedAgent"
def _prepare_context_providers(
self,
context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
) -> AggregateContextProvider | None:
if not context_providers:
return None
if isinstance(context_providers, AggregateContextProvider):
return context_providers
if isinstance(context_providers, ContextProvider):
return AggregateContextProvider([context_providers])
return AggregateContextProvider(context_providers)
python/packages/main/agent_framework/_clients.py
+10 -7
View File
@@ -18,7 +18,7 @@ from ._middleware import (
Middleware,
)
from ._pydantic import AFBaseModel
from ._threads import ChatMessageStore
from ._threads import ChatMessageStoreProtocol
from ._tools import ToolProtocol
from ._types import (
ChatMessage,
@@ -207,9 +207,12 @@ class BaseChatClient(AFBaseModel, ABC):
# This is used for OTel setup, should be overridden in subclasses
def prepare_messages(
self, messages: str | ChatMessage | list[str] | list[ChatMessage]
self, messages: str | ChatMessage | list[str] | list[ChatMessage], chat_options: ChatOptions
) -> MutableSequence[ChatMessage]:
"""Turn the allowed input into a list of chat messages."""
if chat_options.instructions:
system_msg = ChatMessage(role="system", text=chat_options.instructions)
return [system_msg, *prepare_messages(messages)]
return prepare_messages(messages)
def _filter_internal_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
@@ -368,7 +371,7 @@ class BaseChatClient(AFBaseModel, ABC):
user=user,
additional_properties=additional_properties or {},
)
prepped_messages = self.prepare_messages(messages)
prepped_messages = self.prepare_messages(messages, chat_options)
self._prepare_tool_choice(chat_options=chat_options)
filtered_kwargs = self._filter_internal_kwargs(kwargs)
@@ -449,7 +452,7 @@ class BaseChatClient(AFBaseModel, ABC):
user=user,
additional_properties=additional_properties or {},
)
prepped_messages = self.prepare_messages(messages)
prepped_messages = self.prepare_messages(messages, chat_options)
self._prepare_tool_choice(chat_options=chat_options)
filtered_kwargs = self._filter_internal_kwargs(kwargs)
@@ -492,7 +495,7 @@ class BaseChatClient(AFBaseModel, ABC):
| MutableMapping[str, Any]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
middleware: Middleware | list[Middleware] | None = None,
**kwargs: Any,
@@ -503,8 +506,8 @@ class BaseChatClient(AFBaseModel, ABC):
name: The name of the agent.
instructions: The instructions for the agent.
tools: Optional list of tools to associate with the agent.
chat_message_store_factory: Factory function to create an instance of ChatMessageStore. If not provided,
the default in-memory store will be used.
chat_message_store_factory: Factory function to create an instance of ChatMessageStoreProtocol.
If not provided, the default in-memory store will be used.
context_providers: Context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
**kwargs: Additional keyword arguments to pass to the agent.
python/packages/main/agent_framework/_mcp.py
+3
View File
@@ -246,6 +246,7 @@ class MCPTool:
self.request_timeout = request_timeout
self.chat_client = chat_client
self.functions: list[AIFunction[Any, Any]] = []
self.is_connected: bool = False
def __str__(self) -> str:
return f"MCPTool(name={self.name}, description={self.description})"
@@ -282,6 +283,7 @@ class MCPTool:
# If the session is not initialized, we need to reinitialize it
await self.session.initialize()
logger.debug("Connected to MCP server: %s", self.session)
self.is_connected = True
if self.load_tools_flag:
await self.load_tools()
if self.load_prompts_flag:
@@ -434,6 +436,7 @@ class MCPTool:
"""Disconnect from the MCP server."""
await self._exit_stack.aclose()
self.session = None
self.is_connected = False
@abstractmethod
def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
python/packages/main/agent_framework/_memory.py
+83 -38
View File
@@ -6,11 +6,15 @@ from abc import ABC, abstractmethod
from collections.abc import MutableSequence, Sequence
from contextlib import AsyncExitStack
from types import TracebackType
from typing import ClassVar
from typing import Any, Final, cast
from ._pydantic import AFBaseModel
from ._types import ChatMessage, Contents
from ._tools import ToolProtocol
from ._types import ChatMessage
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore[import] # pragma: no cover
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
else:
@@ -21,7 +25,7 @@ else:
__all__ = ["AggregateContextProvider", "Context", "ContextProvider"]
class Context(AFBaseModel):
class Context:
"""A class containing any context that should be provided to the AI model as supplied by an ContextProvider.
Each ContextProvider has the ability to provide its own context for each invocation.
@@ -30,28 +34,39 @@ class Context(AFBaseModel):
This context is per invocation, and will not be stored as part of the chat history.
"""
contents: list[Contents] | None = None
"""
Any content to pass to the AI model in addition to any other prompts
that it may already have (in the case of an agent), or chat history that may already exist.
"""
def __init__(
self,
instructions: str | None = None,
messages: Sequence[ChatMessage] | None = None,
tools: Sequence[ToolProtocol] | None = None,
):
"""Create a new Context object.
Args:
instructions: Instructions to provide to the AI model.
messages: a list of messages.
tools: a list of tools to provide to this run.
"""
self.instructions = instructions
self.messages: Sequence[ChatMessage] = messages or []
self.tools: Sequence[ToolProtocol] = tools or []
# region ContextProvider
class ContextProvider(AFBaseModel, ABC):
class ContextProvider(ABC):
"""Base class for all context providers.
A context provider is a component that can be used to enhance the AI's context management.
It can listen to changes in the conversation and provide additional context to the AI model
just before invocation.
It also has a default memory prompt that can be used by all providers.
"""
# Default prompt to be used by all context providers when assembling memories/instructions
DEFAULT_CONTEXT_PROMPT: ClassVar[str] = (
"## Memories\nConsider the following memories when answering user questions:"
)
DEFAULT_CONTEXT_PROMPT: Final[str] = "## Memories\nConsider the following memories when answering user questions:"
async def thread_created(self, thread_id: str | None) -> None:
"""Called just after a new thread is created.
@@ -65,19 +80,27 @@ class ContextProvider(AFBaseModel, ABC):
"""
pass
async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
"""Called just before messages are added to the chat by any participant.
async def invoked(
self,
request_messages: ChatMessage | Sequence[ChatMessage],
response_messages: ChatMessage | Sequence[ChatMessage] | None = None,
invoke_exception: Exception | None = None,
**kwargs: Any,
) -> None:
"""Called after the agent has received a response from the underlying inference service.
Inheritors can use this method to update their context based on new messages.
You can inspect the request and response messages, and update the state of the context provider
Args:
thread_id: The ID of the thread for the new message.
new_messages: New messages to add.
request_messages: messages that were sent to the model/agent
response_messages: messages that were returned by the model/agent
invoke_exception: exception that was thrown, if any.
kwargs: not used at present.
"""
pass
@abstractmethod
async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
"""Called just before the Model/Agent/etc. is invoked.
Implementers can load any additional context required at this time,
@@ -85,6 +108,7 @@ class ContextProvider(AFBaseModel, ABC):
Args:
messages: The most recent messages that the agent is being invoked with.
kwargs: not used at present.
"""
pass
@@ -125,16 +149,16 @@ class AggregateContextProvider(ContextProvider):
It delegates events to multiple context providers and aggregates responses from those events before returning.
"""
providers: list[ContextProvider]
"""List of registered context providers."""
def __init__(self, context_providers: Sequence[ContextProvider] | None = None) -> None:
"""Initialize AggregateContextProvider with context providers.
def __init__(self, context_providers: ContextProvider | Sequence[ContextProvider] | None = None) -> None:
"""Initialize the AggregateContextProvider with context providers.
Args:
context_providers: Context providers to add.
"""
super().__init__(providers=list(context_providers or [])) # type: ignore
if isinstance(context_providers, ContextProvider):
self.providers = [context_providers]
else:
self.providers = cast(list[ContextProvider], context_providers) or []
self._exit_stack: AsyncExitStack | None = None
def add(self, context_provider: ContextProvider) -> None:
@@ -145,24 +169,44 @@ class AggregateContextProvider(ContextProvider):
"""
self.providers.append(context_provider)
@override
async def thread_created(self, thread_id: str | None = None) -> None:
await asyncio.gather(*[x.thread_created(thread_id) for x in self.providers])
async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
await asyncio.gather(*[x.messages_adding(thread_id, new_messages) for x in self.providers])
@override
async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
contexts = await asyncio.gather(*[provider.invoking(messages, **kwargs) for provider in self.providers])
instructions: str = ""
return_messages: list[ChatMessage] = []
tools: list[ToolProtocol] = []
for ctx in contexts:
if ctx.instructions:
instructions += ctx.instructions
if ctx.messages:
return_messages.extend(ctx.messages)
if ctx.tools:
tools.extend(ctx.tools)
return Context(instructions=instructions, messages=return_messages, tools=tools)
async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
sub_contexts = await asyncio.gather(*[x.model_invoking(messages) for x in self.providers])
combined_context = Context()
# Flatten the list of lists and filter out None values
all_contents = []
for ctx in sub_contexts:
if ctx.contents:
all_contents.extend(ctx.contents)
combined_context.contents = all_contents if all_contents else None
return combined_context
@override
async def invoked(
self,
request_messages: ChatMessage | Sequence[ChatMessage],
response_messages: ChatMessage | Sequence[ChatMessage] | None = None,
invoke_exception: Exception | None = None,
**kwargs: Any,
) -> None:
await asyncio.gather(*[
x.invoked(
request_messages=request_messages,
response_messages=response_messages,
invoke_exception=invoke_exception,
**kwargs,
)
for x in self.providers
])
@override
async def __aenter__(self) -> "Self":
"""Enter async context manager and set up all providers.
@@ -178,6 +222,7 @@ class AggregateContextProvider(ContextProvider):
return self
@override
async def __aexit__(
self,
exc_type: type[BaseException] | None,
python/packages/main/agent_framework/_middleware.py
+4 -2
View File
@@ -35,6 +35,7 @@ class MiddlewareType(Enum):
__all__ = [
"AgentMiddleware",
"AgentMiddlewares",
"AgentRunContext",
"ChatContext",
"ChatMiddleware",
@@ -230,6 +231,7 @@ Middleware: TypeAlias = (
| ChatMiddleware
| ChatMiddlewareCallable
)
AgentMiddlewares: TypeAlias = AgentMiddleware | AgentMiddlewareCallable
# Middleware type markers for decorators
@@ -1009,7 +1011,7 @@ def use_chat_middleware(chat_client_class: type[TChatClient]) -> type[TChatClien
pipeline = ChatMiddlewarePipeline(chat_middleware_list) # type: ignore[arg-type]
context = ChatContext(
chat_client=self,
messages=self.prepare_messages(messages),
messages=self.prepare_messages(messages, chat_options),
chat_options=chat_options,
is_streaming=False,
kwargs=kwargs,
@@ -1059,7 +1061,7 @@ def use_chat_middleware(chat_client_class: type[TChatClient]) -> type[TChatClien
pipeline = ChatMiddlewarePipeline(all_middleware) # type: ignore[arg-type]
context = ChatContext(
chat_client=self,
messages=self.prepare_messages(messages),
messages=self.prepare_messages(messages, chat_options),
chat_options=chat_options,
is_streaming=True,
kwargs=kwargs,
python/packages/main/agent_framework/_threads.py
+232 -248
View File
@@ -1,15 +1,19 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import Sequence
from typing import Any, Protocol, overload
from typing import Any, Protocol, TypeVar
from pydantic import model_validator
from ._memory import AggregateContextProvider
from ._pydantic import AFBaseModel
from ._types import ChatMessage
from .exceptions import AgentThreadException
__all__ = ["AgentThread", "ChatMessageList", "ChatMessageStore"]
__all__ = ["AgentThread", "ChatMessageStore", "ChatMessageStoreProtocol"]
class ChatMessageStore(Protocol):
class ChatMessageStoreProtocol(Protocol):
"""Defines methods for storing and retrieving chat messages associated with a specific thread.
Implementations of this protocol are responsible for managing the storage of chat messages,
@@ -24,7 +28,7 @@ class ChatMessageStore(Protocol):
If the messages stored in the store become very large, it is up to the store to
truncate, summarize or otherwise limit the number of messages returned.
When using implementations of ChatMessageStore, a new one should be created for each thread
When using implementations of ChatMessageStoreProtocol, a new one should be created for each thread
since they may contain state that is specific to a thread.
"""
...
@@ -33,66 +37,187 @@ class ChatMessageStore(Protocol):
"""Adds messages to the store."""
...
async def deserialize_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
"""Deserializes the state into the properties on this store.
@classmethod
async def deserialize(cls, serialized_store_state: Any, **kwargs: Any) -> "ChatMessageStoreProtocol":
"""Creates a new instance of the store from previously serialized state.
This method, together with serialize_state can be used to save and load messages from a persistent store
if this store only has messages in memory.
"""
...
async def serialize_state(self, **kwargs: Any) -> Any:
async def update_from_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
"""Update the current ChatMessageStore instance from serialized state data.
Args:
serialized_store_state: Previously serialized state data containing messages.
**kwargs: Additional arguments for deserialization.
"""
...
async def serialize(self, **kwargs: Any) -> Any:
"""Serializes the current object's state.
This method, together with deserialize_state can be used to save and load messages from a persistent store
This method, together with deserialize can be used to save and load messages from a persistent store
if this store only has messages in memory.
"""
...
class AgentThread(AFBaseModel):
"""Base class for agent threads."""
class AgentThreadState(AFBaseModel):
"""State model for serializing and deserializing thread information.
_service_thread_id: str | None = None
_message_store: ChatMessageStore | None = None
Attributes:
service_thread_id: Optional ID of the thread managed by the agent service.
chat_message_store_state: Optional serialized state of the chat message store.
"""
@overload
def __init__(self) -> None:
"""Initialize an empty AgentThread with no service thread ID or message store."""
...
service_thread_id: str | None = None
chat_message_store_state: Any | None = None
@overload
def __init__(self, service_thread_id: str) -> None:
"""Initialize an AgentThread with a service thread ID.
@model_validator(mode="before")
def validate_only_one(cls, values: dict[str, Any]) -> dict[str, Any]:
if (
isinstance(values, dict)
and values.get("service_thread_id") is not None
and values.get("chat_message_store_state") is not None
):
raise AgentThreadException("Only one of service_thread_id or chat_message_store_state may be set.")
return values
class ChatMessageStoreState(AFBaseModel):
"""State model for serializing and deserializing chat message store data.
Attributes:
messages: List of chat messages stored in the message store.
"""
messages: list[ChatMessage]
TChatMessageStore = TypeVar("TChatMessageStore", bound="ChatMessageStore")
class ChatMessageStore:
"""An in-memory implementation of ChatMessageStoreProtocol that stores messages in a list.
This implementation provides a simple, list-based storage for chat messages
with support for serialization and deserialization. It implements all the
required methods of the ChatMessageStoreProtocol protocol.
The store maintains messages in memory and provides methods to serialize
and deserialize the state for persistence purposes.
Args:
messages: Optional initial list of ChatMessage objects to populate the store.
"""
def __init__(self, messages: Sequence[ChatMessage] | None = None):
"""Create a ChatMessageStore for use in a thread.
Args:
service_thread_id: The ID of the thread managed by the agent service.
messages: The messages to store.
"""
...
self.messages = list(messages) if messages else []
@overload
def __init__(self, *, message_store: ChatMessageStore) -> None:
"""Initialize an AgentThread with a custom message store.
async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
"""Add messages to the store.
Args:
message_store: The ChatMessageStore implementation for managing chat messages.
messages: Sequence of ChatMessage objects to add to the store.
"""
...
self.messages.extend(messages)
def __init__(self, service_thread_id: str | None = None, *, message_store: ChatMessageStore | None = None) -> None:
"""Initialize an AgentThread.
async def list_messages(self) -> list[ChatMessage]:
"""Get all messages from the store in chronological order.
Returns:
List of ChatMessage objects, ordered from oldest to newest.
"""
return self.messages
@classmethod
async def deserialize(
cls: type[TChatMessageStore], serialized_store_state: Any, **kwargs: Any
) -> TChatMessageStore:
"""Create a new ChatMessageStore instance from serialized state data.
Args:
serialized_store_state: Previously serialized state data containing messages.
**kwargs: Additional arguments for deserialization.
Returns:
A new ChatMessageStore instance populated with messages from the serialized state.
"""
state = ChatMessageStoreState.model_validate(serialized_store_state, **kwargs)
if state.messages:
return cls(messages=state.messages)
return cls()
async def update_from_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
"""Update the current ChatMessageStore instance from serialized state data.
Args:
serialized_store_state: Previously serialized state data containing messages.
**kwargs: Additional arguments for deserialization.
"""
if not serialized_store_state:
return
state = ChatMessageStoreState.model_validate(serialized_store_state, **kwargs)
if state.messages:
self.messages = state.messages
async def serialize(self, **kwargs: Any) -> Any:
"""Serialize the current store state for persistence.
Args:
**kwargs: Additional arguments for serialization.
Returns:
Serialized state data that can be used with deserialize_state.
"""
state = ChatMessageStoreState(messages=self.messages)
return state.model_dump(**kwargs)
TAgentThread = TypeVar("TAgentThread", bound="AgentThread")
class AgentThread:
"""The Agent thread class, this can represent both a locally managed thread or a thread managed by the service."""
def __init__(
self,
*,
service_thread_id: str | None = None,
message_store: ChatMessageStoreProtocol | None = None,
context_provider: AggregateContextProvider | None = None,
) -> None:
"""Initialize an AgentThread, do not use this method manually, always use: agent.get_new_thread().
Args:
service_thread_id: Optional ID of the thread managed by the agent service.
message_store: Optional ChatMessageStore implementation for managing chat messages.
context_provider: Optional ContextProvider for the thread.
Note:
Either service_thread_id or message_store may be set, but not both.
"""
super().__init__()
if service_thread_id is not None and message_store is not None:
raise AgentThreadException("Only the service_thread_id or message_store may be set, but not both.")
self.service_thread_id = service_thread_id
self.message_store = message_store
self._service_thread_id = service_thread_id
self._message_store = message_store
self.context_provider = context_provider
@property
def is_initialized(self) -> bool:
"""Indicates if the thread is initialized.
This means either the service_thread_id or the message_store is set.
"""
return self._service_thread_id is not None or self._message_store is not None
@property
def service_thread_id(self) -> str | None:
@@ -105,39 +230,53 @@ class AgentThread(AFBaseModel):
Note that either service_thread_id or message_store may be set, but not both.
"""
if not self._service_thread_id and not service_thread_id:
if service_thread_id is None:
return
if self._message_store is not None:
raise ValueError(
raise AgentThreadException(
"Only the service_thread_id or message_store may be set, "
"but not both and switching from one to another is not supported."
)
self._service_thread_id = service_thread_id
@property
def message_store(self) -> ChatMessageStore | None:
"""Gets the ChatMessageStore used by this thread, when messages should be stored in a custom location."""
def message_store(self) -> ChatMessageStoreProtocol | None:
"""Gets the ChatMessageStoreProtocol used by this thread."""
return self._message_store
@message_store.setter
def message_store(self, message_store: ChatMessageStore | None) -> None:
"""Sets the ChatMessageStore used by this thread, when messages should be stored in a custom location.
def message_store(self, message_store: ChatMessageStoreProtocol | None) -> None:
"""Sets the ChatMessageStoreProtocol used by this thread.
Note that either service_thread_id or message_store may be set, but not both.
"""
if self._message_store is None and message_store is None:
if message_store is None:
return
if self._service_thread_id:
raise ValueError(
if self._service_thread_id is not None:
raise AgentThreadException(
"Only the service_thread_id or message_store may be set, "
"but not both and switching from one to another is not supported."
)
self._message_store = message_store
async def on_new_messages(self, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
"""Invoked when a new message has been contributed to the chat by any participant."""
if self._service_thread_id is not None:
# If the thread messages are stored in the service there is nothing to do here,
# since invoking the service should already update the thread.
return
if self._message_store is None:
# If there is no conversation id, and no store we can
# create a default in memory store.
self._message_store = ChatMessageStore()
# If a store has been provided, we need to add the messages to the store.
if isinstance(new_messages, ChatMessage):
new_messages = [new_messages]
await self._message_store.add_messages(new_messages)
async def serialize(self, **kwargs: Any) -> dict[str, Any]:
"""Serializes the current object's state.
@@ -146,226 +285,71 @@ class AgentThread(AFBaseModel):
"""
chat_message_store_state = None
if self._message_store is not None:
chat_message_store_state = await self._message_store.serialize_state(**kwargs)
chat_message_store_state = await self._message_store.serialize(**kwargs)
state = ThreadState(
state = AgentThreadState(
service_thread_id=self._service_thread_id, chat_message_store_state=chat_message_store_state
)
return state.model_dump()
async def thread_on_new_messages(thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
"""Invoked when a new message has been contributed to the chat by any participant."""
if thread.service_thread_id is not None:
# If the thread messages are stored in the service there is nothing to do here,
# since invoking the service should already update the thread.
return
if thread.message_store is None:
# If there is no conversation id, and no store we can
# create a default in memory store.
thread.message_store = ChatMessageList()
# If a store has been provided, we need to add the messages to the store.
if isinstance(new_messages, ChatMessage):
new_messages = [new_messages]
await thread.message_store.add_messages(new_messages)
async def deserialize_thread_state(
thread: AgentThread,
serialized_thread: dict[str, Any],
**kwargs: Any,
) -> None:
"""Deserializes the state from a dictionary into the thread properties."""
state = ThreadState.model_validate(serialized_thread)
if state.service_thread_id:
thread.service_thread_id = state.service_thread_id
# Since we have an ID, we should not have a chat message store and we can return here.
return
# If we don't have any ChatMessageStore state return here.
if state.chat_message_store_state is None:
return
if thread.message_store is None:
# If we don't have a chat message store yet, create an in-memory one.
thread.message_store = ChatMessageList()
await thread.message_store.deserialize_state(state.chat_message_store_state, **kwargs)
class ThreadState(AFBaseModel):
"""State model for serializing and deserializing thread information.
Attributes:
service_thread_id: Optional ID of the thread managed by the agent service.
chat_message_store_state: Optional serialized state of the chat message store.
"""
service_thread_id: str | None = None
chat_message_store_state: Any | None = None
class StoreState(AFBaseModel):
"""State model for serializing and deserializing chat message store data.
Attributes:
messages: List of chat messages stored in the message store.
"""
messages: list[ChatMessage]
class ChatMessageList:
"""An in-memory implementation of ChatMessageStore that stores messages in a list.
This implementation provides a simple, list-based storage for chat messages
with support for serialization and deserialization. It implements all the
required methods of the ChatMessageStore protocol and provides additional
list-like operations for direct message manipulation.
The store maintains messages in memory and provides methods to serialize
and deserialize the state for persistence purposes.
"""
def __init__(self, messages: Sequence[ChatMessage] | None = None) -> None:
"""Initialize the message store with optional initial messages.
@classmethod
async def deserialize(
cls: type[TAgentThread],
serialized_thread_state: dict[str, Any],
*,
message_store: ChatMessageStoreProtocol | None = None,
**kwargs: Any,
) -> TAgentThread:
"""Deserializes the state from a dictionary into a new AgentThread instance.
Args:
messages: Optional collection of initial ChatMessage objects to store.
"""
self._messages: list[ChatMessage] = []
if messages:
self._messages.extend(messages)
async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
"""Add messages to the store.
Args:
messages: Sequence of ChatMessage objects to add to the store.
"""
self._messages.extend(messages)
async def list_messages(self) -> list[ChatMessage]:
"""Get all messages from the store in chronological order.
Returns:
List of ChatMessage objects, ordered from oldest to newest.
"""
return self._messages
async def deserialize_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
"""Deserialize state data into this store instance.
Args:
serialized_store_state: Previously serialized state data containing messages.
serialized_thread_state: The serialized thread state as a dictionary.
message_store: Optional ChatMessageStoreProtocol to use for managing messages.
If not provided, a new ChatMessageStore will be created if needed.
**kwargs: Additional arguments for deserialization.
"""
if serialized_store_state:
state = StoreState.model_validate(obj=serialized_store_state, **kwargs)
if state.messages:
self._messages.extend(state.messages)
async def serialize_state(self, **kwargs: Any) -> Any:
"""Serialize the current store state for persistence.
Args:
**kwargs: Additional arguments for serialization.
Returns:
Serialized state data that can be used with deserialize_state.
A new AgentThread instance with properties set from the serialized state.
"""
state = StoreState(messages=self._messages)
return state.model_dump(**kwargs)
state = AgentThreadState.model_validate(serialized_thread_state)
def __len__(self) -> int:
"""Return the number of messages in the store.
if state.service_thread_id is not None:
return cls(service_thread_id=state.service_thread_id)
Returns:
The count of messages currently stored.
"""
return len(self._messages)
# If we don't have any ChatMessageStoreProtocol state return here.
if state.chat_message_store_state is None:
return cls()
def __getitem__(self, index: int) -> ChatMessage:
"""Get a message by index.
if message_store is not None:
try:
await message_store.update_from_state(state.chat_message_store_state, **kwargs)
except Exception as ex:
raise AgentThreadException("Failed to deserialize the provided message store.") from ex
return cls(message_store=message_store)
try:
message_store = await ChatMessageStore.deserialize(state.chat_message_store_state, **kwargs)
except Exception as ex:
raise AgentThreadException("Failed to deserialize the message store.") from ex
return cls(message_store=message_store)
Args:
index: The index of the message to retrieve.
async def update_from_thread_state(
self,
serialized_thread_state: dict[str, Any],
**kwargs: Any,
) -> None:
"""Deserializes the state from a dictionary into the thread properties."""
state = AgentThreadState.model_validate(serialized_thread_state)
Returns:
The ChatMessage at the specified index.
"""
return self._messages[index]
def __setitem__(self, index: int, item: ChatMessage) -> None:
"""Set a message at the specified index.
Args:
index: The index at which to set the message.
item: The ChatMessage to set at the specified index.
"""
self._messages[index] = item
def append(self, item: ChatMessage) -> None:
"""Append a message to the end of the store.
Args:
item: The ChatMessage to append.
"""
self._messages.append(item)
def clear(self) -> None:
"""Remove all messages from the store."""
self._messages.clear()
def index(self, item: ChatMessage) -> int:
"""Return the index of the first occurrence of the specified message.
Args:
item: The ChatMessage to find.
Returns:
The index of the first occurrence of the message.
Raises:
ValueError: If the message is not found in the store.
"""
return self._messages.index(item)
def insert(self, index: int, item: ChatMessage) -> None:
"""Insert a message at the specified index.
Args:
index: The index at which to insert the message.
item: The ChatMessage to insert.
"""
self._messages.insert(index, item)
def remove(self, item: ChatMessage) -> None:
"""Remove the first occurrence of the specified message from the store.
Args:
item: The ChatMessage to remove.
Raises:
ValueError: If the message is not found in the store.
"""
self._messages.remove(item)
def pop(self, index: int = -1) -> ChatMessage:
"""Remove and return a message at the specified index.
Args:
index: The index of the message to remove and return. Defaults to -1 (last item).
Returns:
The ChatMessage that was removed.
Raises:
IndexError: If the index is out of range.
"""
return self._messages.pop(index)
if state.service_thread_id is not None:
self.service_thread_id = state.service_thread_id
# Since we have an ID, we should not have a chat message store and we can return here.
return
# If we don't have any ChatMessageStoreProtocol state return here.
if state.chat_message_store_state is None:
return
if self.message_store is not None:
await self.message_store.update_from_state(state.chat_message_store_state, **kwargs)
# If we don't have a chat message store yet, create an in-memory one.
return
# Create the message store from the default.
self.message_store = await ChatMessageStore.deserialize(state.chat_message_store_state, **kwargs) # type: ignore
python/packages/main/agent_framework/_types.py
+4 -1
View File
@@ -1799,6 +1799,7 @@ class ChatOptions(AFBaseModel):
allow_multiple_tool_calls: bool | None = None
conversation_id: str | None = None
frequency_penalty: Annotated[float | None, Field(ge=-2.0, le=2.0)] = None
instructions: str | None = None
logit_bias: MutableMapping[str | int, float] | None = None
max_tokens: Annotated[int | None, Field(gt=0)] = None
metadata: MutableMapping[str, str] | None = None
@@ -1811,7 +1812,7 @@ class ChatOptions(AFBaseModel):
store: bool | None = None
temperature: Annotated[float | None, Field(ge=0.0, le=2.0)] = None
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | Mapping[str, Any] | None = None
tools: list[ToolProtocol | MutableMapping[str, Any]] | None = None
tools: MutableSequence[ToolProtocol | MutableMapping[str, Any]] | None = None
top_p: Annotated[float | None, Field(ge=0.0, le=1.0)] = None
user: str | None = None
@@ -1902,11 +1903,13 @@ class ChatOptions(AFBaseModel):
# tool_choice has a specialized serialize method. Save it here so we can fix it later.
tool_choice = other.tool_choice or self.tool_choice
updated_values = other.model_dump(exclude_none=True, exclude={"tools"})
logit_bias = updated_values.pop("logit_bias", {})
metadata = updated_values.pop("metadata", {})
additional_properties = updated_values.pop("additional_properties", {})
combined = self.model_copy(update=updated_values)
combined.tool_choice = tool_choice
combined.instructions = " ".join([combined.instructions or "", other.instructions or ""])
combined.logit_bias = {**(combined.logit_bias or {}), **logit_bias}
combined.metadata = {**(combined.metadata or {}), **metadata}
combined.additional_properties = {**(combined.additional_properties or {}), **additional_properties}
python/packages/main/agent_framework/_workflow/__init__.py
-2
View File
@@ -187,5 +187,3 @@ __all__ = [
with contextlib.suppress(AttributeError, TypeError, ValueError):
# Rebuild WorkflowExecutor to resolve Workflow forward reference
WorkflowExecutor.model_rebuild()
# Rebuild WorkflowAgent to resolve Workflow forward reference
WorkflowAgent.model_rebuild()
python/packages/main/agent_framework/_workflow/_agent.py
+7 -15
View File
@@ -6,7 +6,7 @@ from collections.abc import AsyncIterable, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
from pydantic import Field
from pydantic import BaseModel
from agent_framework import (
AgentRunResponse,
@@ -21,7 +21,6 @@ from agent_framework import (
UsageDetails,
)
from .._pydantic import AFBaseModel
from ..exceptions import AgentExecutionException
from ._events import (
AgentRunUpdateEvent,
@@ -41,15 +40,10 @@ class WorkflowAgent(BaseAgent):
# Class variable for the request info function name
REQUEST_INFO_FUNCTION_NAME: ClassVar[str] = "request_info"
class RequestInfoFunctionArgs(AFBaseModel):
class RequestInfoFunctionArgs(BaseModel):
request_id: str
data: Any
workflow: "Workflow" = Field(description="The workflow wrapped as an agent")
pending_requests: dict[str, RequestInfoEvent] = Field(
default_factory=dict, description="Pending request info events"
)
def __init__(
self,
workflow: "Workflow",
@@ -70,9 +64,6 @@ class WorkflowAgent(BaseAgent):
"""
if id is None:
id = f"WorkflowAgent_{uuid.uuid4().hex[:8]}"
# Initialize with standard BaseAgent parameters first
kwargs["workflow"] = workflow
# Validate the workflow's start executor can handle agent-facing message inputs
try:
start_executor = workflow.get_start_executor()
@@ -84,6 +75,9 @@ class WorkflowAgent(BaseAgent):
super().__init__(id=id, name=name, description=description, **kwargs)
self.workflow: "Workflow" = workflow
self.pending_requests: dict[str, RequestInfoEvent] = {}
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
@@ -116,8 +110,7 @@ class WorkflowAgent(BaseAgent):
response = self.merge_updates(response_updates, response_id)
# Notify thread of new messages (both input and response messages)
await self._notify_thread_of_new_messages(thread, input_messages)
await self._notify_thread_of_new_messages(thread, response.messages)
await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
return response
@@ -151,8 +144,7 @@ class WorkflowAgent(BaseAgent):
response = self.merge_updates(response_updates, response_id)
# Notify thread of new messages (both input and response messages)
await self._notify_thread_of_new_messages(thread, input_messages)
await self._notify_thread_of_new_messages(thread, response.messages)
await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
async def _run_stream_impl(
self,
python/packages/main/agent_framework/exceptions.py
+6
View File
@@ -49,6 +49,12 @@ class AgentInitializationError(AgentException):
pass
class AgentThreadException(AgentException):
"""An error occurred while managing the agent thread."""
pass
class ChatClientException(AgentFrameworkException):
"""An error occurred while dealing with a chat client."""
python/packages/main/agent_framework/openai/_assistants_client.py
+1 -1
View File
@@ -407,7 +407,7 @@ class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient):
"json_schema": chat_options.response_format.model_json_schema(),
}
instructions: list[str] = []
instructions: list[str] = [chat_options.instructions] if chat_options and chat_options.instructions else []
tool_results: list[FunctionResultContent] | None = None
additional_messages: list[AdditionalMessage] | None = None
python/packages/main/agent_framework/openai/_chat_client.py
+7 -2
View File
@@ -119,7 +119,7 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
# region content creation
def _chat_to_tool_spec(self, tools: list[ToolProtocol | MutableMapping[str, Any]]) -> list[dict[str, Any]]:
def _chat_to_tool_spec(self, tools: Sequence[ToolProtocol | MutableMapping[str, Any]]) -> list[dict[str, Any]]:
chat_tools: list[dict[str, Any]] = []
for tool in tools:
if isinstance(tool, ToolProtocol):
@@ -132,7 +132,9 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
chat_tools.append(tool if isinstance(tool, dict) else dict(tool))
return chat_tools
def _process_web_search_tool(self, tools: list[ToolProtocol | MutableMapping[str, Any]]) -> dict[str, Any] | None:
def _process_web_search_tool(
self, tools: Sequence[ToolProtocol | MutableMapping[str, Any]]
) -> dict[str, Any] | None:
for tool in tools:
if isinstance(tool, HostedWebSearchTool):
# Web search tool requires special handling
@@ -152,6 +154,9 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
def _prepare_options(self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions) -> dict[str, Any]:
# Preprocess web search tool if it exists
options_dict = chat_options.to_provider_settings()
instructions = options_dict.pop("instructions", None)
if instructions:
messages = [ChatMessage(role="system", text=instructions), *messages]
if messages and "messages" not in options_dict:
options_dict["messages"] = self._prepare_chat_history_for_request(messages)
if "messages" not in options_dict:
python/packages/main/agent_framework/openai/_responses_client.py
+3 -1
View File
@@ -172,7 +172,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
# region Prep methods
def _tools_to_response_tools(
self, tools: list[ToolProtocol | MutableMapping[str, Any]]
self, tools: Sequence[ToolProtocol | MutableMapping[str, Any]]
) -> list[ToolParam | dict[str, Any]]:
response_tools: list[ToolParam | dict[str, Any]] = []
for tool in tools:
@@ -314,6 +314,8 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
options_dict["user"] = chat_options.user
# messages
if instructions := options_dict.pop("instructions", None):
messages = [ChatMessage(role="system", text=instructions), *messages]
request_input = self._prepare_chat_messages_for_request(messages)
if not request_input:
raise ServiceInvalidRequestError("Messages are required for chat completions")