github-featured/microsoft--agent-framework
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Python: [BREAKING] PR2 — Wire context provider pipeline, remove old types, update all consumers (#3850)

Browse Source 
* PR2: Wire context provider pipeline and update all internal consumers

- Replace AgentThread with AgentSession across all packages
- Replace ContextProvider with BaseContextProvider across all packages
- Replace context_provider param with context_providers (Sequence)
- Replace thread= with session= in run() signatures
- Replace get_new_thread() with create_session()
- Add get_session(service_session_id) to agent interface
- DurableAgentThread -> DurableAgentSession
- Remove _notify_thread_of_new_messages from WorkflowAgent
- Wire before_run/after_run context provider pipeline in RawAgent
- Auto-inject InMemoryHistoryProvider when no providers configured

* fix: update all tests for context provider pipeline, fix lazy-loaders, remove old test files

* refactor: update all sample files for context provider pipeline (AgentThread→AgentSession, ContextProvider→BaseContextProvider)

* fix: update remaining ag-ui references (client docstring, getting_started sample)

* fix: make get_session service_session_id keyword-only to avoid confusion with session_id

* refactor: rename _RunContext.thread_messages to session_messages

* refactor: remove _threads.py, _memory.py, and old provider files; migrate devui to use plain message lists

* rename: remove _new_ prefix from test files

* refactor: rewrite SlidingWindowChatMessageStore as SlidingWindowHistoryProvider(InMemoryHistoryProvider)

* fix: read full history from session state directly instead of reaching into provider internals

* fix: update stale .pyi stubs, sample imports, and README references for new provider types

* fix: remove stale message_store, _notify_thread_of_new_messages, and session_id.key references in samples

* refactor: merge context_providers and sessions sample folders into sessions, remove aggregate_context_provider

* refactor: UserInfoMemory stores state in session.state instead of instance attributes

* feat: add Pydantic BaseModel support to session state serialization

Pydantic models stored in session.state are now automatically serialized
via model_dump() and restored via model_validate() during to_dict()/from_dict()
round-trips. Models are auto-registered on first serialization; use
register_state_type() for cold-start deserialization.

Also export register_state_type as a public API.

* fix mem0

* Update sample README links and descriptions for session terminology

- Replace 'thread' with 'session' in sample descriptions across all READMEs
- Update file links for renamed samples (mem0_sessions, redis_sessions, etc.)
- Fix Threads section → Sessions section in main samples/README.md
- Update tools, middleware, workflows, durabletask, azure_functions READMEs
- Update architecture diagrams in concepts/tools/README.md
- Update migration guides (autogen, semantic-kernel)

* Fix broken Redis README link to renamed sample

* Fix Mem0 OSS client search: pass scoping params as direct kwargs

AsyncMemory (OSS) expects user_id/agent_id/run_id as direct kwargs,
while AsyncMemoryClient (Platform) expects them in a filters dict.
Adds tests for both client types.

Port of fix from #3844 to new Mem0ContextProvider.

* Fix rebase issues: restore missing _conversation_state.py and checkpoint decode logic

- Add back _conversation_state.py (encode/decode_chat_messages) lost in rebase
- Fix on_checkpoint_restore to decode cache/conversation with decode_chat_messages
- Fix on_checkpoint_restore to use decode_checkpoint_value for pending requests
- Add tests/workflow/__init__.py for relative import support
- Fix test_agent_executor checkpoint selection (checkpoints[1] not superstep)

* Add STORES_BY_DEFAULT ClassVar to skip redundant InMemoryHistoryProvider injection

Chat clients that store history server-side by default (OpenAI Responses API,
Azure AI Agent) now declare STORES_BY_DEFAULT = True. The agent checks this
during auto-injection and skips InMemoryHistoryProvider unless the user
explicitly sets store=False.

* Fix broken markdown links in azure_ai and redis READMEs

* Fix getting-started samples to use session API instead of removed thread/ContextProvider API

* updates to workflow as agent

* fix group chat import

* Rename Thread→Session throughout, fix service_session_id propagation, remove stale AGUIThread

- Fix: Propagate conversation_id from ChatResponse back to session.service_session_id
  in both streaming and non-streaming paths in _agents.py
- Rename AgentThreadException → AgentSessionException
- Remove stale AGUIThread from ag_ui lazy-loader
- Rename use_service_thread → use_service_session in ag-ui package
- Rename test functions from *_thread_* to *_session_*
- Rename sample files from *_thread* to *_session*
- Update docstrings and comments: thread → session
- Update _mcp.py kwargs filter: add 'session' alongside 'thread'
- Fix ContinuationToken docstring example: thread=thread → session=session
- Fix _clients.py docstring: 'Agent threads' → 'Agent sessions'

* Fix broken markdown links after thread→session file renames

* fix azure ai test
This commit is contained in:
Eduard van Valkenburg
2026-02-12 22:00:32 +01:00
committed by GitHub
parent 0c67dbbce5
commit 1e350ea22f
312 changed files with 6669 additions and 11423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/a2a/agent_framework_a2a/_agent.py
+5 -5
View File
@@ -31,7 +31,7 @@ from a2a.types import Role as A2ARole
from agent_framework import (
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
Content,
ContinuationToken,
@@ -211,7 +211,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
continuation_token: A2AContinuationToken | None = None,
background: bool = False,
**kwargs: Any,
@@ -223,7 +223,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
continuation_token: A2AContinuationToken | None = None,
background: bool = False,
**kwargs: Any,
@@ -234,7 +234,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
continuation_token: A2AContinuationToken | None = None,
background: bool = False,
**kwargs: Any,
@@ -246,7 +246,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
Keyword Args:
stream: Whether to stream the response. Defaults to False.
thread: The conversation thread associated with the message(s).
session: The conversation session associated with the message(s).
continuation_token: Optional token to resume a long-running task
instead of starting a new one.
background: When True, in-progress task updates surface continuation
python/packages/ag-ui/agent_framework_ag_ui/_agent.py
+6 -6
View File
@@ -18,7 +18,7 @@ class AgentConfig:
self,
state_schema: Any | None = None,
predict_state_config: dict[str, dict[str, str]] | None = None,
use_service_thread: bool = False,
use_service_session: bool = False,
require_confirmation: bool = True,
):
"""Initialize agent configuration.
@@ -26,12 +26,12 @@ class AgentConfig:
Args:
state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
predict_state_config: Configuration for predictive state updates
use_service_thread: Whether the agent thread is service-managed
use_service_session: Whether the agent session is service-managed
require_confirmation: Whether predictive updates require user confirmation before applying
"""
self.state_schema = self._normalize_state_schema(state_schema)
self.predict_state_config = predict_state_config or {}
self.use_service_thread = use_service_thread
self.use_service_session = use_service_session
self.require_confirmation = require_confirmation
@staticmethod
@@ -77,7 +77,7 @@ class AgentFrameworkAgent:
state_schema: Any | None = None,
predict_state_config: dict[str, dict[str, str]] | None = None,
require_confirmation: bool = True,
use_service_thread: bool = False,
use_service_session: bool = False,
):
"""Initialize the AG-UI compatible agent wrapper.
@@ -88,7 +88,7 @@ class AgentFrameworkAgent:
state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
predict_state_config: Configuration for predictive state updates
require_confirmation: Whether predictive updates require user confirmation before applying
use_service_thread: Whether the agent thread is service-managed
use_service_session: Whether the agent session is service-managed
"""
self.agent = agent
self.name = name or getattr(agent, "name", "agent")
@@ -97,7 +97,7 @@ class AgentFrameworkAgent:
self.config = AgentConfig(
state_schema=state_schema,
predict_state_config=predict_state_config,
use_service_thread=use_service_thread,
use_service_session=use_service_session,
require_confirmation=require_confirmation,
)
python/packages/ag-ui/agent_framework_ag_ui/_client.py
+3 -3
View File
@@ -171,11 +171,11 @@ class AGUIChatClient(
client = AGUIChatClient(endpoint="http://localhost:8888/")
agent = Agent(name="assistant", client=client)
thread = await agent.get_new_thread()
session = agent.create_session()
# Agent automatically maintains history and sends full context
response = await agent.run("Hello!", thread=thread)
response2 = await agent.run("How are you?", thread=thread)
response = await agent.run("Hello!", session=session)
response2 = await agent.run("How are you?", session=session)
Streaming usage:
python/packages/ag-ui/agent_framework_ag_ui/_run.py
+8 -8
View File
@@ -27,7 +27,7 @@ from ag_ui.core import (
ToolCallStartEvent,
)
from agent_framework import (
AgentThread,
AgentSession,
Content,
Message,
SupportsAgentRun,
@@ -809,12 +809,12 @@ async def run_agent_stream(
register_additional_client_tools(agent, client_tools)
tools = merge_tools(server_tools, client_tools)
# Create thread (with service thread support)
if config.use_service_thread:
# Create session (with service session support)
if config.use_service_session:
supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId")
thread = AgentThread(service_thread_id=supplied_thread_id)
session = AgentSession(service_session_id=supplied_thread_id)
else:
thread = AgentThread()
session = AgentSession()
# Inject metadata for AG-UI orchestration (Feature #2: Azure-safe truncation)
base_metadata: dict[str, Any] = {
@@ -823,16 +823,16 @@ async def run_agent_stream(
}
if flow.current_state:
base_metadata["current_state"] = flow.current_state
thread.metadata = _build_safe_metadata(base_metadata) # type: ignore[attr-defined]
session.metadata = _build_safe_metadata(base_metadata) # type: ignore[attr-defined]
# Build run kwargs (Feature #6: Azure store flag when metadata present)
run_kwargs: dict[str, Any] = {"thread": thread}
run_kwargs: dict[str, Any] = {"session": session}
if tools:
run_kwargs["tools"] = tools
# Filter out AG-UI internal metadata keys before passing to chat client
# These are used internally for orchestration and should not be sent to the LLM provider
client_metadata = {
k: v for k, v in (getattr(thread, "metadata", None) or {}).items() if k not in AG_UI_INTERNAL_METADATA_KEYS
k: v for k, v in (getattr(session, "metadata", None) or {}).items() if k not in AG_UI_INTERNAL_METADATA_KEYS
}
safe_metadata = _build_safe_metadata(client_metadata) if client_metadata else {}
if safe_metadata:
python/packages/ag-ui/getting_started/client_with_agent.py
+17 -62
View File
@@ -4,10 +4,10 @@
This demonstrates the HYBRID pattern matching .NET AGUIClient implementation:
1. AgentThread Pattern (like .NET):
- Create thread with agent.get_new_thread()
- Pass thread to agent.run(stream=True) on each turn
- Thread automatically maintains conversation history via message_store
1. AgentSession Pattern (like .NET):
- Create session with agent.create_session()
- Pass session to agent.run(stream=True) on each turn
- Session maintains conversation context via context providers
2. Hybrid Tool Execution:
- AGUIChatClient uses function invocation mixin
@@ -15,7 +15,7 @@ This demonstrates the HYBRID pattern matching .NET AGUIClient implementation:
- Server may also have its own tools that execute server-side
- Both work together: server LLM decides which tool to call, decorator handles client execution
This matches .NET pattern: thread maintains state, tools execute on appropriate side.
This matches .NET pattern: session maintains state, tools execute on appropriate side.
"""
from __future__ import annotations
@@ -59,13 +59,13 @@ async def main():
This matches the .NET pattern from Program.cs where:
- AIAgent agent = chatClient.CreateAIAgent(tools: [...])
- AgentThread thread = agent.GetNewThread()
- RunStreamingAsync(messages, thread)
- AgentSession session = agent.CreateSession()
- RunStreamingAsync(messages, session)
Python equivalent:
- agent = Agent(client=AGUIChatClient(...), tools=[...])
- thread = agent.get_new_thread() # Creates thread with message_store
- agent.run(message, stream=True, thread=thread) # Thread accumulates history
- session = agent.create_session() # Creates session
- agent.run(message, stream=True, session=session) # Session tracks context
"""
server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
@@ -74,7 +74,7 @@ async def main():
print("=" * 70)
print(f"\nServer: {server_url}")
print("\nThis example demonstrates:")
print(" 1. AgentThread maintains conversation state (like .NET)")
print(" 1. AgentSession maintains conversation state (like .NET)")
print(" 2. Client-side tools execute locally via function invocation mixin")
print(" 3. Server may have additional tools that execute server-side")
print(" 4. HYBRID: Client and server tools work together simultaneously\n")
@@ -90,8 +90,8 @@ async def main():
tools=[get_weather],
)
# Create a thread to maintain conversation state (like .NET AgentThread)
thread = agent.get_new_thread()
# Create a session to maintain conversation state (like .NET AgentSession)
session = agent.create_session()
print("=" * 70)
print("CONVERSATION WITH HISTORY")
@@ -99,21 +99,21 @@ async def main():
# Turn 1: Introduce
print("\nUser: My name is Alice and I live in Seattle\n")
async for chunk in agent.run("My name is Alice and I live in Seattle", stream=True, thread=thread):
async for chunk in agent.run("My name is Alice and I live in Seattle", stream=True, session=session):
if chunk.text:
print(chunk.text, end="", flush=True)
print("\n")
# Turn 2: Ask about name (tests history)
print("User: What's my name?\n")
async for chunk in agent.run("What's my name?", stream=True, thread=thread):
async for chunk in agent.run("What's my name?", stream=True, session=session):
if chunk.text:
print(chunk.text, end="", flush=True)
print("\n")
# Turn 3: Ask about location (tests history)
print("User: Where do I live?\n")
async for chunk in agent.run("Where do I live?", stream=True, thread=thread):
async for chunk in agent.run("Where do I live?", stream=True, session=session):
if chunk.text:
print(chunk.text, end="", flush=True)
print("\n")
@@ -123,7 +123,7 @@ async def main():
async for chunk in agent.run(
"What's the weather forecast for today in Seattle?",
stream=True,
thread=thread,
session=session,
):
if chunk.text:
print(chunk.text, end="", flush=True)
@@ -131,56 +131,11 @@ async def main():
# Turn 5: Test server-side tool (get_time_zone is server-side only)
print("User: What time zone is Seattle in?\n")
async for chunk in agent.run("What time zone is Seattle in?", stream=True, thread=thread):
async for chunk in agent.run("What time zone is Seattle in?", stream=True, session=session):
if chunk.text:
print(chunk.text, end="", flush=True)
print("\n")
# Show thread state
if thread.message_store:
def _preview_for_message(m) -> str:
# Prefer plain text when present
if getattr(m, "text", ""):
t = m.text
return (t[:60] + "...") if len(t) > 60 else t
# Build from contents when no direct text
parts: list[str] = []
for c in getattr(m, "contents", []) or []:
content_type = getattr(c, "type", None)
if content_type == "function_call":
args = getattr(c, "arguments", None)
if isinstance(args, dict):
try:
import json as _json
args_str = _json.dumps(args)
except Exception:
args_str = str(args)
else:
args_str = str(args or "{}")
parts.append(f"tool_call {getattr(c, 'name', '?')} {args_str}")
elif content_type == "function_result":
call_id = getattr(c, "call_id", "?")
result = getattr(c, "result", None)
parts.append(f"tool_result[{call_id}]: {str(result)[:40]}")
elif content_type == "text":
text = getattr(c, "text", None)
if text:
parts.append(text)
else:
typename = getattr(c, "type", c.__class__.__name__)
parts.append(f"<{typename}>")
preview = " | ".join(parts) if parts else ""
return (preview[:60] + "...") if len(preview) > 60 else preview
messages = await thread.message_store.list_messages()
print(f"\n[THREAD STATE] {len(messages)} messages in thread's message_store")
for i, msg in enumerate(messages[-6:], 1): # Show last 6
role = msg.role if hasattr(msg.role, "value") else str(msg.role)
text_preview = _preview_for_message(msg)
print(f" {i}. [{role}]: {text_preview}")
except ConnectionError as e:
print(f"\n\033[91mConnection Error: {e}\033[0m")
print("\nMake sure an AG-UI server is running at the specified endpoint.")
python/packages/ag-ui/tests/ag_ui/conftest.py
+10 -10
View File
@@ -11,7 +11,7 @@ import pytest
from agent_framework import (
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseChatClient,
ChatOptions,
ChatResponse,
@@ -49,8 +49,8 @@ class StreamingChatClientStub(
super().__init__(function_middleware=[])
self._stream_fn = stream_fn
self._response_fn = response_fn
self.last_thread: AgentThread | None = None
self.last_service_thread_id: str | None = None
self.last_session: AgentSession | None = None
self.last_service_session_id: str | None = None
@overload
def get_response(
@@ -90,8 +90,8 @@ class StreamingChatClientStub(
options: OptionsCoT | ChatOptions[Any] | None = None,
**kwargs: Any,
) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
self.last_thread = kwargs.get("thread")
self.last_service_thread_id = self.last_thread.service_thread_id if self.last_thread else None
self.last_session = kwargs.get("session")
self.last_service_session_id = self.last_session.service_session_id if self.last_session else None
return cast(
Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
super().get_response(
@@ -178,7 +178,7 @@ class StubAgent(SupportsAgentRun):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@@ -188,7 +188,7 @@ class StubAgent(SupportsAgentRun):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
@@ -197,7 +197,7 @@ class StubAgent(SupportsAgentRun):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
if stream:
@@ -218,8 +218,8 @@ class StubAgent(SupportsAgentRun):
return _get_response()
def get_new_thread(self, **kwargs: Any) -> AgentThread:
return AgentThread()
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
# Fixtures
python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py
+16 -30
View File
@@ -444,13 +444,7 @@ async def test_thread_metadata_tracking(streaming_chat_client_stub):
async for event in wrapper.run_agent(input_data):
events.append(event)
# AG-UI internal metadata should be stored in thread.metadata
thread = agent.client.last_thread
thread_metadata = thread.metadata if thread and hasattr(thread, "metadata") else {}
assert thread_metadata.get("ag_ui_thread_id") == "test_thread_123"
assert thread_metadata.get("ag_ui_run_id") == "test_run_456"
# Internal metadata should NOT be passed to chat client options
# AG-UI internal metadata should NOT be passed to chat client options
options_metadata = captured_options.get("metadata", {})
assert "ag_ui_thread_id" not in options_metadata
assert "ag_ui_run_id" not in options_metadata
@@ -488,15 +482,7 @@ async def test_state_context_injection(streaming_chat_client_stub):
async for event in wrapper.run_agent(input_data):
events.append(event)
# Current state should be stored in thread.metadata
thread = agent.client.last_thread
thread_metadata = thread.metadata if thread and hasattr(thread, "metadata") else {}
current_state = thread_metadata.get("current_state")
if isinstance(current_state, str):
current_state = json.loads(current_state)
assert current_state == {"document": "Test content"}
# Internal metadata should NOT be passed to chat client options
# Current state should NOT be passed to chat client options
options_metadata = captured_options.get("metadata", {})
assert "current_state" not in options_metadata
@@ -611,11 +597,11 @@ async def test_json_decode_error_in_tool_result(streaming_chat_client_stub):
assert len(tool_events) == 0
async def test_agent_with_use_service_thread_is_false(streaming_chat_client_stub):
"""Test that when use_service_thread is False, the AgentThread used to run the agent is NOT set to the service thread ID."""
async def test_agent_with_use_service_session_is_false(streaming_chat_client_stub):
"""Test that when use_service_session is False, the AgentSession used to run the agent is NOT set to the service session ID."""
from agent_framework.ag_ui import AgentFrameworkAgent
request_service_thread_id: str | None = None
request_service_session_id: str | None = None
async def stream_fn(
messages: MutableSequence[Message], chat_options: ChatOptions, **kwargs: Any
@@ -625,42 +611,42 @@ async def test_agent_with_use_service_thread_is_false(streaming_chat_client_stub
)
agent = Agent(client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=False)
wrapper = AgentFrameworkAgent(agent=agent, use_service_session=False)
input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
assert request_service_thread_id is None # type: ignore[attr-defined] (service_thread_id should be set)
assert request_service_session_id is None # type: ignore[attr-defined] (service_session_id should be set)
async def test_agent_with_use_service_thread_is_true(streaming_chat_client_stub):
"""Test that when use_service_thread is True, the AgentThread used to run the agent is set to the service thread ID."""
async def test_agent_with_use_service_session_is_true(streaming_chat_client_stub):
"""Test that when use_service_session is True, the AgentSession used to run the agent is set to the service session ID."""
from agent_framework.ag_ui import AgentFrameworkAgent
request_service_thread_id: str | None = None
request_service_session_id: str | None = None
async def stream_fn(
messages: MutableSequence[Message], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
nonlocal request_service_thread_id
thread = kwargs.get("thread")
request_service_thread_id = thread.service_thread_id if thread else None
nonlocal request_service_session_id
session = kwargs.get("session")
request_service_session_id = session.service_session_id if session else None
yield ChatResponseUpdate(
contents=[Content.from_text(text="Response")], response_id="resp_67890", conversation_id="conv_12345"
)
agent = Agent(client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=True)
wrapper = AgentFrameworkAgent(agent=agent, use_service_session=True)
input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
request_service_thread_id = agent.client.last_service_thread_id
assert request_service_thread_id == "conv_123456" # type: ignore[attr-defined] (service_thread_id should be set)
request_service_session_id = agent.client.last_service_session_id
assert request_service_session_id == "conv_123456" # type: ignore[attr-defined] (service_session_id should be set)
async def test_function_approval_mode_executes_tool(streaming_chat_client_stub):
python/packages/azure-ai-search/agent_framework_azure_ai_search/__init__.py
+1 -3
View File
@@ -2,8 +2,7 @@
import importlib.metadata
from ._context_provider import _AzureAISearchContextProvider
from ._search_provider import AzureAISearchContextProvider, AzureAISearchSettings
from ._context_provider import AzureAISearchContextProvider, AzureAISearchSettings
try:
__version__ = importlib.metadata.version(__name__)
@@ -13,6 +12,5 @@ except importlib.metadata.PackageNotFoundError:
__all__ = [
"AzureAISearchContextProvider",
"AzureAISearchSettings",
"_AzureAISearchContextProvider",
"__version__",
]
python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py
+40 -38
View File
@@ -2,21 +2,20 @@
"""New-pattern Azure AI Search context provider using BaseContextProvider.
This module provides ``_AzureAISearchContextProvider``, a side-by-side implementation of
:class:`AzureAISearchContextProvider` built on the new :class:`BaseContextProvider` hooks
pattern. It will replace the existing class in PR2.
This module provides ``AzureAISearchContextProvider``, built on the new
:class:`BaseContextProvider` hooks pattern.
"""
from __future__ import annotations
import sys
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, ClassVar, Literal
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Message
from agent_framework._logging import get_logger
from agent_framework._sessions import AgentSession, BaseContextProvider, SessionContext
from agent_framework._settings import load_settings
from agent_framework._settings import SecretString, load_settings
from agent_framework.exceptions import ServiceInitializationError
from azure.core.credentials import AzureKeyCredential
from azure.core.credentials_async import AsyncTokenCredential
@@ -43,8 +42,6 @@ from azure.search.documents.models import (
VectorizedQuery,
)
from ._search_provider import AzureAISearchSettings
if TYPE_CHECKING:
from agent_framework._agents import SupportsAgentRun
from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient
@@ -111,16 +108,34 @@ logger = get_logger(__name__)
_DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT = 10
class _AzureAISearchContextProvider(BaseContextProvider):
class AzureAISearchSettings(TypedDict, total=False):
"""Settings for Azure AI Search Context Provider with auto-loading from environment.
The settings are first loaded from environment variables with the prefix 'AZURE_SEARCH_'.
If the environment variables are not found, the settings can be loaded from a .env file.
Keys:
endpoint: Azure AI Search endpoint URL.
Can be set via environment variable AZURE_SEARCH_ENDPOINT.
index_name: Name of the search index.
Can be set via environment variable AZURE_SEARCH_INDEX_NAME.
knowledge_base_name: Name of an existing Knowledge Base (for agentic mode).
Can be set via environment variable AZURE_SEARCH_KNOWLEDGE_BASE_NAME.
api_key: API key for authentication (optional, use managed identity if not provided).
Can be set via environment variable AZURE_SEARCH_API_KEY.
"""
endpoint: str | None
index_name: str | None
knowledge_base_name: str | None
api_key: SecretString | None
class AzureAISearchContextProvider(BaseContextProvider):
"""Azure AI Search context provider using the new BaseContextProvider hooks pattern.
Retrieves relevant context from Azure AI Search using semantic or agentic search
modes. This is the new-pattern equivalent of :class:`AzureAISearchContextProvider`.
Note:
This class uses a temporary ``_`` prefix to coexist with the existing
:class:`AzureAISearchContextProvider`. It will replace the existing class
in PR2.
modes.
"""
_DEFAULT_SEARCH_CONTEXT_PROMPT: ClassVar[str] = "Use the following context to answer the question:"
@@ -179,10 +194,18 @@ class _AzureAISearchContextProvider(BaseContextProvider):
"""
super().__init__(source_id)
# Determine which fields are required based on mode
required: list[str | tuple[str, ...]] = ["endpoint"]
if mode == "semantic":
required.append("index_name")
elif mode == "agentic":
required.append(("index_name", "knowledge_base_name"))
# Load settings from environment/file
settings = load_settings(
AzureAISearchSettings,
env_prefix="AZURE_SEARCH_",
required_fields=required,
endpoint=endpoint,
index_name=index_name,
knowledge_base_name=knowledge_base_name,
@@ -191,32 +214,11 @@ class _AzureAISearchContextProvider(BaseContextProvider):
env_file_encoding=env_file_encoding,
)
if not settings.get("endpoint"):
if mode == "agentic" and settings.get("index_name") and not model_deployment_name:
raise ServiceInitializationError(
"Azure AI Search endpoint is required. Set via 'endpoint' parameter "
"or 'AZURE_SEARCH_ENDPOINT' environment variable."
"model_deployment_name is required for agentic mode when creating Knowledge Base from index."
)
if mode == "semantic":
if not settings.get("index_name"):
raise ServiceInitializationError(
"Azure AI Search index name is required for semantic mode. "
"Set via 'index_name' parameter or 'AZURE_SEARCH_INDEX_NAME' environment variable."
)
elif mode == "agentic":
if settings.get("index_name") and settings.get("knowledge_base_name"):
raise ServiceInitializationError(
"For agentic mode, provide either 'index_name' OR 'knowledge_base_name', not both."
)
if not settings.get("index_name") and not settings.get("knowledge_base_name"):
raise ServiceInitializationError(
"For agentic mode, provide either 'index_name' or 'knowledge_base_name'."
)
if settings.get("index_name") and not model_deployment_name:
raise ServiceInitializationError(
"model_deployment_name is required for agentic mode when creating Knowledge Base from index."
)
resolved_credential: AzureKeyCredential | AsyncTokenCredential
if credential:
resolved_credential = credential
@@ -621,4 +623,4 @@ class _AzureAISearchContextProvider(BaseContextProvider):
return text
__all__ = ["_AzureAISearchContextProvider"]
__all__ = ["AzureAISearchContextProvider"]
python/packages/azure-ai-search/agent_framework_azure_ai_search/_search_provider.py
-991
View File
@@ -1,991 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import sys
from collections.abc import Awaitable, Callable, MutableSequence
from typing import TYPE_CHECKING, Any, Literal
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Context, ContextProvider, Message
from agent_framework._logging import get_logger
from agent_framework._settings import SecretString, load_settings
from agent_framework.exceptions import ServiceInitializationError
from azure.core.credentials import AzureKeyCredential
from azure.core.credentials_async import AsyncTokenCredential
from azure.core.exceptions import ResourceNotFoundError
from azure.search.documents.aio import SearchClient
from azure.search.documents.indexes.aio import SearchIndexClient
from azure.search.documents.indexes.models import (
AzureOpenAIVectorizerParameters,
KnowledgeBase,
KnowledgeBaseAzureOpenAIModel,
KnowledgeRetrievalLowReasoningEffort,
KnowledgeRetrievalMediumReasoningEffort,
KnowledgeRetrievalMinimalReasoningEffort,
KnowledgeRetrievalOutputMode,
KnowledgeRetrievalReasoningEffort,
KnowledgeSourceReference,
SearchIndexKnowledgeSource,
SearchIndexKnowledgeSourceParameters,
)
from azure.search.documents.models import (
QueryCaptionType,
QueryType,
VectorizableTextQuery,
VectorizedQuery,
)
# Type checking imports for optional agentic mode dependencies
if TYPE_CHECKING:
from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient
from azure.search.documents.knowledgebases.models import (
KnowledgeBaseMessage,
KnowledgeBaseMessageTextContent,
KnowledgeBaseRetrievalRequest,
KnowledgeRetrievalIntent,
KnowledgeRetrievalSemanticIntent,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalLowReasoningEffort as KBRetrievalLowReasoningEffort,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalMediumReasoningEffort as KBRetrievalMediumReasoningEffort,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalMinimalReasoningEffort as KBRetrievalMinimalReasoningEffort,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalOutputMode as KBRetrievalOutputMode,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalReasoningEffort as KBRetrievalReasoningEffort,
)
# Runtime imports for agentic mode (optional dependency)
try:
from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient
from azure.search.documents.knowledgebases.models import (
KnowledgeBaseMessage,
KnowledgeBaseMessageTextContent,
KnowledgeBaseRetrievalRequest,
KnowledgeRetrievalIntent,
KnowledgeRetrievalSemanticIntent,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalLowReasoningEffort as KBRetrievalLowReasoningEffort,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalMediumReasoningEffort as KBRetrievalMediumReasoningEffort,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalMinimalReasoningEffort as KBRetrievalMinimalReasoningEffort,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalOutputMode as KBRetrievalOutputMode,
)
from azure.search.documents.knowledgebases.models import (
KnowledgeRetrievalReasoningEffort as KBRetrievalReasoningEffort,
)
_agentic_retrieval_available = True
except ImportError:
_agentic_retrieval_available = False
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, TypedDict # pragma: no cover
else:
from typing_extensions import Self, TypedDict # pragma: no cover
"""Azure AI Search Context Provider for Agent Framework.
This module provides context providers for Azure AI Search integration with two modes:
- Agentic: Recommended for most scenarios. Uses Knowledge Bases for query planning and
multi-hop reasoning. Slightly slower with more token consumption, but more accurate.
- Semantic: Fast hybrid search (vector + keyword) with semantic ranker. Best for simple
queries where speed is critical.
See: https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/foundry-iq-boost-response-relevance-by-36-with-agentic-retrieval/4470720
"""
# Module-level constants
logger = get_logger("agent_framework.azure")
_DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT = 10
class AzureAISearchSettings(TypedDict, total=False):
"""Settings for Azure AI Search Context Provider with auto-loading from environment.
The settings are first loaded from environment variables with the prefix 'AZURE_SEARCH_'.
If the environment variables are not found, the settings can be loaded from a .env file.
Keyword Args:
endpoint: Azure AI Search endpoint URL.
Can be set via environment variable AZURE_SEARCH_ENDPOINT.
index_name: Name of the search index.
Can be set via environment variable AZURE_SEARCH_INDEX_NAME.
knowledge_base_name: Name of an existing Knowledge Base (for agentic mode).
Can be set via environment variable AZURE_SEARCH_KNOWLEDGE_BASE_NAME.
api_key: API key for authentication (optional, use managed identity if not provided).
Can be set via environment variable AZURE_SEARCH_API_KEY.
env_file_path: If provided, the .env settings are read from this file path location.
env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
Examples:
.. code-block:: python
from agent_framework_aisearch import AzureAISearchSettings
# Using environment variables
# Set AZURE_SEARCH_ENDPOINT=https://mysearch.search.windows.net
# Set AZURE_SEARCH_INDEX_NAME=my-index
settings = AzureAISearchSettings()
# Or passing parameters directly
settings = AzureAISearchSettings(
endpoint="https://mysearch.search.windows.net",
index_name="my-index",
)
# Or loading from a .env file
settings = AzureAISearchSettings(env_file_path="path/to/.env")
"""
endpoint: str | None
index_name: str | None
knowledge_base_name: str | None
api_key: SecretString | None
class AzureAISearchContextProvider(ContextProvider):
"""Azure AI Search Context Provider with hybrid search and semantic ranking.
This provider retrieves relevant documents from Azure AI Search to provide context
to the AI agent. It supports two modes:
- **agentic**: Recommended for most scenarios. Uses Knowledge Bases for query planning
and multi-hop reasoning. Slightly slower with more token consumption, but provides
more accurate results (up to 36% improvement in response relevance).
- **semantic** (default): Fast hybrid search combining vector and keyword search
with semantic reranking. Best for simple queries where speed is critical.
Examples:
Using environment variables (recommended):
.. code-block:: python
from agent_framework_aisearch import AzureAISearchContextProvider
from azure.identity.aio import DefaultAzureCredential
# Set AZURE_SEARCH_ENDPOINT and AZURE_SEARCH_INDEX_NAME in environment
search_provider = AzureAISearchContextProvider(credential=DefaultAzureCredential())
Semantic hybrid search with API key:
.. code-block:: python
# Direct API key string
search_provider = AzureAISearchContextProvider(
endpoint="https://mysearch.search.windows.net",
index_name="my-index",
api_key="my-api-key",
mode="semantic",
)
Loading from .env file:
.. code-block:: python
# Load settings from a .env file
search_provider = AzureAISearchContextProvider(
credential=DefaultAzureCredential(), env_file_path="path/to/.env"
)
Agentic retrieval for complex queries:
.. code-block:: python
# Use agentic mode for multi-hop reasoning
# Note: azure_openai_resource_url is the OpenAI endpoint for Knowledge Base model calls,
# which is different from azure_ai_project_endpoint (the AI Foundry project endpoint)
search_provider = AzureAISearchContextProvider(
endpoint="https://mysearch.search.windows.net",
index_name="my-index",
credential=DefaultAzureCredential(),
mode="agentic",
azure_openai_resource_url="https://myresource.openai.azure.com",
model_deployment_name="gpt-4o",
knowledge_base_name="my-knowledge-base",
)
"""
_DEFAULT_SEARCH_CONTEXT_PROMPT = "Use the following context to answer the question:"
def __init__(
self,
endpoint: str | None = None,
index_name: str | None = None,
api_key: str | AzureKeyCredential | None = None,
credential: AsyncTokenCredential | None = None,
*,
mode: Literal["semantic", "agentic"] = "semantic",
top_k: int = 5,
semantic_configuration_name: str | None = None,
vector_field_name: str | None = None,
embedding_function: Callable[[str], Awaitable[list[float]]] | None = None,
context_prompt: str | None = None,
# Agentic mode parameters (Knowledge Base)
azure_openai_resource_url: str | None = None,
model_deployment_name: str | None = None,
model_name: str | None = None,
knowledge_base_name: str | None = None,
retrieval_instructions: str | None = None,
azure_openai_api_key: str | None = None,
knowledge_base_output_mode: Literal["extractive_data", "answer_synthesis"] = "extractive_data",
retrieval_reasoning_effort: Literal["minimal", "medium", "low"] = "minimal",
agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize Azure AI Search Context Provider.
Args:
endpoint: Azure AI Search endpoint URL.
Can also be set via environment variable AZURE_SEARCH_ENDPOINT.
index_name: Name of the search index to query.
Can also be set via environment variable AZURE_SEARCH_INDEX_NAME.
api_key: API key for authentication (string or AzureKeyCredential).
Can also be set via environment variable AZURE_SEARCH_API_KEY.
credential: AsyncTokenCredential for managed identity authentication.
Use this for Entra ID authentication instead of api_key.
mode: Search mode - "semantic" for hybrid search with semantic ranking (fast)
or "agentic" for multi-hop reasoning (slower). Default: "semantic".
top_k: Maximum number of documents to retrieve. Only applies to semantic mode.
In agentic mode, the server-side Knowledge Base determines retrieval based on
query complexity and reasoning effort. Default: 5.
semantic_configuration_name: Name of semantic configuration in the index.
Required for semantic ranking. If None, uses index default.
vector_field_name: Name of the vector field in the index for hybrid search.
Required if using vector search. Default: None (keyword search only).
embedding_function: Async function to generate embeddings for vector search.
Signature: async def embed(text: str) -> list[float]
Required if vector_field_name is specified and no server-side vectorization.
context_prompt: Custom prompt to prepend to retrieved context.
Default: "Use the following context to answer the question:"
azure_openai_resource_url: Azure OpenAI resource URL for Knowledge Base model calls.
Required when using agentic mode with index_name (to auto-create Knowledge Base).
Not required when using an existing knowledge_base_name.
Example: "https://myresource.openai.azure.com"
model_deployment_name: Model deployment name in Azure OpenAI for Knowledge Base.
Required when using agentic mode with index_name (to auto-create Knowledge Base).
Not required when using an existing knowledge_base_name.
model_name: The underlying model name (e.g., "gpt-4o", "gpt-4o-mini").
If not provided, defaults to model_deployment_name. Used for Knowledge Base configuration.
knowledge_base_name: Name of an existing Knowledge Base to use.
Required for agentic mode if not providing index_name.
Supports KBs with any source type (web, blob, index, etc.).
retrieval_instructions: Custom instructions for the Knowledge Base's
retrieval planning. Only used in agentic mode.
azure_openai_api_key: Azure OpenAI API key for Knowledge Base to call the model.
Only needed when using API key authentication instead of managed identity.
knowledge_base_output_mode: Output mode for Knowledge Base retrieval. Only used in agentic mode.
"extractive_data": Returns raw chunks without synthesis (default, recommended for agent integration).
"answer_synthesis": Returns synthesized answer from the LLM.
Some knowledge sources require answer_synthesis mode. Default: "extractive_data".
retrieval_reasoning_effort: Reasoning effort for Knowledge Base query planning. Only used in agentic mode.
"minimal": Fastest, basic query planning.
"medium": Moderate reasoning with some query decomposition.
"low": Lower reasoning effort than medium.
Default: "minimal".
agentic_message_history_count: Number of recent messages from conversation history to send to
the Knowledge Base. This context helps with query planning in agentic mode, allowing the
Knowledge Base to understand the conversation flow and generate better retrieval queries.
There is no technical limit - adjust based on your use case. Default: 10.
env_file_path: Path to environment file for loading settings.
env_file_encoding: Encoding of the environment file.
Examples:
.. code-block:: python
from agent_framework_aisearch import AzureAISearchContextProvider
from azure.identity.aio import DefaultAzureCredential
# Using environment variables
# Set AZURE_SEARCH_ENDPOINT=https://mysearch.search.windows.net
# Set AZURE_SEARCH_INDEX_NAME=my-index
credential = DefaultAzureCredential()
provider = AzureAISearchContextProvider(credential=credential)
# Or passing parameters directly
provider = AzureAISearchContextProvider(
endpoint="https://mysearch.search.windows.net",
index_name="my-index",
credential=credential,
)
# Or loading from a .env file
provider = AzureAISearchContextProvider(credential=credential, env_file_path="path/to/.env")
"""
# Load settings from environment/file
settings = load_settings(
AzureAISearchSettings,
env_prefix="AZURE_SEARCH_",
endpoint=endpoint,
index_name=index_name,
knowledge_base_name=knowledge_base_name,
api_key=api_key if isinstance(api_key, str) else None,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
# Validate required parameters
if not settings.get("endpoint"):
raise ServiceInitializationError(
"Azure AI Search endpoint is required. Set via 'endpoint' parameter "
"or 'AZURE_SEARCH_ENDPOINT' environment variable."
)
# Validate index_name and knowledge_base_name based on mode
# Note: settings["field"] / settings.get("field") contains the resolved value (explicit param OR env var)
if mode == "semantic":
# Semantic mode: always requires index_name
if not settings.get("index_name"):
raise ServiceInitializationError(
"Azure AI Search index name is required for semantic mode. "
"Set via 'index_name' parameter or 'AZURE_SEARCH_INDEX_NAME' environment variable."
)
elif mode == "agentic":
# Agentic mode: requires exactly ONE of index_name or knowledge_base_name
if settings.get("index_name") and settings.get("knowledge_base_name"):
raise ServiceInitializationError(
"For agentic mode, provide either 'index_name' OR 'knowledge_base_name', not both. "
"Use 'index_name' to auto-create a Knowledge Base, or 'knowledge_base_name' to use an existing one."
)
if not settings.get("index_name") and not settings.get("knowledge_base_name"):
raise ServiceInitializationError(
"For agentic mode, provide either 'index_name' (to auto-create Knowledge Base) "
"or 'knowledge_base_name' (to use existing Knowledge Base). "
"Set via parameters or environment variables "
"AZURE_SEARCH_INDEX_NAME / AZURE_SEARCH_KNOWLEDGE_BASE_NAME."
)
# If using index_name to create KB, model config is required
if settings.get("index_name") and not model_deployment_name:
raise ServiceInitializationError(
"model_deployment_name is required for agentic mode when creating Knowledge Base from index. "
"This is the Azure OpenAI deployment used by the Knowledge Base for query planning."
)
# Determine the credential to use
resolved_credential: AzureKeyCredential | AsyncTokenCredential
if credential:
# AsyncTokenCredential takes precedence
resolved_credential = credential
elif isinstance(api_key, AzureKeyCredential):
resolved_credential = api_key
elif resolved_api_key := settings.get("api_key"):
resolved_credential = AzureKeyCredential(resolved_api_key.get_secret_value())
else:
raise ServiceInitializationError(
"Azure credential is required. Provide 'api_key' or 'credential' parameter "
"or set 'AZURE_SEARCH_API_KEY' environment variable."
)
self.endpoint: str = settings["endpoint"] # type: ignore[assignment] # validated above
self.index_name = settings.get("index_name")
self.credential = resolved_credential
self.mode = mode
self.top_k = top_k
self.semantic_configuration_name = semantic_configuration_name
self.vector_field_name = vector_field_name
self.embedding_function = embedding_function
self.context_prompt = context_prompt or self._DEFAULT_SEARCH_CONTEXT_PROMPT
# Agentic mode parameters (Knowledge Base)
self.azure_openai_resource_url = azure_openai_resource_url
self.azure_openai_deployment_name = model_deployment_name
# If model_name not provided, default to deployment name
self.model_name = model_name or model_deployment_name
# Use resolved KB name (from explicit param or env var)
self.knowledge_base_name = settings.get("knowledge_base_name")
self.retrieval_instructions = retrieval_instructions
self.azure_openai_api_key = azure_openai_api_key
self.knowledge_base_output_mode = knowledge_base_output_mode
self.retrieval_reasoning_effort = retrieval_reasoning_effort
self.agentic_message_history_count = agentic_message_history_count
# Determine if using existing Knowledge Base or auto-creating from index
# Since validation ensures exactly one of index_name/knowledge_base_name for agentic mode:
# - knowledge_base_name provided: use existing KB
# - index_name provided: auto-create KB from index
self._use_existing_knowledge_base = False
if mode == "agentic":
if settings.get("knowledge_base_name"):
# Use existing KB directly (supports any source type: web, blob, index, etc.)
self._use_existing_knowledge_base = True
else:
# Auto-generate KB name from index name
self.knowledge_base_name = f"{settings.get('index_name', '')}-kb"
# Auto-discover vector field if not specified
self._auto_discovered_vector_field = False
self._use_vectorizable_query = False # Will be set to True if server-side vectorization detected
if not vector_field_name and mode == "semantic":
# Attempt to auto-discover vector field from index schema
# This will be done lazily on first search to avoid blocking initialization
pass
# Validation
if vector_field_name and not embedding_function:
raise ValueError("embedding_function is required when vector_field_name is specified")
if mode == "agentic":
if not _agentic_retrieval_available:
raise ImportError(
"Agentic retrieval requires azure-search-documents >= 11.7.0b1 with Knowledge Base support. "
"Please upgrade: pip install azure-search-documents>=11.7.0b1"
)
# Only require OpenAI resource URL if NOT using existing KB
# (existing KB already has its model configuration)
# Note: model_deployment_name is already validated at initialization
if not self._use_existing_knowledge_base and not self.azure_openai_resource_url:
raise ValueError(
"azure_openai_resource_url is required for agentic mode when creating Knowledge Base from index. "
"This should be your Azure OpenAI endpoint (e.g., 'https://myresource.openai.azure.com')"
)
# Create search client for semantic mode (only if index_name is available)
self._search_client: SearchClient | None = None
if self.index_name:
self._search_client = SearchClient(
endpoint=self.endpoint,
index_name=self.index_name,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
# Create index client and retrieval client for agentic mode (Knowledge Base)
self._index_client: SearchIndexClient | None = None
self._retrieval_client: KnowledgeBaseRetrievalClient | None = None
if mode == "agentic":
self._index_client = SearchIndexClient(
endpoint=self.endpoint,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
# Retrieval client will be created after Knowledge Base initialization
self._knowledge_base_initialized = False
async def __aenter__(self) -> Self:
"""Async context manager entry."""
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: Any,
) -> None:
"""Async context manager exit - cleanup clients.
Args:
exc_type: Exception type if an error occurred.
exc_val: Exception value if an error occurred.
exc_tb: Exception traceback if an error occurred.
"""
# Close retrieval client if it was created
if self._retrieval_client is not None:
await self._retrieval_client.close()
self._retrieval_client = None
@override
async def invoking(
self,
messages: Message | MutableSequence[Message],
**kwargs: Any,
) -> Context:
"""Retrieve relevant context from Azure AI Search before model invocation.
Args:
messages: User messages to use for context retrieval.
**kwargs: Additional arguments (unused).
Returns:
Context object with retrieved documents as messages.
"""
# Convert to list and filter to USER/ASSISTANT messages with text only
messages_list = [messages] if isinstance(messages, Message) else list(messages)
def get_role_value(role: str | Any) -> str:
return role.value if hasattr(role, "value") else str(role)
filtered_messages = [
msg
for msg in messages_list
if msg and msg.text and msg.text.strip() and get_role_value(msg.role) in ["user", "assistant"]
]
if not filtered_messages:
return Context()
# Perform search based on mode
if self.mode == "semantic":
# Semantic mode: flatten messages to single query
query = "\n".join(msg.text for msg in filtered_messages)
search_result_parts = await self._semantic_search(query)
else: # agentic
# Agentic mode: pass recent messages as conversation history
recent_messages = filtered_messages[-self.agentic_message_history_count :]
search_result_parts = await self._agentic_search(recent_messages)
# Format results as context - return multiple messages for each result part
if not search_result_parts:
return Context()
# Create context messages: first message with prompt, then one message per result part
context_messages = [Message(role="user", text=self.context_prompt)]
context_messages.extend([Message(role="user", text=part) for part in search_result_parts])
return Context(messages=context_messages)
def _find_vector_fields(self, index: Any) -> list[str]:
"""Find all fields that can store vectors (have dimensions defined).
Args:
index: SearchIndex object from Azure Search.
Returns:
List of vector field names.
"""
return [
field.name
for field in index.fields
if field.vector_search_dimensions is not None and field.vector_search_dimensions > 0
]
def _find_vectorizable_fields(self, index: Any, vector_fields: list[str]) -> list[str]:
"""Find vector fields that have auto-vectorization configured.
These are fields that have a vectorizer in their profile, meaning the index
can automatically vectorize text queries without needing a client-side embedding function.
Args:
index: SearchIndex object from Azure Search.
vector_fields: List of vector field names.
Returns:
List of vectorizable field names (subset of vector_fields).
"""
vectorizable_fields: list[str] = []
# Check if index has vector search configuration
if not index.vector_search or not index.vector_search.profiles:
return vectorizable_fields
# For each vector field, check if it has a vectorizer configured
for field in index.fields:
if field.name in vector_fields and field.vector_search_profile_name:
# Find the profile for this field
profile = next(
(p for p in index.vector_search.profiles if p.name == field.vector_search_profile_name), None
)
if profile and hasattr(profile, "vectorizer_name") and profile.vectorizer_name:
# This field has server-side vectorization configured
vectorizable_fields.append(field.name)
return vectorizable_fields
async def _auto_discover_vector_field(self) -> None:
"""Auto-discover vector field from index schema.
Attempts to find vector fields in the index and detect which have server-side
vectorization configured. Prioritizes vectorizable fields (which can auto-embed text)
over regular vector fields (which require client-side embedding).
"""
if self._auto_discovered_vector_field or self.vector_field_name:
return # Already discovered or manually specified
try:
# Use existing index client or create temporary one
if not self._index_client:
self._index_client = SearchIndexClient(
endpoint=self.endpoint,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
index_client = self._index_client
# Get index schema (index_name is guaranteed to be set for semantic mode)
if not self.index_name:
logger.warning("Cannot auto-discover vector field: index_name is not set.")
self._auto_discovered_vector_field = True
return
index = await index_client.get_index(self.index_name)
# Step 1: Find all vector fields
vector_fields = self._find_vector_fields(index)
if not vector_fields:
# No vector fields found - keyword search only
logger.info(f"No vector fields found in index '{self.index_name}'. Using keyword-only search.")
self._auto_discovered_vector_field = True
return
# Step 2: Find which vector fields have server-side vectorization
vectorizable_fields = self._find_vectorizable_fields(index, vector_fields)
# Step 3: Decide which field to use
if vectorizable_fields:
# Prefer vectorizable fields (server-side embedding)
if len(vectorizable_fields) == 1:
self.vector_field_name = vectorizable_fields[0]
self._auto_discovered_vector_field = True
self._use_vectorizable_query = True # Use VectorizableTextQuery
logger.info(
f"Auto-discovered vectorizable field '{self.vector_field_name}' "
f"with server-side vectorization. No embedding_function needed."
)
else:
# Multiple vectorizable fields
logger.warning(
f"Multiple vectorizable fields found: {vectorizable_fields}. "
f"Please specify vector_field_name explicitly. Using keyword-only search."
)
elif len(vector_fields) == 1:
# Single vector field without vectorizer - needs client-side embedding
self.vector_field_name = vector_fields[0]
self._auto_discovered_vector_field = True
self._use_vectorizable_query = False
if not self.embedding_function:
logger.warning(
f"Auto-discovered vector field '{self.vector_field_name}' without server-side vectorization. "
f"Provide embedding_function for vector search, or it will fall back to keyword-only search."
)
self.vector_field_name = None
else:
# Multiple vector fields without vectorizers
logger.warning(
f"Multiple vector fields found: {vector_fields}. "
f"Please specify vector_field_name explicitly. Using keyword-only search."
)
except Exception as e:
# Log warning but continue with keyword search
logger.warning(f"Failed to auto-discover vector field: {e}. Using keyword-only search.")
self._auto_discovered_vector_field = True # Mark as attempted
async def _semantic_search(self, query: str) -> list[str]:
"""Perform semantic hybrid search with semantic ranking.
This is the recommended mode for most use cases. It combines:
- Vector search (if embedding_function provided)
- Keyword search (BM25)
- Semantic reranking (if semantic_configuration_name provided)
Args:
query: Search query text.
Returns:
List of formatted search result strings, one per document.
"""
# Auto-discover vector field if not already done
await self._auto_discover_vector_field()
vector_queries: list[VectorizableTextQuery | VectorizedQuery] = []
# Build vector query based on server-side vectorization or client-side embedding
if self.vector_field_name:
# Use larger k for vector query when semantic reranker is enabled for better ranking quality
vector_k = max(self.top_k, 50) if self.semantic_configuration_name else self.top_k
if self._use_vectorizable_query:
# Server-side vectorization: Index will auto-embed the text query
vector_queries = [
VectorizableTextQuery(
text=query,
k_nearest_neighbors=vector_k,
fields=self.vector_field_name,
)
]
elif self.embedding_function:
# Client-side embedding: We provide the vector
query_vector = await self.embedding_function(query)
vector_queries = [
VectorizedQuery(
vector=query_vector,
k_nearest_neighbors=vector_k,
fields=self.vector_field_name,
)
]
# else: vector_field_name is set but no vectorization available - skip vector search
# Build search parameters
search_params: dict[str, Any] = {
"search_text": query,
"top": self.top_k,
}
if vector_queries:
search_params["vector_queries"] = vector_queries
# Add semantic ranking if configured
if self.semantic_configuration_name:
search_params["query_type"] = QueryType.SEMANTIC
search_params["semantic_configuration_name"] = self.semantic_configuration_name
search_params["query_caption"] = QueryCaptionType.EXTRACTIVE
# Execute search (search client is guaranteed to exist for semantic mode)
if not self._search_client:
raise RuntimeError("Search client is not initialized. This should not happen in semantic mode.")
results = await self._search_client.search(**search_params) # type: ignore[reportUnknownVariableType]
# Format results with citations
formatted_results: list[str] = []
async for doc in results: # type: ignore[reportUnknownVariableType]
# Extract document ID for citation
doc_id = doc.get("id") or doc.get("@search.id") # type: ignore[reportUnknownVariableType]
# Use full document chunks with citation
doc_text: str = self._extract_document_text(doc, doc_id=doc_id) # type: ignore[reportUnknownArgumentType]
if doc_text:
formatted_results.append(doc_text) # type: ignore[reportUnknownArgumentType]
return formatted_results
async def _ensure_knowledge_base(self) -> None:
"""Ensure Knowledge Base and knowledge source are created or use existing KB.
This method is idempotent - it will only create resources if they don't exist.
Note: Azure SDK uses KnowledgeAgent classes internally, but the feature
is marketed as "Knowledge Bases" in Azure AI Search.
"""
if self._knowledge_base_initialized:
return
# Runtime validation
if not self.knowledge_base_name:
raise ValueError("knowledge_base_name is required for agentic mode")
knowledge_base_name = self.knowledge_base_name
# Path 1: Use existing Knowledge Base directly (no index needed)
# This supports KB with any source type (web, blob, index, etc.)
if self._use_existing_knowledge_base:
# Just create the retrieval client - KB already exists with its own sources
if _agentic_retrieval_available and self._retrieval_client is None:
self._retrieval_client = KnowledgeBaseRetrievalClient(
endpoint=self.endpoint,
knowledge_base_name=knowledge_base_name,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
self._knowledge_base_initialized = True
return
# Path 2: Auto-create Knowledge Base from search index
# Requires index_client and OpenAI configuration
if not self._index_client:
raise ValueError("Index client is required when creating Knowledge Base from index")
if not self.azure_openai_resource_url:
raise ValueError("azure_openai_resource_url is required when creating Knowledge Base from index")
if not self.azure_openai_deployment_name:
raise ValueError("model_deployment_name is required when creating Knowledge Base from index")
if not self.index_name:
raise ValueError("index_name is required when creating Knowledge Base from index")
# Step 1: Create or get knowledge source from index
knowledge_source_name = f"{self.index_name}-source"
try:
# Try to get existing knowledge source
await self._index_client.get_knowledge_source(knowledge_source_name)
except ResourceNotFoundError:
# Create new knowledge source if it doesn't exist
knowledge_source = SearchIndexKnowledgeSource(
name=knowledge_source_name,
description=f"Knowledge source for {self.index_name} search index",
search_index_parameters=SearchIndexKnowledgeSourceParameters(
search_index_name=self.index_name,
),
)
await self._index_client.create_knowledge_source(knowledge_source)
# Step 2: Create or update Knowledge Base
# Always create/update to ensure configuration is current
aoai_params = AzureOpenAIVectorizerParameters(
resource_url=self.azure_openai_resource_url,
deployment_name=self.azure_openai_deployment_name,
model_name=self.model_name,
api_key=self.azure_openai_api_key,
)
# Map output mode string to SDK enum
output_mode = (
KnowledgeRetrievalOutputMode.EXTRACTIVE_DATA
if self.knowledge_base_output_mode == "extractive_data"
else KnowledgeRetrievalOutputMode.ANSWER_SYNTHESIS
)
# Map reasoning effort string to SDK class
reasoning_effort_map: dict[str, KnowledgeRetrievalReasoningEffort] = {
"minimal": KnowledgeRetrievalMinimalReasoningEffort(),
"medium": KnowledgeRetrievalMediumReasoningEffort(),
"low": KnowledgeRetrievalLowReasoningEffort(),
}
reasoning_effort = reasoning_effort_map[self.retrieval_reasoning_effort]
knowledge_base = KnowledgeBase(
name=knowledge_base_name,
description=f"Knowledge Base for multi-hop retrieval across {self.index_name}",
knowledge_sources=[
KnowledgeSourceReference(
name=knowledge_source_name,
)
],
models=[KnowledgeBaseAzureOpenAIModel(azure_open_ai_parameters=aoai_params)],
output_mode=output_mode,
retrieval_reasoning_effort=reasoning_effort,
)
await self._index_client.create_or_update_knowledge_base(knowledge_base)
self._knowledge_base_initialized = True
# Create retrieval client now that Knowledge Base is initialized
if _agentic_retrieval_available and self._retrieval_client is None:
self._retrieval_client = KnowledgeBaseRetrievalClient(
endpoint=self.endpoint,
knowledge_base_name=knowledge_base_name,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
async def _agentic_search(self, messages: list[Message]) -> list[str]:
"""Perform agentic retrieval with multi-hop reasoning using Knowledge Bases.
This mode uses query planning and is slightly slower than semantic search,
but provides more accurate results through intelligent retrieval.
This method uses Azure AI Search Knowledge Bases which:
1. Analyze the query and plan sub-queries
2. Retrieve relevant documents across multiple sources
3. Perform multi-hop reasoning with an LLM
4. Synthesize a comprehensive answer with references
Args:
messages: Conversation history to use for retrieval context.
Returns:
List of answer parts from the Knowledge Base, one per content item.
"""
# Ensure Knowledge Base is initialized
await self._ensure_knowledge_base()
# Map reasoning effort string to SDK class (for retrieval requests)
reasoning_effort_map: dict[str, KBRetrievalReasoningEffort] = {
"minimal": KBRetrievalMinimalReasoningEffort(),
"medium": KBRetrievalMediumReasoningEffort(),
"low": KBRetrievalLowReasoningEffort(),
}
reasoning_effort = reasoning_effort_map[self.retrieval_reasoning_effort]
# Map output mode string to SDK enum (for retrieval requests)
output_mode = (
KBRetrievalOutputMode.EXTRACTIVE_DATA
if self.knowledge_base_output_mode == "extractive_data"
else KBRetrievalOutputMode.ANSWER_SYNTHESIS
)
# For minimal reasoning, use intents API; for medium/low, use messages API
if self.retrieval_reasoning_effort == "minimal":
# Minimal reasoning uses intents with a single search query
query = "\n".join(msg.text for msg in messages if msg.text)
intents: list[KnowledgeRetrievalIntent] = [KnowledgeRetrievalSemanticIntent(search=query)]
retrieval_request = KnowledgeBaseRetrievalRequest(
intents=intents,
retrieval_reasoning_effort=reasoning_effort,
output_mode=output_mode,
include_activity=True,
)
else:
# Medium/low reasoning uses messages with conversation history
kb_messages = [
KnowledgeBaseMessage(
role=msg.role if hasattr(msg.role, "value") else str(msg.role),
content=[KnowledgeBaseMessageTextContent(text=msg.text)],
)
for msg in messages
if msg.text
]
retrieval_request = KnowledgeBaseRetrievalRequest(
messages=kb_messages,
retrieval_reasoning_effort=reasoning_effort,
output_mode=output_mode,
include_activity=True,
)
# Use reusable retrieval client
if not self._retrieval_client:
raise RuntimeError("Retrieval client not initialized. Ensure Knowledge Base is set up correctly.")
# Perform retrieval via Knowledge Base
retrieval_result = await self._retrieval_client.retrieve(retrieval_request=retrieval_request)
# Extract answer parts from response
if retrieval_result.response and len(retrieval_result.response) > 0:
# Get the assistant's response (last message)
assistant_message = retrieval_result.response[-1]
if assistant_message.content:
# Extract all text content items as separate parts
answer_parts: list[str] = []
for content_item in assistant_message.content:
# Check if this is a text content item
if isinstance(content_item, KnowledgeBaseMessageTextContent) and content_item.text:
answer_parts.append(content_item.text)
if answer_parts:
return answer_parts
# Fallback if no answer generated
return ["No results found from Knowledge Base."]
def _extract_document_text(self, doc: dict[str, Any], doc_id: str | None = None) -> str:
"""Extract readable text from a search document with optional citation.
Args:
doc: Search result document.
doc_id: Optional document ID for citation.
Returns:
Formatted document text with citation if doc_id provided.
"""
# Try common text field names
text = ""
for field in ["content", "text", "description", "body", "chunk"]:
if doc.get(field):
text = str(doc[field])
break
# Fallback: concatenate all string fields
if not text:
text_parts: list[str] = []
for key, value in doc.items():
if isinstance(value, str) and not key.startswith("@") and key != "id":
text_parts.append(f"{key}: {value}")
text = " | ".join(text_parts) if text_parts else ""
# Add citation if document ID provided
if doc_id and text:
return f"[Source: {doc_id}] {text}"
return text
python/packages/azure-ai-search/tests/test_aisearch_new_context_provider.py → python/packages/azure-ai-search/tests/test_aisearch_context_provider.py
+15 -15
View File
@@ -7,9 +7,9 @@ from unittest.mock import AsyncMock, patch
import pytest
from agent_framework import Message
from agent_framework._sessions import AgentSession, SessionContext
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.exceptions import ServiceInitializationError, SettingNotFoundError
from agent_framework_azure_ai_search._context_provider import _AzureAISearchContextProvider
from agent_framework_azure_ai_search._context_provider import AzureAISearchContextProvider
# -- Helpers -------------------------------------------------------------------
@@ -56,7 +56,7 @@ def mock_search_client_empty() -> AsyncMock:
return client
def _make_provider(**overrides) -> _AzureAISearchContextProvider:
def _make_provider(**overrides) -> AzureAISearchContextProvider:
"""Create a semantic-mode provider with mocked internals (skips auto-discovery)."""
defaults = {
"source_id": "aisearch",
@@ -65,7 +65,7 @@ def _make_provider(**overrides) -> _AzureAISearchContextProvider:
"api_key": "test-key",
}
defaults.update(overrides)
provider = _AzureAISearchContextProvider(**defaults)
provider = AzureAISearchContextProvider(**defaults)
provider._auto_discovered_vector_field = True # skip auto-discovery
return provider
@@ -88,8 +88,8 @@ class TestInitSemantic:
assert provider.source_id == "my-source"
def test_missing_endpoint_raises(self) -> None:
with patch.dict(os.environ, {}, clear=True), pytest.raises(ServiceInitializationError, match="endpoint"):
_AzureAISearchContextProvider(
with patch.dict(os.environ, {}, clear=True), pytest.raises(SettingNotFoundError, match="endpoint"):
AzureAISearchContextProvider(
source_id="s",
endpoint=None,
index_name="idx",
@@ -97,8 +97,8 @@ class TestInitSemantic:
)
def test_missing_index_name_semantic_raises(self) -> None:
with pytest.raises(ServiceInitializationError, match="index name"):
_AzureAISearchContextProvider(
with pytest.raises(SettingNotFoundError, match="index_name"):
AzureAISearchContextProvider(
source_id="s",
endpoint="https://test.search.windows.net",
index_name=None,
@@ -112,7 +112,7 @@ class TestInitSemantic:
"AZURE_SEARCH_API_KEY": "env-key",
}
with patch.dict(os.environ, env, clear=False):
provider = _AzureAISearchContextProvider(source_id="env-test")
provider = AzureAISearchContextProvider(source_id="env-test")
assert provider.endpoint == "https://env.search.windows.net"
assert provider.index_name == "env-index"
@@ -124,8 +124,8 @@ class TestInitAgenticValidation:
"""Initialization validation tests for agentic mode."""
def test_both_index_and_kb_raises(self) -> None:
with pytest.raises(ServiceInitializationError, match="not both"):
_AzureAISearchContextProvider(
with pytest.raises(SettingNotFoundError, match="multiple were set"):
AzureAISearchContextProvider(
source_id="s",
endpoint="https://test.search.windows.net",
index_name="idx",
@@ -137,8 +137,8 @@ class TestInitAgenticValidation:
)
def test_neither_index_nor_kb_raises(self) -> None:
with pytest.raises(ServiceInitializationError, match="provide either"):
_AzureAISearchContextProvider(
with pytest.raises(SettingNotFoundError, match="none was set"):
AzureAISearchContextProvider(
source_id="s",
endpoint="https://test.search.windows.net",
api_key="key",
@@ -147,7 +147,7 @@ class TestInitAgenticValidation:
def test_missing_model_deployment_name_raises(self) -> None:
with pytest.raises(ServiceInitializationError, match="model_deployment_name"):
_AzureAISearchContextProvider(
AzureAISearchContextProvider(
source_id="s",
endpoint="https://test.search.windows.net",
index_name="idx",
@@ -158,7 +158,7 @@ class TestInitAgenticValidation:
def test_vector_field_without_embedding_raises(self) -> None:
with pytest.raises(ValueError, match="embedding_function"):
_AzureAISearchContextProvider(
AzureAISearchContextProvider(
source_id="s",
endpoint="https://test.search.windows.net",
index_name="idx",
python/packages/azure-ai-search/tests/test_search_provider.py
-1018
View File
File diff suppressed because it is too large Load Diff
python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py
+13 -13
View File
@@ -9,7 +9,7 @@ from typing import Any, Generic, cast
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
Agent,
ContextProvider,
BaseContextProvider,
FunctionTool,
MiddlewareTypes,
normalize_tools,
@@ -176,7 +176,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
| None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Create a new agent on the Azure AI service and return a Agent.
@@ -195,7 +195,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: List of middleware to intercept agent and function invocations.
context_provider: Context provider to include during agent invocation.
context_providers: Context providers to include during agent invocation.
Returns:
Agent: A Agent instance configured with the created agent.
@@ -259,7 +259,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
normalized_tools,
default_options=default_options,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
async def get_agent(
@@ -273,7 +273,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
| None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Retrieve an existing agent from the service and return a Agent.
@@ -289,7 +289,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: List of middleware to intercept agent and function invocations.
context_provider: Context provider to include during agent invocation.
context_providers: Context providers to include during agent invocation.
Returns:
Agent: A Agent instance configured with the retrieved agent.
@@ -316,7 +316,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
normalized_tools,
default_options=default_options,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
def as_agent(
@@ -329,7 +329,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
| None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Wrap an existing Agent SDK object as a Agent without making HTTP calls.
@@ -343,7 +343,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: List of middleware to intercept agent and function invocations.
context_provider: Context provider to include during agent invocation.
context_providers: Context providers to include during agent invocation.
Returns:
Agent: A Agent instance configured with the agent.
@@ -373,7 +373,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
normalized_tools,
default_options=default_options,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
def _to_chat_agent_from_agent(
@@ -382,7 +382,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
provided_tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Create a Agent from an Agent SDK object.
@@ -392,7 +392,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: List of middleware to intercept agent and function invocations.
context_provider: Context provider to include during agent invocation.
context_providers: Context providers to include during agent invocation.
"""
# Create the underlying client
client = AzureAIAgentClient(
@@ -416,7 +416,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
tools=merged_tools,
default_options=default_options, # type: ignore[arg-type]
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
def _merge_tools(
python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
+5 -8
View File
@@ -15,14 +15,13 @@ from agent_framework import (
Agent,
Annotation,
BaseChatClient,
BaseContextProvider,
ChatAndFunctionMiddlewareTypes,
ChatMessageStoreProtocol,
ChatMiddlewareLayer,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
ContextProvider,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
@@ -211,6 +210,7 @@ class AzureAIAgentClient(
"""Azure AI Agent Chat client with middleware, telemetry, and function invocation support."""
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai" # type: ignore[reportIncompatibleVariableOverride, misc]
STORES_BY_DEFAULT: ClassVar[bool] = True # type: ignore[reportIncompatibleVariableOverride, misc]
# region Hosted Tool Factory Methods
@@ -1434,8 +1434,7 @@ class AzureAIAgentClient(
| Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: AzureAIAgentOptionsT | Mapping[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> Agent[AzureAIAgentOptionsT]:
@@ -1455,8 +1454,7 @@ class AzureAIAgentClient(
instructions: Optional instructions for the agent.
tools: The tools to use for the request.
default_options: A TypedDict containing chat options.
chat_message_store_factory: Factory function to create an instance of ChatMessageStoreProtocol.
context_provider: Context providers to include during agent invocation.
context_providers: Context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
kwargs: Any additional keyword arguments.
@@ -1470,8 +1468,7 @@ class AzureAIAgentClient(
instructions=instructions,
tools=tools,
default_options=default_options,
chat_message_store_factory=chat_message_store_factory,
context_provider=context_provider,
context_providers=context_providers,
middleware=middleware,
**kwargs,
)
python/packages/azure-ai/agent_framework_azure_ai/_client.py
+4 -8
View File
@@ -9,10 +9,9 @@ from typing import Any, ClassVar, Generic, Literal, TypedDict, TypeVar, cast
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
Agent,
BaseContextProvider,
ChatAndFunctionMiddlewareTypes,
ChatMessageStoreProtocol,
ChatMiddlewareLayer,
ContextProvider,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
@@ -808,8 +807,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
| Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: AzureAIClientOptionsT | Mapping[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> Agent[AzureAIClientOptionsT]:
@@ -829,8 +827,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
instructions: Optional instructions for the agent.
tools: The tools to use for the request.
default_options: A TypedDict containing chat options.
chat_message_store_factory: Factory function to create an instance of ChatMessageStoreProtocol.
context_provider: Context providers to include during agent invocation.
context_providers: Context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
kwargs: Any additional keyword arguments.
@@ -844,8 +841,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
instructions=instructions,
tools=tools,
default_options=default_options,
chat_message_store_factory=chat_message_store_factory,
context_provider=context_provider,
context_providers=context_providers,
middleware=middleware,
**kwargs,
)
python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
+13 -13
View File
@@ -9,7 +9,7 @@ from typing import Any, Generic
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
Agent,
ContextProvider,
BaseContextProvider,
FunctionTool,
MiddlewareTypes,
get_logger,
@@ -168,7 +168,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
| None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Create a new agent on the Azure AI service and return a local Agent wrapper.
@@ -182,7 +182,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: List of middleware to intercept agent and function invocations.
context_provider: Context provider to include during agent invocation.
context_providers: Context providers to include during agent invocation.
Returns:
Agent: A Agent instance configured with the created agent.
@@ -255,7 +255,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
normalized_tools,
default_options=default_options,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
async def get_agent(
@@ -270,7 +270,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
| None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Retrieve an existing agent from the Azure AI service and return a local Agent wrapper.
@@ -284,7 +284,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: List of middleware to intercept agent and function invocations.
context_provider: Context provider to include during agent invocation.
context_providers: Context providers to include during agent invocation.
Returns:
Agent: A Agent instance configured with the retrieved agent.
@@ -317,7 +317,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
normalize_tools(tools),
default_options=default_options,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
def as_agent(
@@ -330,7 +330,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
| None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Wrap an SDK agent version object into a Agent without making HTTP calls.
@@ -342,7 +342,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: List of middleware to intercept agent and function invocations.
context_provider: Context provider to include during agent invocation.
context_providers: Context providers to include during agent invocation.
Returns:
Agent: A Agent instance configured with the agent version.
@@ -361,7 +361,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
normalize_tools(tools),
default_options=default_options,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
def _to_chat_agent_from_details(
@@ -370,7 +370,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
provided_tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Create a Agent from an AgentVersionDetails.
@@ -381,7 +381,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: List of middleware to intercept agent and function invocations.
context_provider: Context provider to include during agent invocation.
context_providers: Context providers to include during agent invocation.
"""
if not isinstance(details.definition, PromptAgentDefinition):
raise ValueError("Agent definition must be PromptAgentDefinition to get a Agent.")
@@ -409,7 +409,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
tools=merged_tools,
default_options=default_options, # type: ignore[arg-type]
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
def _merge_tools(
python/packages/azure-ai/tests/test_azure_ai_agent_client.py
+13 -13
View File
@@ -11,7 +11,7 @@ from agent_framework import (
Agent,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
@@ -1524,24 +1524,24 @@ async def test_azure_ai_chat_client_agent_basic_run_streaming() -> None:
@pytest.mark.flaky
@skip_if_azure_ai_integration_tests_disabled
async def test_azure_ai_chat_client_agent_thread_persistence() -> None:
"""Test Agent thread persistence across runs with AzureAIAgentClient."""
"""Test Agent session persistence across runs with AzureAIAgentClient."""
async with Agent(
client=AzureAIAgentClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# Create a new session that will be reused
session = agent.create_session()
# First message - establish context
first_response = await agent.run(
"Remember this number: 42. What number did I just tell you to remember?", thread=thread
"Remember this number: 42. What number did I just tell you to remember?", session=session
)
assert isinstance(first_response, AgentResponse)
assert "42" in first_response.text
# Second message - test conversation memory
second_response = await agent.run(
"What number did I tell you to remember in my previous message?", thread=thread
"What number did I tell you to remember in my previous message?", session=session
)
assert isinstance(second_response, AgentResponse)
assert "42" in second_response.text
@@ -1555,16 +1555,16 @@ async def test_azure_ai_chat_client_agent_existing_thread_id() -> None:
client=AzureAIAgentClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and get the thread ID
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My name is Alice. Remember this.", thread=thread)
# Start a conversation and get the session ID
session = first_agent.create_session()
first_response = await first_agent.run("My name is Alice. Remember this.", session=session)
# Validate first response
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# The thread ID is set after the first response
existing_thread_id = thread.service_thread_id
existing_thread_id = session.service_session_id
assert existing_thread_id is not None
# Now continue with the same thread ID in a new agent instance
@@ -1572,11 +1572,11 @@ async def test_azure_ai_chat_client_agent_existing_thread_id() -> None:
client=AzureAIAgentClient(thread_id=existing_thread_id, credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Create a thread with the existing ID
thread = AgentThread(service_thread_id=existing_thread_id)
# Create a session with the existing ID
session = AgentSession(service_session_id=existing_thread_id)
# Ask about the previous conversation
response2 = await second_agent.run("What is my name?", thread=thread)
response2 = await second_agent.run("What is my name?", session=session)
# Validate that the agent remembers the previous conversation
assert isinstance(response2, AgentResponse)
python/packages/azure-ai/tests/test_azure_ai_client.py
+15 -15
View File
@@ -1473,39 +1473,39 @@ async def test_integration_agent_hosted_code_interpreter_tool():
@pytest.mark.flaky
@skip_if_azure_ai_integration_tests_disabled
async def test_integration_agent_existing_thread():
"""Test Azure Responses Client agent with existing thread to continue conversations across agent instances."""
# First conversation - capture the thread
preserved_thread = None
async def test_integration_agent_existing_session():
"""Test Azure Responses Client agent with existing session to continue conversations across agent instances."""
# First conversation - capture the session
preserved_session = None
async with (
temporary_chat_client(agent_name="af-int-test-existing-thread") as client,
temporary_chat_client(agent_name="af-int-test-existing-session") as client,
Agent(
client=client,
instructions="You are a helpful assistant with good memory.",
) as first_agent,
):
# Start a conversation and capture the thread
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My hobby is photography. Remember this.", thread=thread, store=True)
# Start a conversation and capture the session
session = first_agent.create_session()
first_response = await first_agent.run("My hobby is photography. Remember this.", session=session, store=True)
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Preserve the thread for reuse
preserved_thread = thread
# Preserve the session for reuse
preserved_session = session
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
# Second conversation - reuse the session in a new agent instance
if preserved_session:
async with (
temporary_chat_client(agent_name="af-int-test-existing-thread-2") as client,
temporary_chat_client(agent_name="af-int-test-existing-session-2") as client,
Agent(
client=client,
instructions="You are a helpful assistant with good memory.",
) as second_agent,
):
# Reuse the preserved thread
second_response = await second_agent.run("What is my hobby?", thread=preserved_thread)
# Reuse the preserved session
second_response = await second_agent.run("What is my hobby?", session=preserved_session)
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
python/packages/azurefunctions/agent_framework_azurefunctions/_app.py
+2 -2
View File
@@ -135,8 +135,8 @@ class AgentFunctionApp(DFAppBase):
@app.orchestration_trigger(context_name="context")
def my_orchestration(context):
writer = app.get_agent(context, "WeatherAgent")
thread = writer.get_new_thread()
forecast_task = writer.run("What's the forecast?", thread=thread)
session = writer.create_session()
forecast_task = writer.run("What's the forecast?", session=session)
forecast = yield forecast_task
return forecast
python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py
+3 -3
View File
@@ -9,7 +9,7 @@ from collections.abc import Callable
from typing import TYPE_CHECKING, Any, TypeAlias
import azure.durable_functions as df
from agent_framework import AgentThread, get_logger
from agent_framework import AgentSession, get_logger
from agent_framework_durabletask import (
DurableAgentExecutor,
RunRequest,
@@ -178,11 +178,11 @@ class AzureFunctionsAgentExecutor(DurableAgentExecutor[AgentTask]):
self,
agent_name: str,
run_request: RunRequest,
thread: AgentThread | None = None,
session: AgentSession | None = None,
) -> AgentTask:
# Resolve session
session_id = self._create_session_id(agent_name, thread)
session_id = self._create_session_id(agent_name, session)
entity_id = df.EntityId(
name=session_id.entity_name,
python/packages/azurefunctions/tests/test_orchestration.py
+17 -17
View File
@@ -214,10 +214,10 @@ class TestAzureFunctionsFireAndForget:
context.call_entity = Mock(return_value=_create_entity_task())
agent = DurableAIAgent(executor, "TestAgent")
thread = agent.get_new_thread()
session = agent.create_session()
# Run with wait_for_response=False
result = agent.run("Test message", thread=thread, options={"wait_for_response": False})
result = agent.run("Test message", session=session, options={"wait_for_response": False})
# Verify signal_entity was called and call_entity was not
assert context.signal_entity.call_count == 1
@@ -232,9 +232,9 @@ class TestAzureFunctionsFireAndForget:
context.signal_entity = Mock()
agent = DurableAIAgent(executor, "TestAgent")
thread = agent.get_new_thread()
session = agent.create_session()
result = agent.run("Test message", thread=thread, options={"wait_for_response": False})
result = agent.run("Test message", session=session, options={"wait_for_response": False})
# Task should be immediately complete
assert isinstance(result, AgentTask)
@@ -246,9 +246,9 @@ class TestAzureFunctionsFireAndForget:
context.signal_entity = Mock()
agent = DurableAIAgent(executor, "TestAgent")
thread = agent.get_new_thread()
session = agent.create_session()
result = agent.run("Test message", thread=thread, options={"wait_for_response": False})
result = agent.run("Test message", session=session, options={"wait_for_response": False})
# Get the result
response = result.result
@@ -267,9 +267,9 @@ class TestAzureFunctionsFireAndForget:
context.call_entity = Mock(return_value=_create_entity_task())
agent = DurableAIAgent(executor, "TestAgent")
thread = agent.get_new_thread()
session = agent.create_session()
result = agent.run("Test message", thread=thread, options={"wait_for_response": True})
result = agent.run("Test message", session=session, options={"wait_for_response": True})
# Verify call_entity was called and signal_entity was not
assert context.call_entity.call_count == 1
@@ -298,15 +298,15 @@ class TestOrchestrationIntegration:
# Create agent directly with executor (not via app.get_agent)
agent = DurableAIAgent(executor, "WriterAgent")
# Create thread
thread = agent.get_new_thread()
# Create session
session = agent.create_session()
# First call - returns AgentTask
task1 = agent.run("Write something", thread=thread)
task1 = agent.run("Write something", session=session)
assert isinstance(task1, AgentTask)
# Second call - returns AgentTask
task2 = agent.run("Improve: something", thread=thread)
task2 = agent.run("Improve: something", session=session)
assert isinstance(task2, AgentTask)
# Verify both calls used the same entity (same session key)
@@ -315,7 +315,7 @@ class TestOrchestrationIntegration:
# EntityId format is @dafx-writeragent@<uuid_hex>
expected_entity_id = f"@dafx-writeragent@{uuid_hexes[0]}"
assert entity_calls[0]["entity_id"] == expected_entity_id
# generate_unique_id called 3 times: thread + 2 correlation IDs
# generate_unique_id called 3 times: session + 2 correlation IDs
assert executor.generate_unique_id.call_count == 3
def test_multiple_agents_in_orchestration(self, executor_with_multiple_uuids: tuple[Any, Mock, list[str]]) -> None:
@@ -334,12 +334,12 @@ class TestOrchestrationIntegration:
writer = DurableAIAgent(executor, "WriterAgent")
editor = DurableAIAgent(executor, "EditorAgent")
writer_thread = writer.get_new_thread()
editor_thread = editor.get_new_thread()
writer_session = writer.create_session()
editor_session = editor.create_session()
# Call both agents - returns AgentTasks
writer_task = writer.run("Write", thread=writer_thread)
editor_task = editor.run("Edit", thread=editor_thread)
writer_task = writer.run("Write", session=writer_session)
editor_task = editor.run("Edit", session=editor_session)
assert isinstance(writer_task, AgentTask)
assert isinstance(editor_task, AgentTask)
python/packages/claude/agent_framework_claude/_agent.py
+26 -35
View File
@@ -12,12 +12,13 @@ from agent_framework import (
AgentMiddlewareTypes,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
BaseContextProvider,
Content,
ContextProvider,
FunctionTool,
Message,
ResponseStream,
get_logger,
normalize_messages,
)
@@ -184,9 +185,9 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
.. code-block:: python
async with ClaudeAgent() as agent:
thread = agent.get_new_thread()
await agent.run("Remember my name is Alice", thread=thread)
response = await agent.run("What's my name?", thread=thread)
session = agent.create_session()
await agent.run("Remember my name is Alice", session=session)
response = await agent.run("What's my name?", session=session)
# Claude will remember "Alice" from the same session
With Agent Framework tools:
@@ -214,7 +215,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
id: str | None = None,
name: str | None = None,
description: str | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[AgentMiddlewareTypes] | None = None,
tools: FunctionTool
| Callable[..., Any]
@@ -237,7 +238,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
id: Unique identifier for the agent.
name: Name of the agent.
description: Description of the agent.
context_provider: Context provider for the agent.
context_providers: Context providers for the agent.
middleware: List of middleware.
tools: Tools for the agent. Can be:
- Strings for built-in tools (e.g., "Read", "Write", "Bash", "Glob")
@@ -250,7 +251,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
id=id,
name=name,
description=description,
context_provider=context_provider,
context_providers=context_providers,
middleware=middleware,
)
@@ -559,7 +560,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: OptionsT | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]: ...
@@ -570,7 +571,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: OptionsT | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AgentResponse[Any]: ...
@@ -580,7 +581,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: OptionsT | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate] | Awaitable[AgentResponse[Any]]:
@@ -592,46 +593,36 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
Keyword Args:
stream: If True, returns an async iterable of updates. If False (default),
returns an awaitable AgentResponse.
thread: The conversation thread. If thread has service_thread_id set,
session: The conversation session. If session has service_session_id set,
the agent will resume that session.
options: Runtime options (model, permission_mode can be changed per-request).
kwargs: Additional keyword arguments.
Returns:
When stream=True: An AsyncIterable[AgentResponseUpdate] for streaming updates.
When stream=True: An ResponseStream for streaming updates.
When stream=False: An Awaitable[AgentResponse] with the complete response.
"""
if stream:
return self._run_streaming(messages, thread=thread, options=options, **kwargs)
return self._run_non_streaming(messages, thread=thread, options=options, **kwargs)
async def _run_non_streaming(
self,
messages: str | Message | Sequence[str | Message] | None = None,
*,
thread: AgentThread | None = None,
options: OptionsT | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AgentResponse[Any]:
"""Internal non-streaming implementation."""
thread = thread or self.get_new_thread()
return await AgentResponse.from_update_generator(
self._run_streaming(messages, thread=thread, options=options, **kwargs)
response = ResponseStream(
self._get_stream(messages, session=session, options=options, **kwargs),
finalizer=AgentResponse.from_updates,
)
if stream:
return response
return response.get_final_response()
async def _run_streaming(
async def _get_stream(
self,
messages: str | Message | Sequence[str | Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: OptionsT | MutableMapping[str, Any] | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
"""Internal streaming implementation."""
thread = thread or self.get_new_thread()
session = session or self.create_session()
# Ensure we're connected to the right session
await self._ensure_session(thread.service_thread_id)
await self._ensure_session(session.service_session_id)
if not self._client:
raise ServiceException("Claude SDK client not initialized.")
@@ -696,6 +687,6 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
raise ServiceException(f"Claude API error: {error_msg}")
session_id = message.session_id
# Update thread with session ID
# Update session with session ID
if session_id:
thread.service_thread_id = session_id
session.service_session_id = session_id
python/packages/claude/tests/test_claude_agent.py
+18 -26
View File
@@ -4,7 +4,7 @@ from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import AgentResponseUpdate, AgentThread, Content, Message, tool
from agent_framework import AgentResponseUpdate, AgentSession, Content, Message, tool
from agent_framework._settings import load_settings
from agent_framework_claude import ClaudeAgent, ClaudeAgentOptions, ClaudeAgentSettings
@@ -267,12 +267,12 @@ class TestClaudeAgentRun:
with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client):
agent = ClaudeAgent()
thread = agent.get_new_thread()
await agent.run("Hello", thread=thread)
assert thread.service_thread_id == "test-session-id"
session = agent.create_session()
await agent.run("Hello", session=session)
assert session.service_session_id == "test-session-id"
async def test_run_with_thread(self) -> None:
"""Test run with existing thread."""
async def test_run_with_session(self) -> None:
"""Test run with existing session."""
from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
from claude_agent_sdk.types import StreamEvent
@@ -302,9 +302,9 @@ class TestClaudeAgentRun:
with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client):
agent = ClaudeAgent()
thread = agent.get_new_thread()
thread.service_thread_id = "existing-session"
await agent.run("Hello", thread=thread)
session = agent.create_session()
session.service_session_id = "existing-session"
await agent.run("Hello", session=session)
# region Test ClaudeAgent Run Stream
@@ -440,26 +440,18 @@ class TestClaudeAgentRunStream:
class TestClaudeAgentSessionManagement:
"""Tests for ClaudeAgent session management."""
def test_get_new_thread(self) -> None:
"""Test get_new_thread creates a new thread."""
def test_create_session(self) -> None:
"""Test create_session creates a new session."""
agent = ClaudeAgent()
thread = agent.get_new_thread()
assert isinstance(thread, AgentThread)
assert thread.service_thread_id is None
session = agent.create_session()
assert isinstance(session, AgentSession)
assert session.service_session_id is None
def test_get_new_thread_with_service_thread_id(self) -> None:
"""Test get_new_thread with existing service_thread_id."""
def test_create_session_with_service_session_id(self) -> None:
"""Test create_session with existing service_session_id."""
agent = ClaudeAgent()
thread = agent.get_new_thread(service_thread_id="existing-session-123")
assert isinstance(thread, AgentThread)
assert thread.service_thread_id == "existing-session-123"
def test_thread_inherits_context_provider(self) -> None:
"""Test that thread inherits context provider."""
mock_provider = MagicMock()
agent = ClaudeAgent(context_provider=mock_provider)
thread = agent.get_new_thread()
assert thread.context_provider == mock_provider
session = agent.create_session(session_id="existing-session-123")
assert isinstance(session, AgentSession)
async def test_ensure_session_creates_client(self) -> None:
"""Test _ensure_session creates client when not started."""
python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py
+22 -22
View File
@@ -9,10 +9,10 @@ from agent_framework import (
AgentMiddlewareTypes,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
BaseContextProvider,
Content,
ContextProvider,
Message,
ResponseStream,
normalize_messages,
@@ -59,7 +59,7 @@ class CopilotStudioAgent(BaseAgent):
id: str | None = None,
name: str | None = None,
description: str | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: list[AgentMiddlewareTypes] | None = None,
environment_id: str | None = None,
agent_identifier: str | None = None,
@@ -87,7 +87,7 @@ class CopilotStudioAgent(BaseAgent):
id: id of the CopilotAgent
name: Name of the CopilotAgent
description: Description of the CopilotAgent
context_provider: Context Provider, to be used by the copilot agent.
context_providers: Context Providers, to be used by the copilot agent.
middleware: Agent middleware used by the agent, should be a list of AgentMiddlewareTypes.
environment_id: Environment ID of the Power Platform environment containing
the Copilot Studio app. Can also be set via COPILOTSTUDIOAGENT__ENVIRONMENTID
@@ -118,7 +118,7 @@ class CopilotStudioAgent(BaseAgent):
id=id,
name=name,
description=description,
context_provider=context_provider,
context_providers=context_providers,
middleware=middleware,
)
if not client:
@@ -190,7 +190,7 @@ class CopilotStudioAgent(BaseAgent):
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: Literal[False] = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse]: ...
@@ -200,7 +200,7 @@ class CopilotStudioAgent(BaseAgent):
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ...
@@ -209,7 +209,7 @@ class CopilotStudioAgent(BaseAgent):
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
"""Get a response from the agent.
@@ -223,7 +223,7 @@ class CopilotStudioAgent(BaseAgent):
Keyword Args:
stream: Whether to stream the response. Defaults to False.
thread: The conversation thread associated with the message(s).
session: The conversation session associated with the message(s).
kwargs: Additional keyword arguments.
Returns:
@@ -231,26 +231,26 @@ class CopilotStudioAgent(BaseAgent):
When stream=True: A ResponseStream of AgentResponseUpdate items.
"""
if stream:
return self._run_stream_impl(messages=messages, thread=thread, **kwargs)
return self._run_impl(messages=messages, thread=thread, **kwargs)
return self._run_stream_impl(messages=messages, session=session, **kwargs)
return self._run_impl(messages=messages, session=session, **kwargs)
async def _run_impl(
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> AgentResponse:
"""Non-streaming implementation of run."""
if not thread:
thread = self.get_new_thread()
thread.service_thread_id = await self._start_new_conversation()
if not session:
session = self.create_session()
session.service_session_id = await self._start_new_conversation()
input_messages = normalize_messages(messages)
question = "\n".join([message.text for message in input_messages])
activities = self.client.ask_question(question, thread.service_thread_id)
activities = self.client.ask_question(question, session.service_session_id)
response_messages: list[Message] = []
response_id: str | None = None
@@ -263,22 +263,22 @@ class CopilotStudioAgent(BaseAgent):
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
"""Streaming implementation of run."""
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
nonlocal thread
if not thread:
thread = self.get_new_thread()
thread.service_thread_id = await self._start_new_conversation()
nonlocal session
if not session:
session = self.create_session()
session.service_session_id = await self._start_new_conversation()
input_messages = normalize_messages(messages)
question = "\n".join([message.text for message in input_messages])
activities = self.client.ask_question(question, thread.service_thread_id)
activities = self.client.ask_question(question, session.service_session_id)
async for message in self._process_activities(activities, streaming=True):
yield AgentResponseUpdate(
python/packages/copilotstudio/tests/test_copilot_agent.py
+11 -11
View File
@@ -4,7 +4,7 @@ from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from agent_framework import AgentResponse, AgentResponseUpdate, AgentThread, Content, Message
from agent_framework import AgentResponse, AgentResponseUpdate, AgentSession, Content, Message
from agent_framework.exceptions import ServiceException, ServiceInitializationError
from microsoft_agents.copilotstudio.client import CopilotClient
@@ -165,10 +165,10 @@ class TestCopilotStudioAgent:
assert content.text == "Test response"
assert response.messages[0].role == "assistant"
async def test_run_with_thread(self, mock_copilot_client: MagicMock, mock_activity: MagicMock) -> None:
"""Test run method with existing thread."""
async def test_run_with_session(self, mock_copilot_client: MagicMock, mock_activity: MagicMock) -> None:
"""Test run method with existing session."""
agent = CopilotStudioAgent(client=mock_copilot_client)
thread = AgentThread()
session = AgentSession()
conversation_activity = MagicMock()
conversation_activity.conversation.id = "test-conversation-id"
@@ -176,11 +176,11 @@ class TestCopilotStudioAgent:
mock_copilot_client.start_conversation.return_value = create_async_generator([conversation_activity])
mock_copilot_client.ask_question.return_value = create_async_generator([mock_activity])
response = await agent.run("test message", thread=thread)
response = await agent.run("test message", session=session)
assert isinstance(response, AgentResponse)
assert len(response.messages) == 1
assert thread.service_thread_id == "test-conversation-id"
assert session.service_session_id == "test-conversation-id"
async def test_run_start_conversation_failure(self, mock_copilot_client: MagicMock) -> None:
"""Test run method when conversation start fails."""
@@ -217,10 +217,10 @@ class TestCopilotStudioAgent:
assert response_count == 1
async def test_run_streaming_with_thread(self, mock_copilot_client: MagicMock) -> None:
"""Test run(stream=True) method with existing thread."""
async def test_run_streaming_with_session(self, mock_copilot_client: MagicMock) -> None:
"""Test run(stream=True) method with existing session."""
agent = CopilotStudioAgent(client=mock_copilot_client)
thread = AgentThread()
session = AgentSession()
conversation_activity = MagicMock()
conversation_activity.conversation.id = "test-conversation-id"
@@ -235,7 +235,7 @@ class TestCopilotStudioAgent:
mock_copilot_client.ask_question.return_value = create_async_generator([typing_activity])
response_count = 0
async for response in agent.run("test message", thread=thread, stream=True):
async for response in agent.run("test message", session=session, stream=True):
assert isinstance(response, AgentResponseUpdate)
content = response.contents[0]
assert content.type == "text"
@@ -243,7 +243,7 @@ class TestCopilotStudioAgent:
response_count += 1
assert response_count == 1
assert thread.service_thread_id == "test-conversation-id"
assert session.service_session_id == "test-conversation-id"
async def test_run_streaming_no_typing_activity(self, mock_copilot_client: MagicMock) -> None:
"""Test run(stream=True) method with non-typing activity."""
python/packages/core/AGENTS.md
+6 -11
View File
@@ -12,8 +12,7 @@ agent_framework/
├── _types.py # Core types (Message, ChatResponse, Content, etc.)
├── _tools.py # Tool definitions and function invocation
├── _middleware.py # Middleware system for request/response interception
├── _threads.py # AgentThread and message store abstractions
├── _memory.py # Context providers for memory/RAG
├── _sessions.py # AgentSession and context provider abstractions
├── _mcp.py # Model Context Protocol support
├── _workflows/ # Workflow orchestration (sequential, concurrent, handoff, etc.)
├── openai/ # Built-in OpenAI client
@@ -57,16 +56,12 @@ agent_framework/
- **`FunctionMiddleware`** - Intercepts function/tool invocations
- **`AgentContext`** / **`ChatContext`** / **`FunctionInvocationContext`** - Context objects passed through middleware
### Threads (`_threads.py`)
### Sessions (`_sessions.py`)
- **`AgentThread`** - Manages conversation history for an agent
- **`ChatMessageStoreProtocol`** - Protocol for persistent message storage
- **`ChatMessageStore`** - Default in-memory implementation
### Memory (`_memory.py`)
- **`ContextProvider`** - Protocol for providing additional context to agents (RAG, memory systems)
- **`Context`** - Container for context data
- **`AgentSession`** - Manages conversation state and session metadata
- **`SessionContext`** - Context object for session-scoped data during agent runs
- **`BaseContextProvider`** - Base class for context providers (RAG, memory systems)
- **`BaseHistoryProvider`** - Base class for conversation history storage
### Workflows (`_workflows/`)
python/packages/core/agent_framework/__init__.py
+1 -2
View File
@@ -13,10 +13,9 @@ from ._agents import * # noqa: F403
from ._clients import * # noqa: F403
from ._logging import * # noqa: F403
from ._mcp import * # noqa: F403
from ._memory import * # noqa: F403
from ._middleware import * # noqa: F403
from ._sessions import * # noqa: F403
from ._telemetry import * # noqa: F403
from ._threads import * # noqa: F403
from ._tools import * # noqa: F403
from ._types import * # noqa: F403
from ._workflows import * # noqa: F403
python/packages/core/agent_framework/_agents.py
+214 -290
View File
@@ -31,10 +31,9 @@ from pydantic import BaseModel, Field, create_model
from ._clients import BaseChatClient, SupportsChatGetResponse
from ._logging import get_logger
from ._mcp import LOG_LEVEL_MAPPING, MCPTool
from ._memory import Context, ContextProvider
from ._middleware import AgentMiddlewareLayer, MiddlewareTypes
from ._serialization import SerializationMixin
from ._threads import AgentThread, ChatMessageStoreProtocol
from ._sessions import AgentSession, BaseContextProvider, BaseHistoryProvider, InMemoryHistoryProvider, SessionContext
from ._tools import (
FunctionInvocationLayer,
FunctionTool,
@@ -49,7 +48,7 @@ from ._types import (
map_chat_to_agent_update,
normalize_messages,
)
from .exceptions import AgentExecutionException, AgentInitializationError
from .exceptions import AgentExecutionException
from .observability import AgentTelemetryLayer
if sys.version_info >= (3, 13):
@@ -57,9 +56,9 @@ if sys.version_info >= (3, 13):
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
pass # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore[import] # pragma: no cover
pass # type: ignore[import] # pragma: no cover
if sys.version_info >= (3, 11):
from typing import Self, TypedDict # pragma: no cover
else:
@@ -68,14 +67,9 @@ else:
if TYPE_CHECKING:
from ._types import ChatOptions
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None, covariant=True)
ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel)
logger = get_logger("agent_framework")
ThreadTypeT = TypeVar("ThreadTypeT", bound="AgentThread")
ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel)
OptionsCoT = TypeVar(
"OptionsCoT",
bound=TypedDict, # type: ignore[valid-type]
@@ -155,9 +149,10 @@ def _sanitize_agent_name(agent_name: str | None) -> str | None:
class _RunContext(TypedDict):
thread: AgentThread
session: AgentSession | None
session_context: SessionContext
input_messages: list[Message]
thread_messages: list[Message]
session_messages: list[Message]
agent_name: str
chat_options: dict[str, Any]
filtered_kwargs: dict[str, Any]
@@ -197,7 +192,7 @@ class SupportsAgentRun(Protocol):
self.name = "Custom Agent"
self.description = "A fully custom agent implementation"
async def run(self, messages=None, *, stream=False, thread=None, **kwargs):
async def run(self, messages=None, *, stream=False, session=None, **kwargs):
if stream:
# Your custom streaming implementation
async def _stream():
@@ -212,9 +207,15 @@ class SupportsAgentRun(Protocol):
return AgentResponse(messages=[], response_id="custom-response")
def get_new_thread(self, **kwargs):
# Return your own thread implementation
return {"id": "custom-thread", "messages": []}
def create_session(self, **kwargs):
from agent_framework import AgentSession
return AgentSession(**kwargs)
def get_session(self, *, service_session_id, **kwargs):
from agent_framework import AgentSession
return AgentSession(service_session_id=service_session_id, **kwargs)
# Verify the instance satisfies the protocol
@@ -232,7 +233,7 @@ class SupportsAgentRun(Protocol):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]:
"""Get a response from the agent (non-streaming)."""
@@ -244,7 +245,7 @@ class SupportsAgentRun(Protocol):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Get a streaming response from the agent."""
@@ -255,7 +256,7 @@ class SupportsAgentRun(Protocol):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Get a response from the agent.
@@ -269,7 +270,7 @@ class SupportsAgentRun(Protocol):
Keyword Args:
stream: Whether to stream the response. Defaults to False.
thread: The conversation thread associated with the message(s).
session: The conversation session associated with the message(s).
kwargs: Additional keyword arguments.
Returns:
@@ -279,8 +280,12 @@ class SupportsAgentRun(Protocol):
"""
...
def get_new_thread(self, **kwargs: Any) -> AgentThread:
"""Creates a new conversation thread for the agent."""
def create_session(self, **kwargs: Any) -> AgentSession:
"""Creates a new conversation session."""
...
def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
"""Gets or creates a session for a service-managed session ID."""
...
@@ -294,7 +299,7 @@ class BaseAgent(SerializationMixin):
For most use cases, prefer :class:`Agent` which includes all standard layers.
This class provides core functionality for agent implementations, including
context providers, middleware support, and thread management.
context providers, middleware support, and session management.
Note:
BaseAgent cannot be instantiated directly as it doesn't implement the
@@ -304,12 +309,12 @@ class BaseAgent(SerializationMixin):
Examples:
.. code-block:: python
from agent_framework import BaseAgent, AgentThread, AgentResponse
from agent_framework import BaseAgent, AgentSession, AgentResponse
# Create a concrete subclass that implements the protocol
class SimpleAgent(BaseAgent):
async def run(self, messages=None, *, stream=False, thread=None, **kwargs):
async def run(self, messages=None, *, stream=False, session=None, **kwargs):
if stream:
async def _stream():
@@ -345,7 +350,7 @@ class BaseAgent(SerializationMixin):
id: str | None = None,
name: str | None = None,
description: str | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
additional_properties: MutableMapping[str, Any] | None = None,
**kwargs: Any,
@@ -357,7 +362,7 @@ class BaseAgent(SerializationMixin):
a new UUID will be generated.
name: The name of the agent, can be None.
description: The description of the agent.
context_provider: The context provider to include during agent invocation.
context_providers: Context providers to include during agent invocation.
middleware: List of middleware.
additional_properties: Additional properties set on the agent.
kwargs: Additional keyword arguments (merged into additional_properties).
@@ -367,7 +372,7 @@ class BaseAgent(SerializationMixin):
self.id = id
self.name = name
self.description = description
self.context_provider = context_provider
self.context_providers: list[BaseContextProvider] = list(context_providers or [])
self.middleware: list[MiddlewareTypes] | None = (
cast(list[MiddlewareTypes], middleware) if middleware is not None else None
)
@@ -376,56 +381,53 @@ class BaseAgent(SerializationMixin):
self.additional_properties: dict[str, Any] = cast(dict[str, Any], additional_properties or {})
self.additional_properties.update(kwargs)
async def _notify_thread_of_new_messages(
self,
thread: AgentThread,
input_messages: Message | Sequence[Message],
response_messages: Message | Sequence[Message],
**kwargs: Any,
) -> None:
"""Notify the thread of new messages.
This also calls the invoked method of a potential context provider on the thread.
Args:
thread: The thread to notify of new messages.
input_messages: The input messages to notify about.
response_messages: The response messages to notify about.
**kwargs: Any extra arguments to pass from the agent run.
"""
if isinstance(input_messages, Message) or len(input_messages) > 0:
await thread.on_new_messages(input_messages)
if isinstance(response_messages, Message) 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, **kwargs)
def get_new_thread(self, **kwargs: Any) -> AgentThread:
"""Return a new AgentThread instance that is compatible with the agent.
Keyword Args:
kwargs: Additional keyword arguments passed to AgentThread.
Returns:
A new AgentThread instance configured with the agent's context provider.
"""
return AgentThread(**kwargs, context_provider=self.context_provider)
async def deserialize_thread(self, serialized_thread: Any, **kwargs: Any) -> AgentThread:
"""Deserialize a thread from its serialized state.
Args:
serialized_thread: The serialized thread data.
def create_session(self, *, session_id: str | None = None, **kwargs: Any) -> AgentSession:
"""Create a new lightweight session.
Keyword Args:
session_id: Optional session ID (generated if not provided).
kwargs: Additional keyword arguments.
Returns:
A new AgentThread instance restored from the serialized state.
A new AgentSession instance.
"""
thread: AgentThread = self.get_new_thread()
await thread.update_from_thread_state(serialized_thread, **kwargs)
return thread
return AgentSession(session_id=session_id)
def get_session(self, *, service_session_id: str, session_id: str | None = None, **kwargs: Any) -> AgentSession:
"""Get or create a session for a service-managed session ID.
Args:
service_session_id: The service-managed session ID.
Keyword Args:
session_id: Optional local session ID (generated if not provided).
kwargs: Additional keyword arguments.
Returns:
A new AgentSession instance with service_session_id set.
"""
return AgentSession(session_id=session_id, service_session_id=service_session_id)
async def _run_after_providers(
self,
*,
session: AgentSession | None,
context: SessionContext,
) -> None:
"""Run after_run on all context providers in reverse order.
Keyword Args:
session: The conversation session.
context: The invocation context with response populated.
"""
state = session.state if session else {}
for provider in reversed(self.context_providers):
await provider.after_run(
agent=self, # type: ignore[arg-type]
session=session, # type: ignore[arg-type]
context=context,
state=state,
)
def as_tool(
self,
@@ -621,8 +623,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
| Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any]
| None = None,
default_options: OptionsCoT | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
**kwargs: Any,
) -> None:
"""Initialize a Agent instance.
@@ -636,9 +637,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
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_provider: The context providers to include during agent invocation.
context_providers: Context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
default_options: A TypedDict containing chat options. When using a typed agent like
``Agent[OpenAIChatOptions]``, this enables IDE autocomplete for
@@ -649,19 +648,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
These can be overridden at runtime via the ``options`` parameter of ``run()``.
tools: The tools to use for the request.
kwargs: Any additional keyword arguments. Will be stored as ``additional_properties``.
Raises:
AgentInitializationError: If both conversation_id and chat_message_store_factory are provided.
"""
# Extract conversation_id from options for validation
opts = dict(default_options) if default_options else {}
conversation_id = opts.get("conversation_id")
if conversation_id is not None and chat_message_store_factory is not None:
raise AgentInitializationError(
"Cannot specify both conversation_id and chat_message_store_factory. "
"Use conversation_id for service-managed threads or chat_message_store_factory for local storage."
)
if not isinstance(client, FunctionInvocationLayer) and isinstance(client, BaseChatClient):
logger.warning(
@@ -672,11 +660,10 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
id=id,
name=name,
description=description,
context_provider=context_provider,
context_providers=context_providers,
**kwargs,
)
self.client = client
self.chat_message_store_factory = chat_message_store_factory
# Get tools from options or named parameter (named param takes precedence)
tools_ = tools if tools is not None else opts.pop("tools", None)
@@ -704,7 +691,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
self.default_options: dict[str, Any] = {
"model_id": opts.pop("model_id", None) or (getattr(self.client, "model_id", None)),
"allow_multiple_tool_calls": opts.pop("allow_multiple_tool_calls", None),
"conversation_id": conversation_id,
"conversation_id": opts.pop("conversation_id", None),
"frequency_penalty": opts.pop("frequency_penalty", None),
"instructions": instructions_,
"logit_bias": opts.pop("logit_bias", None),
@@ -779,7 +766,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
tools: FunctionTool
| Callable[..., Any]
| MutableMapping[str, Any]
@@ -796,7 +783,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
tools: FunctionTool
| Callable[..., Any]
| MutableMapping[str, Any]
@@ -813,7 +800,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
tools: FunctionTool
| Callable[..., Any]
| MutableMapping[str, Any]
@@ -829,7 +816,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
tools: FunctionTool
| Callable[..., Any]
| MutableMapping[str, Any]
@@ -852,7 +839,9 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
stream: Whether to stream the response. Defaults to False.
Keyword Args:
thread: The thread to use for the agent.
session: The session to use for the agent.
If None, and no settings for the chat client that indicate otherwise,
the run will be stateless.
tools: The tools to use for this specific run (merged with default tools).
options: A TypedDict containing chat options. When using a typed agent like
``Agent[OpenAIChatOptions]``, this enables IDE autocomplete for
@@ -871,13 +860,13 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
async def _run_non_streaming() -> AgentResponse[Any]:
ctx = await self._prepare_run_context(
messages=messages,
thread=thread,
session=session,
tools=tools,
options=options,
kwargs=kwargs,
)
response = await self.client.get_response( # type: ignore[call-overload]
messages=ctx["thread_messages"],
messages=ctx["session_messages"],
stream=False,
options=ctx["chat_options"],
**ctx["filtered_kwargs"],
@@ -886,12 +875,11 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
if not response:
raise AgentExecutionException("Chat client did not return a response.")
await self._finalize_response_and_update_thread(
await self._finalize_response(
response=response,
agent_name=ctx["agent_name"],
thread=ctx["thread"],
input_messages=ctx["input_messages"],
kwargs=ctx["finalize_kwargs"],
session=ctx["session"],
session_context=ctx["session_context"],
)
response_format = ctx["chat_options"].get("response_format")
if not (
@@ -923,33 +911,39 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
if ctx is None:
return # No context available (shouldn't happen in normal flow)
# Update thread with conversation_id
await self._update_thread_with_type_and_conversation_id(ctx["thread"], response.response_id)
# Ensure author names are set for all messages
for message in response.messages:
if message.author_name is None:
message.author_name = ctx["agent_name"]
# Notify thread of new messages
await self._notify_thread_of_new_messages(
ctx["thread"],
ctx["input_messages"],
response.messages,
**{k: v for k, v in ctx["finalize_kwargs"].items() if k != "thread"},
# Propagate conversation_id back to session from streaming updates
sess = ctx["session"]
if sess and not sess.service_session_id and response.raw_representation:
raw_items = response.raw_representation if isinstance(response.raw_representation, list) else []
for item in raw_items:
if hasattr(item, "conversation_id") and item.conversation_id:
sess.service_session_id = item.conversation_id
break
# Run after_run providers (reverse order)
session_context = ctx["session_context"]
session_context._response = AgentResponse( # type: ignore[assignment]
messages=response.messages,
response_id=response.response_id,
)
await self._run_after_providers(session=ctx["session"], context=session_context)
async def _get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]:
ctx_holder["ctx"] = await self._prepare_run_context(
messages=messages,
thread=thread,
session=session,
tools=tools,
options=options,
kwargs=kwargs,
)
ctx: _RunContext = ctx_holder["ctx"] # type: ignore[assignment] # Safe: we just assigned it
return self.client.get_response( # type: ignore[call-overload, no-any-return]
messages=ctx["thread_messages"],
messages=ctx["session_messages"],
stream=True,
options=ctx["chat_options"],
**ctx["filtered_kwargs"],
@@ -984,7 +978,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
self,
*,
messages: str | Message | Sequence[str | Message] | None,
thread: AgentThread | None,
session: AgentSession | None,
tools: FunctionTool
| Callable[..., Any]
| MutableMapping[str, Any]
@@ -1000,8 +994,23 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
tools_ = tools if tools is not None else opts.pop("tools", None)
input_messages = normalize_messages(messages)
thread, run_chat_options, thread_messages = await self._prepare_thread_and_messages(
thread=thread, input_messages=input_messages, **kwargs
# Auto-inject InMemoryHistoryProvider when session is provided, no context providers
# registered, and no service-side storage indicators
if (
session is not None
and not self.context_providers
and not session.service_session_id
and not opts.get("conversation_id")
and not opts.get("store")
and not (getattr(self.client, "STORES_BY_DEFAULT", False) and opts.get("store") is not False)
):
self.context_providers.append(InMemoryHistoryProvider("memory"))
session_context, chat_options = await self._prepare_session_and_messages(
session=session,
input_messages=input_messages,
options=opts,
)
# Normalize tools
@@ -1028,7 +1037,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
# Build options dict from run() options merged with provided options
run_opts: dict[str, Any] = {
"model_id": opts.pop("model_id", None),
"conversation_id": thread.service_thread_id,
"conversation_id": session.service_session_id if session else opts.pop("conversation_id", None),
"allow_multiple_tool_calls": opts.pop("allow_multiple_tool_calls", None),
"additional_function_arguments": opts.pop("additional_function_arguments", None),
"frequency_penalty": opts.pop("frequency_penalty", None),
@@ -1049,103 +1058,129 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
}
# Remove None values and merge with chat_options
run_opts = {k: v for k, v in run_opts.items() if v is not None}
co = _merge_options(run_chat_options, run_opts)
co = _merge_options(chat_options, run_opts)
# Ensure thread is forwarded in kwargs for tool invocation
# Build session_messages from session context: context messages + input messages
session_messages: list[Message] = session_context.get_messages(include_input=True)
# Ensure session is forwarded in kwargs for tool invocation
finalize_kwargs = dict(kwargs)
finalize_kwargs["thread"] = thread
finalize_kwargs["session"] = session
# Filter chat_options from kwargs to prevent duplicate keyword argument
filtered_kwargs = {k: v for k, v in finalize_kwargs.items() if k != "chat_options"}
return {
"thread": thread,
"session": session,
"session_context": session_context,
"input_messages": input_messages,
"thread_messages": thread_messages,
"session_messages": session_messages,
"agent_name": agent_name,
"chat_options": co,
"filtered_kwargs": filtered_kwargs,
"finalize_kwargs": finalize_kwargs,
}
async def _finalize_response_and_update_thread(
async def _finalize_response(
self,
response: ChatResponse,
agent_name: str,
thread: AgentThread,
input_messages: list[Message],
kwargs: dict[str, Any],
session: AgentSession | None,
session_context: SessionContext,
) -> None:
"""Finalize response by updating thread and setting author names.
"""Finalize response by setting author names and running after_run providers.
Args:
response: The chat response to finalize.
agent_name: The name of the agent to set as author.
thread: The conversation thread.
input_messages: The input messages.
kwargs: Additional keyword arguments.
session: The conversation session.
session_context: The invocation context.
"""
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:
if message.author_name is None:
message.author_name = agent_name
# 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,
response.messages,
**{k: v for k, v in kwargs.items() if k != "thread"},
# Propagate conversation_id back to session (e.g. thread ID from Assistants API)
if session and response.conversation_id and not session.service_session_id:
session.service_session_id = response.conversation_id
# Set the response on the context for after_run providers
session_context._response = AgentResponse( # type: ignore[assignment]
messages=response.messages,
response_id=response.response_id,
)
@override
def get_new_thread(
# Run after_run providers (reverse order)
await self._run_after_providers(session=session, context=session_context)
async def _prepare_session_and_messages(
self,
*,
service_thread_id: str | None = None,
**kwargs: Any,
) -> AgentThread:
"""Get a new conversation thread for the agent.
session: AgentSession | None,
input_messages: list[Message] | None = None,
options: dict[str, Any] | None = None,
) -> tuple[SessionContext, dict[str, Any]]:
"""Prepare the session context and messages for agent execution.
If you supply a service_thread_id, the thread will be marked as service managed.
If you don't supply a service_thread_id but have a conversation_id configured on the agent,
that conversation_id will be used to create a service-managed thread.
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.
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 include a thread_id and that will be set.
Otherwise a message store is created from the default factory.
Runs the before_run pipeline on all context providers and assembles
the chat options from default options and provider-contributed context.
Keyword Args:
service_thread_id: Optional service managed thread ID.
kwargs: Not used at present.
session: The conversation session (None for stateless invocation).
input_messages: Messages to process.
options: Runtime options dict (already copied, safe to mutate).
Returns:
A new AgentThread instance.
A tuple containing:
- The SessionContext with provider context populated
- The merged chat options dict
"""
if service_thread_id is not None:
return AgentThread(
service_thread_id=service_thread_id,
context_provider=self.context_provider,
# Create a shallow copy of options and deep copy non-tool values
if self.default_options:
chat_options: dict[str, Any] = {}
for key, value in self.default_options.items():
if key == "tools":
chat_options[key] = list(value) if value else []
else:
chat_options[key] = deepcopy(value)
else:
chat_options = {}
session_context = SessionContext(
session_id=session.session_id if session else None,
service_session_id=session.service_session_id if session else None,
input_messages=input_messages or [],
options=options or {},
)
# Run before_run providers (forward order, skip BaseHistoryProvider with load_messages=False)
state = session.state if session else {}
for provider in self.context_providers:
if isinstance(provider, BaseHistoryProvider) and not provider.load_messages:
continue
await provider.before_run(
agent=self, # type: ignore[arg-type]
session=session, # type: ignore[arg-type]
context=session_context,
state=state,
)
if self.default_options.get("conversation_id") is not None:
return AgentThread(
service_thread_id=self.default_options["conversation_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)
# Merge provider-contributed tools into chat_options
if session_context.tools:
if chat_options.get("tools") is not None:
chat_options["tools"].extend(session_context.tools)
else:
chat_options["tools"] = list(session_context.tools)
# Merge provider-contributed instructions into chat_options
if session_context.instructions:
combined_instructions = "\n".join(session_context.instructions)
if "instructions" in chat_options:
chat_options["instructions"] = f"{chat_options['instructions']}\n{combined_instructions}"
else:
chat_options["instructions"] = combined_instructions
return session_context, chat_options
def as_mcp_server(
self,
@@ -1256,115 +1291,6 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
return server
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.
Args:
thread: The thread to update.
response_conversation_id: The conversation ID from the response, if any.
Raises:
AgentExecutionException: If conversation ID is missing for service-managed thread.
"""
if response_conversation_id is None and thread.service_thread_id is not None:
# We were passed a thread that is service managed, but we got no conversation id back from the chat client,
# meaning the service doesn't support service managed threads,
# so the thread cannot be used with this service.
raise AgentExecutionException(
"Service did not return a valid conversation id when using a service managed thread."
)
if response_conversation_id is not None:
# 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
# with the custom message_store so that it has somewhere to store the chat history.
thread.message_store = self.chat_message_store_factory()
async def _prepare_thread_and_messages(
self,
*,
thread: AgentThread | None,
input_messages: list[Message] | None = None,
**kwargs: Any,
) -> tuple[AgentThread, dict[str, Any], list[Message]]:
"""Prepare the thread and messages for agent execution.
This method prepares the conversation thread, merges context provider data,
and assembles the final message list for the chat client.
Keyword Args:
thread: The conversation thread.
input_messages: Messages to process.
**kwargs: Any extra arguments to pass from the agent run.
Returns:
A tuple containing:
- The validated or created thread
- The merged chat options
- The complete list of messages for the chat client
Raises:
AgentExecutionException: If the conversation IDs on the thread and agent don't match.
"""
# Create a shallow copy of options and deep copy non-tool values
# Tools containing HTTP clients or other non-copyable objects cannot be deep copied
if self.default_options:
chat_options: dict[str, Any] = {}
for key, value in self.default_options.items():
if key == "tools":
# Keep tool references as-is (don't deep copy)
chat_options[key] = list(value) if value else []
else:
# Deep copy other options to prevent mutation
chat_options[key] = deepcopy(value)
else:
chat_options = {}
thread = thread or self.get_new_thread()
if thread.service_thread_id and thread.context_provider:
await thread.context_provider.thread_created(thread.service_thread_id)
thread_messages: list[Message] = []
if thread.message_store:
thread_messages.extend(await thread.message_store.list_messages() or [])
context: Context | None = None
if self.context_provider:
# Note: We don't use 'async with' here because the context provider's lifecycle
# should be managed by the user (via async with) or persist across multiple invocations.
# Using async with here would close resources (like retrieval clients) after each query.
context = await self.context_provider.invoking(input_messages or [], **kwargs)
if context:
if context.messages:
thread_messages.extend(context.messages)
if context.tools:
if chat_options.get("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 "instructions" not in chat_options
else f"{chat_options['instructions']}\n{context.instructions}"
)
thread_messages.extend(input_messages or [])
if (
thread.service_thread_id
and chat_options.get("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:
"""Get the agent name for message attribution.
@@ -1404,8 +1330,7 @@ class Agent(
| Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any]
| None = None,
default_options: OptionsCoT | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> None:
@@ -1418,8 +1343,7 @@ class Agent(
description=description,
tools=tools,
default_options=default_options,
chat_message_store_factory=chat_message_store_factory,
context_provider=context_provider,
context_providers=context_providers,
middleware=middleware,
**kwargs,
)
python/packages/core/agent_framework/_clients.py
+12 -10
View File
@@ -28,9 +28,7 @@ from typing import (
from pydantic import BaseModel
from ._logging import get_logger
from ._memory import ContextProvider
from ._serialization import SerializationMixin
from ._threads import ChatMessageStoreProtocol
from ._tools import (
FunctionInvocationConfiguration,
FunctionTool,
@@ -264,7 +262,15 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
OTEL_PROVIDER_NAME: ClassVar[str] = "unknown"
DEFAULT_EXCLUDE: ClassVar[set[str]] = {"additional_properties"}
# This is used for OTel setup, should be overridden in subclasses
STORES_BY_DEFAULT: ClassVar[bool] = False
"""Whether this client stores conversation history server-side by default.
Clients that use server-side storage (e.g., OpenAI Responses API with ``store=True``
as default, Azure AI Agent sessions) should override this to ``True``.
When ``True``, the agent skips auto-injecting ``InMemoryHistoryProvider`` unless the
user explicitly sets ``store=False``.
"""
# OTEL_PROVIDER_NAME is used for OTel setup, should be overridden in subclasses
def __init__(
self,
@@ -448,8 +454,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
| Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
default_options: OptionsCoT | Mapping[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[Any] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
@@ -471,9 +476,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
including temperature, max_tokens, model_id, tool_choice, and more.
Note: response_format typing does not flow into run outputs when set via default_options,
and dict literals are accepted without specialized option typing.
chat_message_store_factory: Factory function to create an instance of ChatMessageStoreProtocol.
If not provided, the default in-memory store will be used.
context_provider: Context providers to include during agent invocation.
context_providers: Context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
function_invocation_configuration: Optional function invocation configuration override.
kwargs: Any additional keyword arguments. Will be stored as ``additional_properties``.
@@ -509,8 +512,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
instructions=instructions,
tools=tools,
default_options=cast(Any, default_options),
chat_message_store_factory=chat_message_store_factory,
context_provider=context_provider,
context_providers=context_providers,
middleware=middleware,
function_invocation_configuration=function_invocation_configuration,
**kwargs,
python/packages/core/agent_framework/_mcp.py
+51 -26
View File
@@ -102,21 +102,31 @@ def _parse_prompt_result_from_mcp(
if isinstance(content, types.TextContent):
parts.append(content.text)
elif isinstance(content, (types.ImageContent, types.AudioContent)):
parts.append(json.dumps({
"type": "image" if isinstance(content, types.ImageContent) else "audio",
"data": content.data,
"mimeType": content.mimeType,
}, default=str))
parts.append(
json.dumps(
{
"type": "image" if isinstance(content, types.ImageContent) else "audio",
"data": content.data,
"mimeType": content.mimeType,
},
default=str,
)
)
elif isinstance(content, types.EmbeddedResource):
match content.resource:
case types.TextResourceContents():
parts.append(content.resource.text)
case types.BlobResourceContents():
parts.append(json.dumps({
"type": "blob",
"data": content.resource.blob,
"mimeType": content.resource.mimeType,
}, default=str))
parts.append(
json.dumps(
{
"type": "blob",
"data": content.resource.blob,
"mimeType": content.resource.mimeType,
},
default=str,
)
)
else:
parts.append(str(content))
if not parts:
@@ -159,27 +169,42 @@ def _parse_tool_result_from_mcp(
case types.TextContent():
parts.append(item.text)
case types.ImageContent() | types.AudioContent():
parts.append(json.dumps({
"type": "image" if isinstance(item, types.ImageContent) else "audio",
"data": item.data,
"mimeType": item.mimeType,
}, default=str))
parts.append(
json.dumps(
{
"type": "image" if isinstance(item, types.ImageContent) else "audio",
"data": item.data,
"mimeType": item.mimeType,
},
default=str,
)
)
case types.ResourceLink():
parts.append(json.dumps({
"type": "resource_link",
"uri": str(item.uri),
"mimeType": item.mimeType,
}, default=str))
parts.append(
json.dumps(
{
"type": "resource_link",
"uri": str(item.uri),
"mimeType": item.mimeType,
},
default=str,
)
)
case types.EmbeddedResource():
match item.resource:
case types.TextResourceContents():
parts.append(item.resource.text)
case types.BlobResourceContents():
parts.append(json.dumps({
"type": "blob",
"data": item.resource.blob,
"mimeType": item.resource.mimeType,
}, default=str))
parts.append(
json.dumps(
{
"type": "blob",
"data": item.resource.blob,
"mimeType": item.resource.mimeType,
},
default=str,
)
)
case _:
parts.append(str(item))
if not parts:
@@ -847,7 +872,7 @@ class MCPTool:
k: v
for k, v in kwargs.items()
if k
not in {"chat_options", "tools", "tool_choice", "thread", "conversation_id", "options", "response_format"}
not in {"chat_options", "tools", "tool_choice", "session", "thread", "conversation_id", "options", "response_format"}
}
parser = self.parse_tool_results or _parse_tool_result_from_mcp
python/packages/core/agent_framework/_memory.py
-181
View File
@@ -1,181 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import sys
from abc import ABC, abstractmethod
from collections.abc import MutableSequence, Sequence
from types import TracebackType
from typing import TYPE_CHECKING, Any, Final
from ._types import Message
if TYPE_CHECKING:
from ._tools import FunctionTool
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
else:
from typing_extensions import Self # pragma: no cover
# region Context
__all__ = ["Context", "ContextProvider"]
class Context:
"""A class containing any context that should be provided to the AI model as supplied by a ContextProvider.
Each ContextProvider has the ability to provide its own context for each invocation.
The Context class contains the additional context supplied by the ContextProvider.
This context will be combined with context supplied by other providers before being passed to the AI model.
This context is per invocation, and will not be stored as part of the chat history.
Examples:
.. code-block:: python
from agent_framework import Context, Message
# Create context with instructions
context = Context(
instructions="Use a professional tone when responding.",
messages=[Message(content="Previous context", role="user")],
tools=[my_tool],
)
# Access context properties
print(context.instructions)
print(len(context.messages))
"""
def __init__(
self,
instructions: str | None = None,
messages: Sequence[Message] | None = None,
tools: Sequence[FunctionTool] | None = None,
):
"""Create a new Context object.
Args:
instructions: The instructions to provide to the AI model.
messages: The list of messages to include in the context.
tools: The list of tools to provide to this run.
"""
self.instructions = instructions
self.messages: Sequence[Message] = messages or []
self.tools: Sequence[FunctionTool] = tools or []
# region ContextProvider
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.
Note:
ContextProvider is an abstract base class. You must subclass it and implement
the ``invoking()`` method to create a custom context provider. Ideally, you should
also implement the ``invoked()`` and ``thread_created()`` methods to track conversation
state, but these are optional.
Examples:
.. code-block:: python
from agent_framework import ContextProvider, Context, Message
class CustomContextProvider(ContextProvider):
async def invoking(self, messages, **kwargs):
# Add custom instructions before each invocation
return Context(instructions="Always be concise and helpful.", messages=[], tools=[])
# Use with a chat agent
async with CustomContextProvider() as provider:
agent = Agent(client=client, name="assistant", context_provider=provider)
"""
# Default prompt to be used by all context providers when assembling memories/instructions
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.
Implementers can use this method to perform any operations required at the creation
of a new thread. For example, checking long-term storage for any data that is relevant
to the current session.
Args:
thread_id: The ID of the new thread.
"""
pass
async def invoked(
self,
request_messages: Message | Sequence[Message],
response_messages: Message | Sequence[Message] | None = None,
invoke_exception: Exception | None = None,
**kwargs: Any,
) -> None:
"""Called after the agent has received a response from the underlying inference service.
You can inspect the request and response messages, and update the state of the context provider.
Args:
request_messages: The messages that were sent to the model/agent.
response_messages: The messages that were returned by the model/agent.
invoke_exception: The exception that was thrown, if any.
Keyword Args:
kwargs: Additional keyword arguments (not used at present).
"""
pass
@abstractmethod
async def invoking(self, messages: Message | MutableSequence[Message], **kwargs: Any) -> Context:
"""Called just before the model/agent is invoked.
Implementers can load any additional context required at this time,
and they should return any context that should be passed to the agent.
Args:
messages: The most recent messages that the agent is being invoked with.
Keyword Args:
kwargs: Additional keyword arguments (not used at present).
Returns:
A Context object containing instructions, messages, and tools to include.
"""
pass
async def __aenter__(self) -> Self:
"""Enter the async context manager.
Override this method to perform any setup operations when the context provider is entered.
Returns:
The ContextProvider instance for chaining.
"""
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
"""Exit the async context manager.
Override this method to perform any cleanup operations when the context provider is exited.
Args:
exc_type: The exception type if an exception occurred, None otherwise.
exc_val: The exception value if an exception occurred, None otherwise.
exc_tb: The exception traceback if an exception occurred, None otherwise.
"""
pass
python/packages/core/agent_framework/_middleware.py
+13 -13
View File
@@ -36,7 +36,7 @@ if TYPE_CHECKING:
from ._agents import SupportsAgentRun
from ._clients import SupportsChatGetResponse
from ._threads import AgentThread
from ._sessions import AgentSession
from ._tools import FunctionTool
from ._types import ChatOptions, ChatResponse, ChatResponseUpdate
@@ -118,7 +118,7 @@ class AgentContext:
Attributes:
agent: The agent being invoked.
messages: The messages being sent to the agent.
thread: The agent thread for this invocation, if any.
session: The agent session for this invocation, if any.
options: The options for the agent invocation as a dict.
stream: Whether this is a streaming invocation.
metadata: Metadata dictionary for sharing data between agent middleware.
@@ -138,7 +138,7 @@ class AgentContext:
async def process(self, context: AgentContext, call_next):
print(f"Agent: {context.agent.name}")
print(f"Messages: {len(context.messages)}")
print(f"Thread: {context.thread}")
print(f"Session: {context.session}")
print(f"Streaming: {context.stream}")
# Store metadata
@@ -156,7 +156,7 @@ class AgentContext:
*,
agent: SupportsAgentRun,
messages: list[Message],
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: Mapping[str, Any] | None = None,
stream: bool = False,
metadata: Mapping[str, Any] | None = None,
@@ -175,7 +175,7 @@ class AgentContext:
Args:
agent: The agent being invoked.
messages: The messages being sent to the agent.
thread: The agent thread for this invocation, if any.
session: The agent session for this invocation, if any.
options: The options for the agent invocation as a dict.
stream: Whether this is a streaming invocation.
metadata: Metadata dictionary for sharing data between agent middleware.
@@ -187,7 +187,7 @@ class AgentContext:
"""
self.agent = agent
self.messages = messages
self.thread = thread
self.session = session
self.options = options
self.stream = stream
self.metadata = metadata if metadata is not None else {}
@@ -1098,7 +1098,7 @@ class AgentMiddlewareLayer:
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
options: ChatOptions[ResponseModelBoundT],
**kwargs: Any,
@@ -1110,7 +1110,7 @@ class AgentMiddlewareLayer:
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
options: ChatOptions[None] | None = None,
**kwargs: Any,
@@ -1122,7 +1122,7 @@ class AgentMiddlewareLayer:
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
options: ChatOptions[Any] | None = None,
**kwargs: Any,
@@ -1133,7 +1133,7 @@ class AgentMiddlewareLayer:
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
options: ChatOptions[Any] | None = None,
**kwargs: Any,
@@ -1157,12 +1157,12 @@ class AgentMiddlewareLayer:
# Execute with middleware if available
if not pipeline.has_middlewares:
return super().run(messages, stream=stream, thread=thread, options=options, **combined_kwargs) # type: ignore[misc, no-any-return]
return super().run(messages, stream=stream, session=session, options=options, **combined_kwargs) # type: ignore[misc, no-any-return]
context = AgentContext(
agent=self, # type: ignore[arg-type]
messages=prepare_messages(messages), # type: ignore[arg-type]
thread=thread,
session=session,
options=options,
stream=stream,
kwargs=combined_kwargs,
@@ -1197,7 +1197,7 @@ class AgentMiddlewareLayer:
return super().run( # type: ignore[misc, no-any-return]
context.messages,
stream=context.stream,
thread=context.thread,
session=context.session,
options=context.options,
**context.kwargs,
)
python/packages/core/agent_framework/_serialization.py
+8 -31
View File
@@ -166,45 +166,22 @@ class SerializationMixin:
during deserialization via the ``dependencies`` parameter.
Examples:
**Nested object serialization with agent thread management:**
**Nested object serialization:**
.. code-block:: python
from agent_framework import Message
from agent_framework._threads import AgentThreadState, ChatMessageStoreState
from agent_framework._sessions import AgentSession
# ChatMessageStoreState handles nested Message serialization
store_state = ChatMessageStoreState(
messages=[
Message(role="user", text="Hello agent"),
Message(role="assistant", text="Hi! How can I help?"),
]
)
# AgentSession uses SerializationMixin for state serialization
session = AgentSession(session_id="test")
# Nested serialization: messages are automatically converted to dicts
store_dict = store_state.to_dict()
# Result: {
# "type": "chat_message_store_state",
# "messages": [
# {"type": "chat_message", "role": {...}, "contents": [...]},
# {"type": "chat_message", "role": {...}, "contents": [...]}
# ]
# }
# Serialization produces a clean dict representation
session_dict = session.to_dict()
# AgentThreadState contains nested ChatMessageStoreState
thread_state = AgentThreadState(chat_message_store_state=store_state)
# Deep serialization: nested SerializationMixin objects are handled automatically
thread_dict = thread_state.to_dict()
# The chat_message_store_state and its nested messages are all serialized
# Reconstruction from nested dictionaries with automatic type conversion
# The __init__ method handles MutableMapping -> object conversion:
reconstructed = AgentThreadState.from_dict({
"chat_message_store_state": {"messages": [{"role": "user", "text": "Hello again"}]}
})
# chat_message_store_state becomes ChatMessageStoreState instance automatically
# Reconstruction from dictionaries
restored = AgentSession.from_dict(session_dict)
**Framework tools with exclusion patterns:**
python/packages/core/agent_framework/_sessions.py
+50 -5
View File
@@ -30,6 +30,7 @@ __all__ = [
"BaseHistoryProvider",
"InMemoryHistoryProvider",
"SessionContext",
"register_state_type",
]
@@ -37,16 +38,50 @@ __all__ = [
_STATE_TYPE_REGISTRY: dict[str, type] = {}
def _register_state_type(cls: type) -> None:
"""Register a type for automatic deserialization in session state."""
def register_state_type(cls: type) -> None:
"""Register a type for automatic deserialization in session state.
Call this for any custom type (including Pydantic models) that you store
in ``session.state`` and want to survive ``to_dict()`` / ``from_dict()``
round-trips. Types with ``to_dict``/``from_dict`` methods or Pydantic
``BaseModel`` subclasses are handled automatically.
The type identifier defaults to ``cls.__name__.lower()`` but can be
overridden by defining a ``_get_type_identifier`` classmethod.
Note:
Pydantic models are auto-registered on first serialization, but
pre-registering ensures deserialization works even if the model
hasn't been serialized in this process yet (e.g. cold-start restore).
Args:
cls: The type to register.
"""
type_id: str = getattr(cls, "_get_type_identifier", lambda: cls.__name__.lower())()
_STATE_TYPE_REGISTRY[type_id] = cls
# Keep internal alias for framework use
_register_state_type = register_state_type
def _serialize_value(value: Any) -> Any:
"""Serialize a single value, handling objects with to_dict()."""
"""Serialize a single value, handling objects with to_dict() and Pydantic models."""
if hasattr(value, "to_dict") and callable(value.to_dict):
return value.to_dict() # pyright: ignore[reportUnknownMemberType]
# Pydantic BaseModel support — import lazily to avoid hard dep at module level
try:
from pydantic import BaseModel
if isinstance(value, BaseModel):
data = value.model_dump()
type_id: str = getattr(value.__class__, "_get_type_identifier", lambda: value.__class__.__name__.lower())()
data["type"] = type_id
# Auto-register for round-trip deserialization
_STATE_TYPE_REGISTRY.setdefault(type_id, value.__class__)
return data
except ImportError:
pass
if isinstance(value, list):
return [_serialize_value(item) for item in value] # pyright: ignore[reportUnknownVariableType]
if isinstance(value, dict):
@@ -59,8 +94,18 @@ def _deserialize_value(value: Any) -> Any:
if isinstance(value, dict) and "type" in value:
type_id = str(value["type"]) # pyright: ignore[reportUnknownArgumentType]
cls = _STATE_TYPE_REGISTRY.get(type_id)
if cls is not None and hasattr(cls, "from_dict"):
return cls.from_dict(value) # type: ignore[union-attr]
if cls is not None:
if hasattr(cls, "from_dict"):
return cls.from_dict(value) # type: ignore[union-attr]
# Pydantic BaseModel support
try:
from pydantic import BaseModel
if issubclass(cls, BaseModel):
data = {k: v for k, v in value.items() if k != "type"}
return cls.model_validate(data)
except ImportError:
pass
if isinstance(value, list):
return [_deserialize_value(item) for item in value] # pyright: ignore[reportUnknownVariableType]
if isinstance(value, dict):
python/packages/core/agent_framework/_settings.py
+38 -20
View File
@@ -12,14 +12,17 @@ Usage::
class MySettings(TypedDict, total=False):
api_key: str | None # optional — resolves to None if not set
model_id: str | None # optional by default
source_a: str | None
source_b: str | None
# Make model_id required at call time:
# Make model_id required; require exactly one of source_a / source_b:
settings = load_settings(
MySettings,
env_prefix="MY_APP_",
required_fields=["model_id"],
required_fields=["model_id", ("source_a", "source_b")],
model_id="gpt-4",
source_a="value",
)
settings["api_key"] # type-checked dict access
settings["model_id"] # str | None per type, but guaranteed not None at runtime
@@ -167,7 +170,7 @@ def load_settings(
env_prefix: str = "",
env_file_path: str | None = None,
env_file_encoding: str | None = None,
required_fields: Sequence[str] | None = None,
required_fields: Sequence[str | tuple[str, ...]] | None = None,
**overrides: Any,
) -> SettingsT:
"""Load settings from environment variables, a ``.env`` file, and explicit overrides.
@@ -181,18 +184,19 @@ def load_settings(
4. Default values — fields with class-level defaults on the TypedDict, or
``None`` for optional fields.
Fields listed in *required_fields* are validated after resolution. If any
required field resolves to ``None``, a ``SettingNotFoundError`` is raised.
This allows callers to decide which fields are required based on runtime
context (e.g. ``endpoint`` is only required when no pre-built client is
provided).
Entries in *required_fields* are validated after resolution:
- A **string** entry means the field must resolve to a non-``None`` value.
- A **tuple** entry means exactly one field in the group must be non-``None``
(mutually exclusive).
Args:
settings_type: A ``TypedDict`` class describing the settings schema.
env_prefix: Prefix for environment variable lookup (e.g. ``"OPENAI_"``).
env_file_path: Path to ``.env`` file. Defaults to ``".env"`` when omitted.
env_file_encoding: Encoding for reading the ``.env`` file. Defaults to ``"utf-8"``.
required_fields: Field names that must resolve to a non-``None`` value.
required_fields: Field names (``str``) that must resolve to a non-``None``
value, or tuples of field names where exactly one must be set.
**overrides: Field values. ``None`` values are ignored so that callers can
forward optional parameters without masking env-var / default resolution.
@@ -200,7 +204,8 @@ def load_settings(
A populated dict matching *settings_type*.
Raises:
SettingNotFoundError: If a required field could not be resolved from any source.
SettingNotFoundError: If a required field could not be resolved from any
source, or if a mutually exclusive constraint is violated.
ServiceInitializationError: If an override value has an incompatible type.
"""
encoding = env_file_encoding or "utf-8"
@@ -215,7 +220,6 @@ def load_settings(
# Get field type hints from the TypedDict
hints = get_type_hints(settings_type)
required: set[str] = set(required_fields) if required_fields else set()
result: dict[str, Any] = {}
for field_name, field_type in hints.items():
@@ -249,14 +253,28 @@ def load_settings(
result[field_name] = None
# Validate required fields after all resolution
if required:
for field_name in required:
if result.get(field_name) is None:
env_var_name = f"{env_prefix}{field_name.upper()}"
raise SettingNotFoundError(
f"Required setting '{field_name}' was not provided. "
f"Set it via the '{field_name}' parameter or the "
f"'{env_var_name}' environment variable."
)
if required_fields:
for entry in required_fields:
if isinstance(entry, str):
# Single required field
if result.get(entry) is None:
env_var_name = f"{env_prefix}{entry.upper()}"
raise SettingNotFoundError(
f"Required setting '{entry}' was not provided. "
f"Set it via the '{entry}' parameter or the "
f"'{env_var_name}' environment variable."
)
else:
# Mutually exclusive group — exactly one must be set
set_fields = [f for f in entry if result.get(f) is not None]
if len(set_fields) == 0:
names = ", ".join(f"'{f}'" for f in entry)
raise SettingNotFoundError(f"Exactly one of {names} must be provided, but none was set.")
if len(set_fields) > 1:
all_names = ", ".join(f"'{f}'" for f in entry)
set_names = ", ".join(f"'{f}'" for f in set_fields)
raise SettingNotFoundError(
f"Only one of {all_names} may be provided, but multiple were set: {set_names}."
)
return result # type: ignore[return-value]
python/packages/core/agent_framework/_threads.py
-507
View File
@@ -1,507 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
from collections.abc import MutableMapping, Sequence
from typing import Any, Protocol, TypeVar
from ._memory import ContextProvider
from ._serialization import SerializationMixin
from ._types import Message
from .exceptions import AgentThreadException
__all__ = ["AgentThread", "ChatMessageStore", "ChatMessageStoreProtocol"]
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,
including handling large volumes of data by truncating or summarizing messages as necessary.
Examples:
.. code-block:: python
from agent_framework import Message
class MyMessageStore:
def __init__(self):
self._messages = []
async def list_messages(self) -> list[Message]:
return self._messages
async def add_messages(self, messages: Sequence[Message]) -> None:
self._messages.extend(messages)
@classmethod
async def deserialize(cls, serialized_store_state, **kwargs):
store = cls()
store._messages = serialized_store_state.get("messages", [])
return store
async def update_from_state(self, serialized_store_state, **kwargs) -> None:
self._messages = serialized_store_state.get("messages", [])
async def serialize(self, **kwargs):
return {"messages": self._messages}
# Use the custom store
store = MyMessageStore()
"""
async def list_messages(self) -> list[Message]:
"""Gets all the messages from the store that should be used for the next agent invocation.
Messages are returned in ascending chronological order, with the oldest message first.
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 ``ChatMessageStoreProtocol``, a new one should be created for each thread
since they may contain state that is specific to a thread.
"""
...
async def add_messages(self, messages: Sequence[Message]) -> None:
"""Adds messages to the store.
Args:
messages: The sequence of Message objects to add to the store.
"""
...
@classmethod
async def deserialize(
cls, serialized_store_state: MutableMapping[str, Any], **kwargs: Any
) -> ChatMessageStoreProtocol:
"""Creates a new instance of the store from previously serialized state.
This method, together with ``serialize()`` can be used to save and load messages from a persistent store
if this store only has messages in memory.
Args:
serialized_store_state: The previously serialized state data containing messages.
Keyword Args:
**kwargs: Additional arguments for deserialization.
Returns:
A new instance of the store populated with messages from the serialized state.
"""
...
async def update_from_state(self, serialized_store_state: MutableMapping[str, Any], **kwargs: Any) -> None:
"""Update the current ChatMessageStore instance from serialized state data.
Args:
serialized_store_state: Previously serialized state data containing messages.
Keyword Args:
kwargs: Additional arguments for deserialization.
"""
...
async def serialize(self, **kwargs: Any) -> dict[str, Any]:
"""Serializes the current object's state.
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.
Keyword Args:
kwargs: Additional arguments for serialization.
Returns:
The serialized state data that can be used with ``deserialize()``.
"""
...
class ChatMessageStoreState(SerializationMixin):
"""State model for serializing and deserializing chat message store data.
Attributes:
messages: List of chat messages stored in the message store.
"""
def __init__(
self,
messages: Sequence[Message] | Sequence[MutableMapping[str, Any]] | None = None,
**kwargs: Any,
) -> None:
"""Create the store state.
Args:
messages: a list of messages or a list of the dict representation of messages.
Keyword Args:
**kwargs: not used for this, but might be used by subclasses.
"""
if not messages:
self.messages: list[Message] = []
return
if not isinstance(messages, list):
raise TypeError("Messages should be a list")
new_messages: list[Message] = []
for msg in messages:
if isinstance(msg, Message):
new_messages.append(msg)
else:
new_messages.append(Message.from_dict(msg))
self.messages = new_messages
class AgentThreadState(SerializationMixin):
"""State model for serializing and deserializing thread information."""
def __init__(
self,
*,
service_thread_id: str | None = None,
chat_message_store_state: ChatMessageStoreState | MutableMapping[str, Any] | None = None,
) -> None:
"""Create a AgentThread state.
Keyword Args:
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.
"""
if service_thread_id is not None and chat_message_store_state is not None:
raise AgentThreadException("A thread cannot have both a service_thread_id and a chat_message_store.")
self.service_thread_id = service_thread_id
self.chat_message_store_state: ChatMessageStoreState | None = None
if chat_message_store_state is not None:
if isinstance(chat_message_store_state, dict):
self.chat_message_store_state = ChatMessageStoreState.from_dict(chat_message_store_state)
elif isinstance(chat_message_store_state, ChatMessageStoreState):
self.chat_message_store_state = chat_message_store_state
else:
raise TypeError("Could not parse ChatMessageStoreState.")
ChatMessageStoreT = TypeVar("ChatMessageStoreT", 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.
Examples:
.. code-block:: python
from agent_framework import ChatMessageStore, Message
# Create an empty store
store = ChatMessageStore()
# Add messages
message = Message(role="user", text="Hello")
await store.add_messages([message])
# Retrieve messages
messages = await store.list_messages()
# Serialize for persistence
state = await store.serialize()
# Deserialize from saved state
restored_store = await ChatMessageStore.deserialize(state)
"""
def __init__(self, messages: Sequence[Message] | None = None):
"""Create a ChatMessageStore for use in a thread.
Args:
messages: The messages to store.
"""
self.messages = list(messages) if messages else []
async def add_messages(self, messages: Sequence[Message]) -> None:
"""Add messages to the store.
Args:
messages: Sequence of Message objects to add to the store.
"""
self.messages.extend(messages)
async def list_messages(self) -> list[Message]:
"""Get all messages from the store in chronological order.
Returns:
List of Message objects, ordered from oldest to newest.
"""
return self.messages
@classmethod
async def deserialize(
cls: type[ChatMessageStoreT], serialized_store_state: MutableMapping[str, Any], **kwargs: Any
) -> ChatMessageStoreT:
"""Create a new ChatMessageStore instance from serialized state data.
Args:
serialized_store_state: Previously serialized state data containing messages.
Keyword Args:
**kwargs: Additional arguments for deserialization.
Returns:
A new ChatMessageStore instance populated with messages from the serialized state.
"""
state = ChatMessageStoreState.from_dict(serialized_store_state, **kwargs)
if state.messages:
return cls(messages=state.messages)
return cls()
async def update_from_state(self, serialized_store_state: MutableMapping[str, Any], **kwargs: Any) -> None:
"""Update the current ChatMessageStore instance from serialized state data.
Args:
serialized_store_state: Previously serialized state data containing messages.
Keyword Args:
**kwargs: Additional arguments for deserialization.
"""
if not serialized_store_state:
return
state = ChatMessageStoreState.from_dict(serialized_store_state, **kwargs)
if state.messages:
self.messages = state.messages
async def serialize(self, **kwargs: Any) -> dict[str, Any]:
"""Serialize the current store state for persistence.
Keyword Args:
**kwargs: Additional arguments for serialization.
Returns:
Serialized state data that can be used with deserialize_state.
"""
state = ChatMessageStoreState(messages=self.messages)
return state.to_dict()
AgentThreadT = TypeVar("AgentThreadT", bound="AgentThread")
class AgentThread:
"""The Agent thread class, this can represent both a locally managed thread or a thread managed by the service.
An ``AgentThread`` maintains the conversation state and message history for an agent interaction.
It can either use a service-managed thread (via ``service_thread_id``) or a local message store
(via ``message_store``), but not both.
Examples:
.. code-block:: python
from agent_framework import Agent, ChatMessageStore
from agent_framework.openai import OpenAIChatClient
client = OpenAIChatClient(model="gpt-4o")
# Create agent with service-managed threads using a service_thread_id
service_agent = Agent(name="assistant", client=client)
service_thread = await service_agent.get_new_thread(service_thread_id="thread_abc123")
# Create agent with service-managed threads using conversation_id
conversation_agent = Agent(name="assistant", client=client, conversation_id="thread_abc123")
conversation_thread = await conversation_agent.get_new_thread()
# Create agent with custom message store factory
local_agent = Agent(name="assistant", client=client, chat_message_store_factory=ChatMessageStore)
local_thread = await local_agent.get_new_thread()
# Serialize and restore thread state
state = await local_thread.serialize()
restored_thread = await local_agent.deserialize_thread(state)
"""
def __init__(
self,
*,
service_thread_id: str | None = None,
message_store: ChatMessageStoreProtocol | None = None,
context_provider: ContextProvider | None = None,
) -> None:
"""Initialize an AgentThread, do not use this method manually, always use: ``agent.get_new_thread()``.
Args:
service_thread_id: The optional ID of the thread managed by the agent service.
message_store: The optional ChatMessageStore implementation for managing chat messages.
context_provider: The optional ContextProvider for the thread.
Note:
Either ``service_thread_id`` or ``message_store`` may be set, but not both.
"""
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.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:
"""Gets the ID of the current thread to support cases where the thread is owned by the agent service."""
return self._service_thread_id
@service_thread_id.setter
def service_thread_id(self, service_thread_id: str | None) -> None:
"""Sets the ID of the current thread to support cases where the thread is owned by the agent service.
Note:
Either ``service_thread_id`` or ``message_store`` may be set, but not both.
"""
if service_thread_id is None:
return
if self._message_store 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._service_thread_id = service_thread_id
@property
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: ChatMessageStoreProtocol | None) -> None:
"""Sets the ``ChatMessageStoreProtocol`` used by this thread.
Note:
Either ``service_thread_id`` or ``message_store`` may be set, but not both.
"""
if message_store is None:
return
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: Message | Sequence[Message]) -> None:
"""Invoked when a new message has been contributed to the chat by any participant.
Args:
new_messages: The new Message or sequence of Message objects to add to the thread.
"""
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, Message):
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.
Keyword Args:
**kwargs: Arguments for serialization.
"""
chat_message_store_state = None
if self._message_store is not None:
chat_message_store_state = await self._message_store.serialize(**kwargs)
state = AgentThreadState(
service_thread_id=self._service_thread_id, chat_message_store_state=chat_message_store_state
)
return state.to_dict(exclude_none=False)
@classmethod
async def deserialize(
cls: type[AgentThreadT],
serialized_thread_state: MutableMapping[str, Any],
*,
message_store: ChatMessageStoreProtocol | None = None,
**kwargs: Any,
) -> AgentThreadT:
"""Deserializes the state from a dictionary into a new AgentThread instance.
Args:
serialized_thread_state: The serialized thread state as a dictionary.
Keyword Args:
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.
Returns:
A new AgentThread instance with properties set from the serialized state.
"""
state = AgentThreadState.from_dict(serialized_thread_state)
if state.service_thread_id is not None:
return cls(service_thread_id=state.service_thread_id)
# If we don't have any ChatMessageStoreProtocol state return here.
if state.chat_message_store_state is None:
return cls()
if message_store is not None:
try:
await message_store.add_messages(state.chat_message_store_state.messages, **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 = ChatMessageStore(messages=state.chat_message_store_state.messages, **kwargs)
except Exception as ex:
raise AgentThreadException("Failed to deserialize the message store.") from ex
return cls(message_store=message_store)
async def update_from_thread_state(
self,
serialized_thread_state: MutableMapping[str, Any],
**kwargs: Any,
) -> None:
"""Deserializes the state from a dictionary into the thread properties.
Args:
serialized_thread_state: The serialized thread state as a dictionary.
Keyword Args:
**kwargs: Additional arguments for deserialization.
"""
state = AgentThreadState.from_dict(serialized_thread_state)
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.add_messages(state.chat_message_store_state.messages, **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 = ChatMessageStore(messages=state.chat_message_store_state.messages, **kwargs)
python/packages/core/agent_framework/_tools.py
+2 -2
View File
@@ -468,7 +468,7 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]):
"chat_options",
"tools",
"tool_choice",
"thread",
"session",
"conversation_id",
"options",
"response_format",
@@ -1897,7 +1897,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
config=self.function_invocation_configuration,
middleware_pipeline=function_middleware_pipeline,
)
filtered_kwargs = {k: v for k, v in kwargs.items() if k != "thread"}
filtered_kwargs = {k: v for k, v in kwargs.items() if k != "session"}
# Make options mutable so we can update conversation_id during function invocation loop
mutable_options: dict[str, Any] = dict(options) if options else {}
# Remove additional_function_arguments from options passed to underlying chat client
python/packages/core/agent_framework/_types.py
+1 -1
View File
@@ -1791,7 +1791,7 @@ class ContinuationToken(TypedDict):
# Restore and resume
token = json.loads(token_json)
response = await agent.run(
thread=thread,
session=session,
options={"continuation_token": token},
)
"""
python/packages/core/agent_framework/_workflows/_agent.py
+88 -125
View File
@@ -6,22 +6,22 @@ import json
import logging
import sys
import uuid
from collections.abc import AsyncIterable, Awaitable
from collections.abc import AsyncIterable, Awaitable, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload
from agent_framework import (
from .._agents import BaseAgent
from .._sessions import AgentSession, BaseContextProvider, BaseHistoryProvider, SessionContext
from .._types import (
AgentResponse,
AgentResponseUpdate,
AgentThread,
BaseAgent,
Content,
Message,
ResponseStream,
UsageDetails,
add_usage_details,
)
from .._types import add_usage_details
from ..exceptions import AgentExecutionException
from ._checkpoint import CheckpointStorage
from ._events import (
@@ -79,6 +79,7 @@ class WorkflowAgent(BaseAgent):
id: str | None = None,
name: str | None = None,
description: str | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
**kwargs: Any,
) -> None:
"""Initialize the WorkflowAgent.
@@ -90,6 +91,7 @@ class WorkflowAgent(BaseAgent):
id: Unique identifier for the agent. If None, will be generated.
name: Optional name for the agent.
description: Optional description of the agent.
context_providers: Optional sequence of context providers for the agent.
**kwargs: Additional keyword arguments passed to BaseAgent.
Note:
@@ -110,7 +112,7 @@ class WorkflowAgent(BaseAgent):
if not any(is_type_compatible(list[Message], input_type) for input_type in start_executor.input_types):
raise ValueError("Workflow's start executor cannot handle list[Message]")
super().__init__(id=id, name=name, description=description, **kwargs)
super().__init__(id=id, name=name, description=description, context_providers=context_providers, **kwargs)
self._workflow: Workflow = workflow
self._pending_requests: dict[str, WorkflowEvent[Any]] = {}
@@ -127,22 +129,22 @@ class WorkflowAgent(BaseAgent):
@overload
def run(
self,
messages: str | Message | list[str] | list[Message] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]: ...
) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ...
@overload
async def run(
self,
messages: str | Message | list[str] | list[Message] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
@@ -150,14 +152,14 @@ class WorkflowAgent(BaseAgent):
def run(
self,
messages: str | Message | list[str] | list[Message] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate] | Awaitable[AgentResponse]:
) -> ResponseStream[AgentResponseUpdate, AgentResponse] | Awaitable[AgentResponse]:
"""Get a response from the workflow agent.
Args:
@@ -167,7 +169,7 @@ class WorkflowAgent(BaseAgent):
Keyword Args:
stream: If True, returns an async iterable of updates. If False (default),
returns an awaitable AgentResponse.
thread: The conversation thread. If None, a new thread will be created.
session: The agent session for conversation context.
checkpoint_id: ID of checkpoint to restore from. If provided, the workflow
resumes from this checkpoint instead of starting fresh.
checkpoint_storage: Runtime checkpoint storage. When provided with checkpoint_id,
@@ -184,82 +186,21 @@ class WorkflowAgent(BaseAgent):
or AgentResponseUpdate objects. Request info events (type='request_info') will be
converted to function call and approval request contents.
"""
if messages is None:
messages = []
response_id = str(uuid.uuid4())
if stream:
return self._run_streaming(
messages=messages,
thread=thread,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
**kwargs,
return ResponseStream(
self._run_stream_impl(messages, response_id, session, checkpoint_id, checkpoint_storage, **kwargs),
finalizer=AgentResponse.from_updates,
)
return self._run_non_streaming(
messages=messages,
thread=thread,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
**kwargs,
)
async def _run_non_streaming(
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AgentResponse:
"""Internal non-streaming implementation."""
input_messages = normalize_messages_input(messages)
thread = thread or self.get_new_thread()
response_id = str(uuid.uuid4())
response = await self._run_impl(
input_messages, response_id, thread, checkpoint_id, checkpoint_storage, **kwargs
)
# Notify thread of new messages (both input and response messages)
await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
return response
async def _run_streaming(
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
"""Internal streaming implementation.
Yields AgentResponseUpdate objects. Output events (type='output') from the workflow
are converted to updates. Request info events (type='request_info') are converted
to function call and approval request contents.
"""
input_messages = normalize_messages_input(messages)
thread = thread or self.get_new_thread()
response_updates: list[AgentResponseUpdate] = []
response_id = str(uuid.uuid4())
async for update in self._run_stream_impl(
input_messages, response_id, thread, checkpoint_id, checkpoint_storage, **kwargs
):
response_updates.append(update)
yield update
# Convert updates to final response.
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, response.messages)
return self._run_impl(messages, response_id, session, checkpoint_id, checkpoint_storage, **kwargs)
async def _run_impl(
self,
input_messages: list[Message],
messages: str | Message | Sequence[str | Message],
response_id: str,
thread: AgentThread,
session: AgentSession | None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
@@ -267,9 +208,9 @@ class WorkflowAgent(BaseAgent):
"""Internal implementation of non-streaming execution.
Args:
input_messages: Normalized input messages to process.
messages: Normalized input messages to process.
response_id: The unique response ID for this workflow execution.
thread: The conversation thread containing message history.
session: The agent session for conversation context.
checkpoint_id: ID of checkpoint to restore from.
checkpoint_storage: Runtime checkpoint storage.
**kwargs: Additional keyword arguments passed through to the underlying
@@ -278,20 +219,44 @@ class WorkflowAgent(BaseAgent):
Returns:
An AgentResponse representing the workflow execution results.
"""
input_messages = normalize_messages_input(messages)
# run the context providers with the session
session_context = SessionContext(
session_id=session.session_id if session else None,
service_session_id=session.service_session_id if session else None,
input_messages=input_messages or [],
options={},
)
state = session.state if session else {}
for provider in self.context_providers:
if isinstance(provider, BaseHistoryProvider) and not provider.load_messages:
continue
await provider.before_run(
agent=self, # type: ignore[arg-type]
session=session, # type: ignore[arg-type]
context=session_context,
state=state,
)
# combine the messages
session_messages: list[Message] = session_context.get_messages(include_input=True)
output_events: list[WorkflowEvent[Any]] = []
async for event in self._run_core(
input_messages, thread, checkpoint_id, checkpoint_storage, streaming=False, **kwargs
session_messages, checkpoint_id, checkpoint_storage, streaming=False, **kwargs
):
if event.type == "output" or event.type == "request_info":
output_events.append(event)
return self._convert_workflow_events_to_agent_response(response_id, output_events)
result = self._convert_workflow_events_to_agent_response(response_id, output_events)
await self._run_after_providers(session=session, context=session_context)
return result
async def _run_stream_impl(
self,
input_messages: list[Message],
messages: str | Message | Sequence[str | Message],
response_id: str,
thread: AgentThread,
session: AgentSession | None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
@@ -299,9 +264,9 @@ class WorkflowAgent(BaseAgent):
"""Internal implementation of streaming execution.
Args:
input_messages: Normalized input messages to process.
messages: Input messages to process.
response_id: The unique response ID for this workflow execution.
thread: The conversation thread containing message history.
session: The agent session for conversation context.
checkpoint_id: ID of checkpoint to restore from.
checkpoint_storage: Runtime checkpoint storage.
**kwargs: Additional keyword arguments passed through to the underlying
@@ -310,17 +275,39 @@ class WorkflowAgent(BaseAgent):
Yields:
AgentResponseUpdate objects representing the workflow execution progress.
"""
input_messages = normalize_messages_input(messages)
# run the context providers with the session
session_context = SessionContext(
session_id=session.session_id if session else None,
service_session_id=session.service_session_id if session else None,
input_messages=input_messages or [],
options={},
)
state = session.state if session else {}
for provider in self.context_providers:
if isinstance(provider, BaseHistoryProvider) and not provider.load_messages:
continue
await provider.before_run(
agent=self, # type: ignore[arg-type]
session=session, # type: ignore[arg-type]
context=session_context,
state=state,
)
# combine the messages
session_messages: list[Message] = session_context.get_messages(include_input=True)
async for event in self._run_core(
input_messages, thread, checkpoint_id, checkpoint_storage, streaming=True, **kwargs
session_messages, checkpoint_id, checkpoint_storage, streaming=True, **kwargs
):
updates = self._convert_workflow_event_to_agent_response_updates(response_id, event)
for update in updates:
yield update
await self._run_after_providers(session=session, context=session_context)
async def _run_core(
self,
input_messages: list[Message],
thread: AgentThread,
input_messages: Sequence[Message],
checkpoint_id: str | None,
checkpoint_storage: CheckpointStorage | None,
streaming: bool,
@@ -330,7 +317,6 @@ class WorkflowAgent(BaseAgent):
Args:
input_messages: Normalized input messages to process.
thread: The conversation thread containing message history.
checkpoint_id: ID of checkpoint to restore from.
checkpoint_storage: Runtime checkpoint storage.
streaming: Whether to use streaming workflow methods.
@@ -371,10 +357,9 @@ class WorkflowAgent(BaseAgent):
yield event
else:
conversation_messages = await self._build_conversation_messages(thread, input_messages)
if streaming:
async for event in self.workflow.run(
message=conversation_messages,
message=input_messages,
stream=True,
checkpoint_storage=checkpoint_storage,
**kwargs,
@@ -382,7 +367,7 @@ class WorkflowAgent(BaseAgent):
yield event
else:
for event in await self.workflow.run(
message=conversation_messages,
message=input_messages,
checkpoint_storage=checkpoint_storage,
**kwargs,
):
@@ -390,29 +375,7 @@ class WorkflowAgent(BaseAgent):
# endregion Run Methods
async def _build_conversation_messages(
self,
thread: AgentThread,
input_messages: list[Message],
) -> list[Message]:
"""Build the complete conversation by prepending thread history to input messages.
Args:
thread: The conversation thread containing message history.
input_messages: The new input messages to append.
Returns:
A list of Message objects representing the full conversation.
"""
conversation_messages: list[Message] = []
if thread.message_store:
history = await thread.message_store.list_messages()
if history:
conversation_messages.extend(history)
conversation_messages.extend(input_messages)
return conversation_messages
def _process_pending_requests(self, input_messages: list[Message]) -> dict[str, Any]:
def _process_pending_requests(self, input_messages: Sequence[Message]) -> dict[str, Any]:
"""Process pending requests by extracting function responses and updating state.
Args:
@@ -669,7 +632,7 @@ class WorkflowAgent(BaseAgent):
# Ignore workflow-internal events
return []
def _extract_function_responses(self, input_messages: list[Message]) -> dict[str, Any]:
def _extract_function_responses(self, input_messages: Sequence[Message]) -> dict[str, Any]:
"""Extract function responses from input messages."""
function_responses: dict[str, Any] = {}
for message in input_messages:
python/packages/core/agent_framework/_workflows/_agent_executor.py
+41 -31
View File
@@ -11,7 +11,7 @@ from typing_extensions import Never
from agent_framework import Content
from .._agents import SupportsAgentRun
from .._threads import AgentThread
from .._sessions import AgentSession
from .._types import AgentResponse, AgentResponseUpdate, Message
from ._agent_utils import resolve_agent_id
from ._const import WORKFLOW_RUN_KWARGS_KEY
@@ -81,14 +81,14 @@ class AgentExecutor(Executor):
self,
agent: SupportsAgentRun,
*,
agent_thread: AgentThread | None = None,
session: AgentSession | None = None,
id: str | None = None,
):
"""Initialize the executor with a unique identifier.
Args:
agent: The agent to be wrapped by this executor.
agent_thread: The thread to use for running the agent. If None, a new thread will be created.
session: The session to use for running the agent. If None, a new session will be created.
id: A unique identifier for the executor. If None, the agent's name will be used if available.
"""
# Prefer provided id; else use agent.name if present; else generate deterministic prefix
@@ -97,7 +97,7 @@ class AgentExecutor(Executor):
raise ValueError("Agent must have a non-empty name or id or an explicit id must be provided.")
super().__init__(exec_id)
self._agent = agent
self._agent_thread = agent_thread or self._agent.get_new_thread()
self._session = session or self._agent.create_session()
self._pending_agent_requests: dict[str, Content] = {}
self._pending_responses_to_agent: list[Content] = []
@@ -205,35 +205,33 @@ class AgentExecutor(Executor):
async def on_checkpoint_save(self) -> dict[str, Any]:
"""Capture current executor state for checkpointing.
NOTE: if the thread storage is on the server side, the full thread state
may not be serialized locally. Therefore, we are relying on the server-side
to ensure the thread state is preserved and immutable across checkpoints.
This is not the case for AzureAI Agents, but works for the Responses API.
NOTE: if the session uses service-side storage, the full session state
may not be serialized locally.
Returns:
Dict containing serialized cache and thread state
Dict containing serialized cache and session state
"""
# Check if using AzureAIAgentClient with server-side thread and warn about checkpointing limitations
if is_chat_agent(self._agent) and self._agent_thread.service_thread_id is not None:
# Check if using AzureAIAgentClient with server-side session and warn about checkpointing limitations
if is_chat_agent(self._agent) and self._session.service_session_id is not None:
client_class_name = self._agent.client.__class__.__name__
client_module = self._agent.client.__class__.__module__
if client_class_name == "AzureAIAgentClient" and "azure_ai" in client_module:
logger.warning(
"Checkpointing an AgentExecutor with AzureAIAgentClient that uses server-side threads. "
"Currently, checkpointing does not capture messages from server-side threads "
"(service_thread_id: %s). The thread state in checkpoints is not immutable and can be "
"Checkpointing an AgentExecutor with AzureAIAgentClient that uses server-side sessions. "
"Currently, checkpointing does not capture messages from server-side sessions "
"(service_session_id: %s). The session state in checkpoints is not immutable and can be "
"modified by subsequent runs. If you need reliable checkpointing with Azure AI agents, "
"consider implementing a custom executor and managing the thread state yourself.",
self._agent_thread.service_thread_id,
"consider implementing a custom executor and managing the session state yourself.",
self._session.service_session_id,
)
serialized_thread = await self._agent_thread.serialize()
serialized_session = self._session.to_dict()
return {
"cache": self._cache,
"full_conversation": self._full_conversation,
"agent_thread": serialized_thread,
"agent_session": serialized_session,
"pending_agent_requests": self._pending_agent_requests,
"pending_responses_to_agent": self._pending_responses_to_agent,
}
@@ -246,22 +244,34 @@ class AgentExecutor(Executor):
state: Checkpoint data dict
"""
cache_payload = state.get("cache")
self._cache = cache_payload or []
if cache_payload:
try:
self._cache = cache_payload
except Exception as exc:
logger.warning("Failed to restore cache: %s", exc)
self._cache = []
else:
self._cache = []
full_conversation_payload = state.get("full_conversation")
self._full_conversation = full_conversation_payload or []
thread_payload = state.get("agent_thread")
if thread_payload:
if full_conversation_payload:
try:
# Deserialize the thread state directly
self._agent_thread = await AgentThread.deserialize(thread_payload)
self._full_conversation = full_conversation_payload
except Exception as exc:
logger.warning("Failed to restore agent thread: %s", exc)
self._agent_thread = self._agent.get_new_thread()
logger.warning("Failed to restore full conversation: %s", exc)
self._full_conversation = []
else:
self._agent_thread = self._agent.get_new_thread()
self._full_conversation = []
session_payload = state.get("agent_session")
if session_payload:
try:
self._session = AgentSession.from_dict(session_payload)
except Exception as exc:
logger.warning("Failed to restore agent session: %s", exc)
self._session = self._agent.create_session()
else:
self._session = self._agent.create_session()
pending_requests_payload = state.get("pending_agent_requests")
if pending_requests_payload:
@@ -321,7 +331,7 @@ class AgentExecutor(Executor):
response = await self._agent.run(
self._cache,
stream=False,
thread=self._agent_thread,
session=self._session,
options=options,
**run_kwargs,
)
@@ -352,7 +362,7 @@ class AgentExecutor(Executor):
async for update in self._agent.run(
self._cache,
stream=True,
thread=self._agent_thread,
session=self._session,
options=options,
**run_kwargs,
):
python/packages/core/agent_framework/ag_ui/__init__.py
-1
View File
@@ -8,7 +8,6 @@ PACKAGE_NAME = "agent-framework-ag-ui"
_IMPORTS = [
"__version__",
"AgentFrameworkAgent",
"AGUIThread",
"add_agent_framework_fastapi_endpoint",
"AGUIChatClient",
"AGUIEventConverter",
python/packages/core/agent_framework/azure/_responses_client.py
+4 -11
View File
@@ -83,8 +83,7 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
env_file_encoding: str | None = None,
instruction_role: str | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration
| None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
"""Initialize an Azure OpenAI Responses client.
@@ -190,9 +189,7 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
deployment_name = str(model_id)
# Project client path: create OpenAI client from an Azure AI Foundry project
if async_client is None and (
project_client is not None or project_endpoint is not None
):
if async_client is None and (project_client is not None or project_endpoint is not None):
async_client = self._create_client_from_project(
project_client=project_client,
project_endpoint=project_endpoint,
@@ -221,9 +218,7 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
and (hostname := urlparse(str(azure_openai_settings["endpoint"])).hostname)
and hostname.endswith(".openai.azure.com")
):
azure_openai_settings["base_url"] = urljoin(
str(azure_openai_settings["endpoint"]), "/openai/v1/"
)
azure_openai_settings["base_url"] = urljoin(str(azure_openai_settings["endpoint"]), "/openai/v1/")
if not azure_openai_settings["responses_deployment_name"]:
raise ServiceInitializationError(
@@ -236,9 +231,7 @@ class AzureOpenAIResponsesClient( # type: ignore[misc]
endpoint=azure_openai_settings["endpoint"],
base_url=azure_openai_settings["base_url"],
api_version=azure_openai_settings["api_version"], # type: ignore
api_key=azure_openai_settings["api_key"].get_secret_value()
if azure_openai_settings["api_key"]
else None,
api_key=azure_openai_settings["api_key"].get_secret_value() if azure_openai_settings["api_key"] else None,
ad_token=ad_token,
ad_token_provider=ad_token_provider,
token_endpoint=azure_openai_settings["token_endpoint"],
python/packages/core/agent_framework/exceptions.py
+2 -2
View File
@@ -49,8 +49,8 @@ class AgentInitializationError(AgentException):
pass
class AgentThreadException(AgentException):
"""An error occurred while managing the agent thread."""
class AgentSessionException(AgentException):
"""An error occurred while managing the agent session."""
pass
python/packages/core/agent_framework/mem0/__init__.py
+1 -1
View File
@@ -5,7 +5,7 @@ from typing import Any
IMPORT_PATH = "agent_framework_mem0"
PACKAGE_NAME = "agent-framework-mem0"
_IMPORTS = ["__version__", "Mem0Provider"]
_IMPORTS = ["__version__", "Mem0ContextProvider"]
def __getattr__(name: str) -> Any:
python/packages/core/agent_framework/mem0/__init__.pyi
+2 -2
View File
@@ -1,11 +1,11 @@
# Copyright (c) Microsoft. All rights reserved.
from agent_framework_mem0 import (
Mem0Provider,
Mem0ContextProvider,
__version__,
)
__all__ = [
"Mem0Provider",
"Mem0ContextProvider",
"__version__",
]
python/packages/core/agent_framework/observability.py
+8 -8
View File
@@ -39,7 +39,7 @@ if TYPE_CHECKING: # pragma: no cover
from ._agents import SupportsAgentRun
from ._clients import SupportsChatGetResponse
from ._threads import AgentThread
from ._sessions import AgentSession
from ._tools import FunctionTool
from ._types import (
AgentResponse,
@@ -1280,7 +1280,7 @@ class AgentTelemetryLayer:
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@@ -1290,7 +1290,7 @@ class AgentTelemetryLayer:
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
@@ -1299,7 +1299,7 @@ class AgentTelemetryLayer:
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
"""Trace agent runs with OpenTelemetry spans and metrics."""
@@ -1312,7 +1312,7 @@ class AgentTelemetryLayer:
return super_run( # type: ignore[no-any-return]
messages=messages,
stream=stream,
thread=thread,
session=session,
**kwargs,
)
@@ -1327,7 +1327,7 @@ class AgentTelemetryLayer:
agent_id=getattr(self, "id", "unknown"),
agent_name=getattr(self, "name", None) or getattr(self, "id", "unknown"),
agent_description=getattr(self, "description", None),
thread_id=thread.service_thread_id if thread else None,
thread_id=session.service_session_id if session else None,
all_options=merged_options,
**kwargs,
)
@@ -1336,7 +1336,7 @@ class AgentTelemetryLayer:
run_result = super_run(
messages=messages,
stream=True,
thread=thread,
session=session,
**kwargs,
)
if isinstance(run_result, ResponseStream):
@@ -1423,7 +1423,7 @@ class AgentTelemetryLayer:
response = await super_run(
messages=messages,
stream=False,
thread=thread,
session=session,
**kwargs,
)
except Exception as exception:
python/packages/core/agent_framework/openai/_assistant_provider.py
+13 -13
View File
@@ -13,8 +13,8 @@ from pydantic import BaseModel
from agent_framework._settings import SecretString, load_settings
from .._agents import Agent
from .._memory import ContextProvider
from .._middleware import MiddlewareTypes
from .._sessions import BaseContextProvider
from .._tools import FunctionTool
from .._types import normalize_tools
from ..exceptions import ServiceInitializationError
@@ -208,7 +208,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
metadata: dict[str, str] | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Create a new assistant on OpenAI and return a Agent.
@@ -230,7 +230,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
These options are applied to every run unless overridden.
Include ``response_format`` here for structured output responses.
middleware: MiddlewareTypes for the Agent.
context_provider: Context provider for the Agent.
context_providers: Context providers for the Agent.
Returns:
A Agent instance wrapping the created assistant.
@@ -304,7 +304,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
tools=normalized_tools,
instructions=instructions,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
default_options=default_options,
)
@@ -316,7 +316,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
instructions: str | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Retrieve an existing assistant by ID and return a Agent.
@@ -335,7 +335,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: MiddlewareTypes for the Agent.
context_provider: Context provider for the Agent.
context_providers: Context providers for the Agent.
Returns:
A Agent instance wrapping the retrieved assistant.
@@ -371,7 +371,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
instructions=instructions,
default_options=default_options,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
def as_agent(
@@ -382,7 +382,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
instructions: str | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
) -> Agent[OptionsCoT]:
"""Wrap an existing SDK Assistant object as a Agent.
@@ -400,7 +400,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: MiddlewareTypes for the Agent.
context_provider: Context provider for the Agent.
context_providers: Context providers for the Agent.
Returns:
A Agent instance wrapping the assistant.
@@ -437,7 +437,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
instructions=instructions,
default_options=default_options,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
)
def _validate_function_tools(
@@ -524,7 +524,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
tools: list[FunctionTool | MutableMapping[str, Any]] | None,
instructions: str | None,
middleware: Sequence[MiddlewareTypes] | None,
context_provider: ContextProvider | None,
context_providers: Sequence[BaseContextProvider] | None,
default_options: OptionsCoT | None = None,
**kwargs: Any,
) -> Agent[OptionsCoT]:
@@ -535,7 +535,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
tools: Tools for the agent.
instructions: Instructions override.
middleware: MiddlewareTypes for the agent.
context_provider: Context provider for the agent.
context_providers: Context providers for the agent.
default_options: Default chat options for the agent (may include response_format).
**kwargs: Additional arguments passed to Agent.
@@ -563,7 +563,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
instructions=final_instructions,
tools=tools if tools else None,
middleware=middleware,
context_provider=context_provider,
context_providers=context_providers,
default_options=default_options, # type: ignore[arg-type]
**kwargs,
)
python/packages/core/agent_framework/openai/_responses_client.py
+3 -1
View File
@@ -13,7 +13,7 @@ from collections.abc import (
)
from datetime import datetime, timezone
from itertools import chain
from typing import TYPE_CHECKING, Any, Generic, Literal, NoReturn, TypedDict, cast
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, NoReturn, TypedDict, cast
from openai import AsyncOpenAI, BadRequestError
from openai.types.responses.file_search_tool_param import FileSearchToolParam
@@ -238,6 +238,8 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
Use ``OpenAIResponsesClient`` instead for a fully-featured client with all layers applied.
"""
STORES_BY_DEFAULT: ClassVar[bool] = True # type: ignore[reportIncompatibleVariableOverride, misc]
FILE_SEARCH_MAX_RESULTS: int = 50
# region Inner Methods
python/packages/core/agent_framework/redis/__init__.py
+1 -1
View File
@@ -5,7 +5,7 @@ from typing import Any
IMPORT_PATH = "agent_framework_redis"
PACKAGE_NAME = "agent-framework-redis"
_IMPORTS = ["__version__", "RedisProvider", "RedisChatMessageStore"]
_IMPORTS = ["__version__", "RedisContextProvider", "RedisHistoryProvider"]
def __getattr__(name: str) -> Any:
python/packages/core/agent_framework/redis/__init__.pyi
+4 -4
View File
@@ -1,13 +1,13 @@
# Copyright (c) Microsoft. All rights reserved.
from agent_framework_redis import (
RedisChatMessageStore,
RedisProvider,
RedisContextProvider,
RedisHistoryProvider,
__version__,
)
__all__ = [
"RedisChatMessageStore",
"RedisProvider",
"RedisContextProvider",
"RedisHistoryProvider",
"__version__",
]
python/packages/core/tests/azure/test_azure_assistants_client.py
+24 -24
View File
@@ -12,7 +12,7 @@ from agent_framework import (
Agent,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
ChatResponse,
ChatResponseUpdate,
Message,
@@ -433,70 +433,70 @@ async def test_azure_assistants_agent_basic_run_streaming():
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_thread_persistence():
"""Test Agent thread persistence across runs with AzureOpenAIAssistantsClient."""
async def test_azure_assistants_agent_session_persistence():
"""Test Agent session persistence across runs with AzureOpenAIAssistantsClient."""
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# Create a new session that will be reused
session = agent.create_session()
# First message - establish context
first_response = await agent.run(
"Remember this number: 42. What number did I just tell you to remember?", thread=thread
"Remember this number: 42. What number did I just tell you to remember?", session=session
)
assert isinstance(first_response, AgentResponse)
assert "42" in first_response.text
# Second message - test conversation memory
second_response = await agent.run(
"What number did I tell you to remember in my previous message?", thread=thread
"What number did I tell you to remember in my previous message?", session=session
)
assert isinstance(second_response, AgentResponse)
assert "42" in second_response.text
# Verify thread has been populated with conversation ID
assert thread.service_thread_id is not None
# Verify session has been populated with conversation ID
assert session.service_session_id is not None
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_existing_thread_id():
"""Test Agent with existing thread ID to continue conversations across agent instances."""
# First, create a conversation and capture the thread ID
existing_thread_id = None
async def test_azure_assistants_agent_existing_session_id():
"""Test Agent with existing session ID to continue conversations across agent instances."""
# First, create a conversation and capture the session ID
existing_session_id = None
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Start a conversation and get the thread ID
thread = agent.get_new_thread()
response1 = await agent.run("What's the weather in Paris?", thread=thread)
# Start a conversation and get the session ID
session = agent.create_session()
response1 = await agent.run("What's the weather in Paris?", session=session)
# Validate first response
assert isinstance(response1, AgentResponse)
assert response1.text is not None
assert any(word in response1.text.lower() for word in ["weather", "paris"])
# The thread ID is set after the first response
existing_thread_id = thread.service_thread_id
assert existing_thread_id is not None
# The session ID is set after the first response
existing_session_id = session.service_session_id
assert existing_session_id is not None
# Now continue with the same thread ID in a new agent instance
# Now continue with the same session ID in a new agent instance
async with Agent(
client=AzureOpenAIAssistantsClient(thread_id=existing_thread_id, credential=AzureCliCredential()),
client=AzureOpenAIAssistantsClient(thread_id=existing_session_id, credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Create a thread with the existing ID
thread = AgentThread(service_thread_id=existing_thread_id)
# Create a session with the existing ID
session = AgentSession(service_session_id=existing_session_id)
# Ask about the previous conversation
response2 = await agent.run("What was the last city I asked about?", thread=thread)
response2 = await agent.run("What was the last city I asked about?", session=session)
# Validate that the agent remembers the previous conversation
assert isinstance(response2, AgentResponse)
python/packages/core/tests/azure/test_azure_chat_client.py
+19 -19
View File
@@ -800,23 +800,23 @@ async def test_azure_openai_chat_client_agent_basic_run_streaming():
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_thread_persistence():
"""Test Azure OpenAI chat client agent thread persistence across runs with AzureOpenAIChatClient."""
async def test_azure_openai_chat_client_agent_session_persistence():
"""Test Azure OpenAI chat client agent session persistence across runs with AzureOpenAIChatClient."""
async with Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# Create a new session that will be reused
session = agent.create_session()
# First interaction
response1 = await agent.run("My name is Alice. Remember this.", thread=thread)
response1 = await agent.run("My name is Alice. Remember this.", session=session)
assert isinstance(response1, AgentResponse)
assert response1.text is not None
# Second interaction - test memory
response2 = await agent.run("What is my name?", thread=thread)
response2 = await agent.run("What is my name?", session=session)
assert isinstance(response2, AgentResponse)
assert response2.text is not None
@@ -825,33 +825,33 @@ async def test_azure_openai_chat_client_agent_thread_persistence():
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_existing_thread():
"""Test Azure OpenAI chat client agent with existing thread to continue conversations across agent instances."""
# First conversation - capture the thread
preserved_thread = None
async def test_azure_openai_chat_client_agent_existing_session():
"""Test Azure OpenAI chat client agent with existing session to continue conversations across agent instances."""
# First conversation - capture the session
preserved_session = None
async with Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the thread
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My name is Alice. Remember this.", thread=thread)
# Start a conversation and capture the session
session = first_agent.create_session()
first_response = await first_agent.run("My name is Alice. Remember this.", session=session)
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Preserve the thread for reuse
preserved_thread = thread
# Preserve the session for reuse
preserved_session = session
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
# Second conversation - reuse the session in a new agent instance
if preserved_session:
async with Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved thread
second_response = await second_agent.run("What is my name?", thread=preserved_thread)
# Reuse the preserved session
second_response = await second_agent.run("What is my name?", session=preserved_session)
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
python/packages/core/tests/azure/test_azure_responses_client.py
+13 -13
View File
@@ -537,33 +537,33 @@ async def test_integration_client_agent_hosted_code_interpreter_tool():
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_integration_client_agent_existing_thread():
"""Test Azure Responses Client agent with existing thread to continue conversations across agent instances."""
# First conversation - capture the thread
preserved_thread = None
async def test_integration_client_agent_existing_session():
"""Test Azure Responses Client agent with existing session to continue conversations across agent instances."""
# First conversation - capture the session
preserved_session = None
async with Agent(
client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the thread
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My hobby is photography. Remember this.", thread=thread, store=True)
# Start a conversation and capture the session
session = first_agent.create_session()
first_response = await first_agent.run("My hobby is photography. Remember this.", session=session, store=True)
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Preserve the thread for reuse
preserved_thread = thread
# Preserve the session for reuse
preserved_session = session
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
# Second conversation - reuse the session in a new agent instance
if preserved_session:
async with Agent(
client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved thread
second_response = await second_agent.run("What is my hobby?", thread=preserved_thread)
# Reuse the preserved session
second_response = await second_agent.run("What is my hobby?", session=preserved_session)
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
python/packages/core/tests/core/conftest.py
+13 -13
View File
@@ -13,7 +13,7 @@ from pytest import fixture
from agent_framework import (
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseChatClient,
ChatMiddlewareLayer,
ChatResponse,
@@ -261,7 +261,7 @@ def chat_client_base(enable_function_calling: bool, max_iterations: int) -> Mock
# region Agents
class MockAgentThread(AgentThread):
class MockAgentSession(AgentSession):
pass
@@ -284,41 +284,41 @@ class MockAgent(SupportsAgentRun):
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
stream: bool = False,
**kwargs: Any,
) -> Awaitable[AgentResponse] | AsyncIterable[AgentResponseUpdate]:
if stream:
return self._run_stream_impl(messages=messages, thread=thread, **kwargs)
return self._run_impl(messages=messages, thread=thread, **kwargs)
return self._run_stream_impl(messages=messages, session=session, **kwargs)
return self._run_impl(messages=messages, session=session, **kwargs)
async def _run_impl(
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> AgentResponse:
logger.debug(f"Running mock agent, with: {messages=}, {thread=}, {kwargs=}")
logger.debug(f"Running mock agent, with: {messages=}, {session=}, {kwargs=}")
return AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("Response")])])
async def _run_stream_impl(
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
logger.debug(f"Running mock agent stream, with: {messages=}, {thread=}, {kwargs=}")
logger.debug(f"Running mock agent stream, with: {messages=}, {session=}, {kwargs=}")
yield AgentResponseUpdate(contents=[Content.from_text("Response")])
def get_new_thread(self) -> AgentThread:
return MockAgentThread()
def create_session(self) -> AgentSession:
return MockAgentSession()
@fixture
def agent_thread() -> AgentThread:
return MockAgentThread()
def agent_session() -> AgentSession:
return MockAgentSession()
@fixture
python/packages/core/tests/core/test_agents.py
+193 -186
View File
@@ -1,7 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
import contextlib
from collections.abc import AsyncIterable, MutableSequence, Sequence
from collections.abc import AsyncIterable, MutableSequence
from typing import Any
from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4
@@ -13,13 +13,11 @@ from agent_framework import (
Agent,
AgentResponse,
AgentResponseUpdate,
AgentThread,
ChatMessageStore,
AgentSession,
BaseContextProvider,
ChatOptions,
ChatResponse,
Content,
Context,
ContextProvider,
FunctionTool,
Message,
SupportsAgentRun,
@@ -28,11 +26,10 @@ from agent_framework import (
)
from agent_framework._agents import _merge_options, _sanitize_agent_name
from agent_framework._mcp import MCPTool
from agent_framework.exceptions import AgentExecutionException, AgentInitializationError
def test_agent_thread_type(agent_thread: AgentThread) -> None:
assert isinstance(agent_thread, AgentThread)
def test_agent_session_type(agent_session: AgentSession) -> None:
assert isinstance(agent_session, AgentSession)
def test_agent_type(agent: SupportsAgentRun) -> None:
@@ -93,38 +90,42 @@ async def test_chat_client_agent_run_streaming(client: SupportsChatGetResponse)
assert result.text == "test streaming response another update"
async def test_chat_client_agent_get_new_thread(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_create_session(client: SupportsChatGetResponse) -> None:
agent = Agent(client=client)
thread = agent.get_new_thread()
session = agent.create_session()
assert isinstance(thread, AgentThread)
assert isinstance(session, AgentSession)
async def test_chat_client_agent_prepare_thread_and_messages(client: SupportsChatGetResponse) -> None:
agent = Agent(client=client)
async def test_chat_client_agent_prepare_session_and_messages(client: SupportsChatGetResponse) -> None:
from agent_framework._sessions import InMemoryHistoryProvider
agent = Agent(client=client, context_providers=[InMemoryHistoryProvider("memory")])
message = Message(role="user", text="Hello")
thread = AgentThread(message_store=ChatMessageStore(messages=[message]))
session = AgentSession()
session.state["memory"] = {"messages": [message]}
_, _, result_messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=thread,
session_context, _ = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
session=session,
input_messages=[Message(role="user", text="Test")],
)
result_messages = session_context.get_messages(include_input=True)
assert len(result_messages) == 2
assert result_messages[0] == message
assert result_messages[0].text == "Hello"
assert result_messages[1].text == "Test"
async def test_prepare_thread_does_not_mutate_agent_chat_options(client: SupportsChatGetResponse) -> None:
async def test_prepare_session_does_not_mutate_agent_chat_options(client: SupportsChatGetResponse) -> None:
tool = {"type": "code_interpreter"}
agent = Agent(client=client, tools=[tool])
assert agent.default_options.get("tools") is not None
base_tools = agent.default_options["tools"]
thread = agent.get_new_thread()
session = agent.create_session()
_, prepared_chat_options, _ = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=thread,
_, prepared_chat_options = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
session=session,
input_messages=[Message(role="user", text="Test")],
)
@@ -135,7 +136,7 @@ async def test_prepare_thread_does_not_mutate_agent_chat_options(client: Support
assert len(agent.default_options["tools"]) == 1
async def test_chat_client_agent_update_thread_id(chat_client_base: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_run_with_session(chat_client_base: SupportsChatGetResponse) -> None:
mock_response = ChatResponse(
messages=[Message(role="assistant", contents=[Content.from_text("test response")])],
conversation_id="123",
@@ -145,25 +146,24 @@ async def test_chat_client_agent_update_thread_id(chat_client_base: SupportsChat
client=chat_client_base,
tools={"type": "code_interpreter"},
)
thread = agent.get_new_thread()
session = agent.get_session(service_session_id="123")
result = await agent.run("Hello", thread=thread)
result = await agent.run("Hello", session=session)
assert result.text == "test response"
assert thread.service_thread_id == "123"
assert session.service_session_id == "123"
async def test_chat_client_agent_update_thread_messages(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_update_session_messages(client: SupportsChatGetResponse) -> None:
agent = Agent(client=client)
thread = agent.get_new_thread()
session = agent.create_session()
result = await agent.run("Hello", thread=thread)
result = await agent.run("Hello", session=session)
assert result.text == "test response"
assert thread.service_thread_id is None
assert thread.message_store is not None
assert session.service_session_id is None
chat_messages: list[Message] = await thread.message_store.list_messages()
chat_messages: list[Message] = session.state.get("memory", {}).get("messages", [])
assert chat_messages is not None
assert len(chat_messages) == 2
@@ -171,12 +171,12 @@ async def test_chat_client_agent_update_thread_messages(client: SupportsChatGetR
assert chat_messages[1].text == "test response"
async def test_chat_client_agent_update_thread_conversation_id_missing(client: SupportsChatGetResponse) -> None:
async def test_chat_client_agent_update_session_conversation_id_missing(client: SupportsChatGetResponse) -> None:
agent = Agent(client=client)
thread = AgentThread(service_thread_id="123")
session = agent.get_session(service_session_id="123")
with raises(AgentExecutionException, match="Service did not return a valid conversation id"):
await agent._update_thread_with_type_and_conversation_id(thread, None) # type: ignore[reportPrivateUsage]
# With the session-based API, service_session_id is managed directly on the session
assert session.service_session_id == "123"
async def test_chat_client_agent_default_author_name(client: SupportsChatGetResponse) -> None:
@@ -214,54 +214,41 @@ async def test_chat_client_agent_author_name_is_used_from_response(chat_client_b
# Mock context provider for testing
class MockContextProvider(ContextProvider):
class MockContextProvider(BaseContextProvider):
def __init__(self, messages: list[Message] | None = None) -> None:
super().__init__(source_id="mock")
self.context_messages = messages
self.thread_created_called = False
self.invoked_called = False
self.invoking_called = False
self.thread_created_thread_id = None
self.invoked_thread_id = None
self.before_run_called = False
self.after_run_called = False
self.new_messages: list[Message] = []
self.last_service_session_id: str | None = None
async def thread_created(self, thread_id: str | None) -> None:
self.thread_created_called = True
self.thread_created_thread_id = thread_id
async def before_run(self, *, agent: Any, session: Any, context: Any, state: Any) -> None:
self.before_run_called = True
if self.context_messages:
context.extend_messages(self, self.context_messages)
async def invoked(
self,
request_messages: Message | Sequence[Message],
response_messages: Message | Sequence[Message] | None = None,
invoke_exception: Any = None,
**kwargs: Any,
) -> None:
self.invoked_called = True
if isinstance(request_messages, Message):
self.new_messages.append(request_messages)
else:
self.new_messages.extend(request_messages)
if isinstance(response_messages, Message):
self.new_messages.append(response_messages)
else:
self.new_messages.extend(response_messages)
async def invoking(self, messages: Message | MutableSequence[Message], **kwargs: Any) -> Context:
self.invoking_called = True
return Context(messages=self.context_messages)
async def after_run(self, *, agent: Any, session: Any, context: Any, state: Any) -> None:
self.after_run_called = True
if session:
self.last_service_session_id = session.service_session_id
if context.response:
self.new_messages.extend(context.input_messages)
self.new_messages.extend(context.response.messages)
async def test_chat_agent_context_providers_model_invoking(client: SupportsChatGetResponse) -> None:
"""Test that context providers' invoking is called during agent run."""
async def test_chat_agent_context_providers_model_before_run(client: SupportsChatGetResponse) -> None:
"""Test that context providers' before_run is called during agent run."""
mock_provider = MockContextProvider(messages=[Message(role="system", text="Test context instructions")])
agent = Agent(client=client, context_provider=mock_provider)
agent = Agent(client=client, context_providers=[mock_provider])
await agent.run("Hello")
assert mock_provider.invoking_called
assert mock_provider.before_run_called
async def test_chat_agent_context_providers_thread_created(chat_client_base: SupportsChatGetResponse) -> None:
"""Test that context providers' thread_created is called during agent run."""
async def test_chat_agent_context_providers_after_run(chat_client_base: SupportsChatGetResponse) -> None:
"""Test that context providers' after_run is called during agent run."""
mock_provider = MockContextProvider()
chat_client_base.run_responses = [
ChatResponse(
@@ -270,22 +257,23 @@ async def test_chat_agent_context_providers_thread_created(chat_client_base: Sup
)
]
agent = Agent(client=chat_client_base, context_provider=mock_provider)
agent = Agent(client=chat_client_base, context_providers=[mock_provider])
await agent.run("Hello")
session = agent.get_session(service_session_id="test-thread-id")
await agent.run("Hello", session=session)
assert mock_provider.thread_created_called
assert mock_provider.thread_created_thread_id == "test-thread-id"
assert mock_provider.after_run_called
assert mock_provider.last_service_session_id == "test-thread-id"
async def test_chat_agent_context_providers_messages_adding(client: SupportsChatGetResponse) -> None:
"""Test that context providers' invoked is called during agent run."""
"""Test that context providers' after_run is called during agent run."""
mock_provider = MockContextProvider()
agent = Agent(client=client, context_provider=mock_provider)
agent = Agent(client=client, context_providers=[mock_provider])
await agent.run("Hello")
assert mock_provider.invoked_called
assert mock_provider.after_run_called
# Should be called with both input and response messages
assert len(mock_provider.new_messages) >= 2
@@ -293,12 +281,13 @@ async def test_chat_agent_context_providers_messages_adding(client: SupportsChat
async def test_chat_agent_context_instructions_in_messages(client: SupportsChatGetResponse) -> None:
"""Test that AI context instructions are included in messages."""
mock_provider = MockContextProvider(messages=[Message(role="system", text="Context-specific instructions")])
agent = Agent(client=client, instructions="Agent instructions", context_provider=mock_provider)
agent = Agent(client=client, instructions="Agent instructions", context_providers=[mock_provider])
# We need to test the _prepare_thread_and_messages method directly
_, _, messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=None, input_messages=[Message(role="user", text="Hello")]
# We need to test the _prepare_session_and_messages method directly
session_context, _ = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
session=None, input_messages=[Message(role="user", text="Hello")]
)
messages = session_context.get_messages(include_input=True)
# Should have context instructions, and user message
assert len(messages) == 2
@@ -312,11 +301,12 @@ async def test_chat_agent_context_instructions_in_messages(client: SupportsChatG
async def test_chat_agent_no_context_instructions(client: SupportsChatGetResponse) -> None:
"""Test behavior when AI context has no instructions."""
mock_provider = MockContextProvider()
agent = Agent(client=client, instructions="Agent instructions", context_provider=mock_provider)
agent = Agent(client=client, instructions="Agent instructions", context_providers=[mock_provider])
_, _, messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=None, input_messages=[Message(role="user", text="Hello")]
session_context, _ = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
session=None, input_messages=[Message(role="user", text="Hello")]
)
messages = session_context.get_messages(include_input=True)
# Should have agent instructions and user message only
assert len(messages) == 1
@@ -327,7 +317,7 @@ async def test_chat_agent_no_context_instructions(client: SupportsChatGetRespons
async def test_chat_agent_run_stream_context_providers(client: SupportsChatGetResponse) -> None:
"""Test that context providers work with run method."""
mock_provider = MockContextProvider(messages=[Message(role="system", text="Stream context instructions")])
agent = Agent(client=client, context_provider=mock_provider)
agent = Agent(client=client, context_providers=[mock_provider])
# Collect all stream updates and get final response
stream = agent.run("Hello", stream=True)
@@ -338,14 +328,12 @@ async def test_chat_agent_run_stream_context_providers(client: SupportsChatGetRe
await stream.get_final_response()
# Verify context provider was called
assert mock_provider.invoking_called
# no conversation id is created, so no need to thread_create to be called.
assert not mock_provider.thread_created_called
assert mock_provider.invoked_called
assert mock_provider.before_run_called
assert mock_provider.after_run_called
async def test_chat_agent_context_providers_with_thread_service_id(chat_client_base: SupportsChatGetResponse) -> None:
"""Test context providers with service-managed thread."""
async def test_chat_agent_context_providers_with_service_session_id(chat_client_base: SupportsChatGetResponse) -> None:
"""Test context providers with service-managed session."""
mock_provider = MockContextProvider()
chat_client_base.run_responses = [
ChatResponse(
@@ -354,14 +342,14 @@ async def test_chat_agent_context_providers_with_thread_service_id(chat_client_b
)
]
agent = Agent(client=chat_client_base, context_provider=mock_provider)
agent = Agent(client=chat_client_base, context_providers=[mock_provider])
# Use existing service-managed thread
thread = agent.get_new_thread(service_thread_id="existing-thread-id")
await agent.run("Hello", thread=thread)
# Use existing service-managed session
session = agent.get_session(service_session_id="existing-thread-id")
await agent.run("Hello", session=session)
# invoked should be called with the service thread ID from response
assert mock_provider.invoked_called
# after_run should be called
assert mock_provider.after_run_called
# Tests for as_tool method
@@ -562,16 +550,16 @@ async def test_chat_agent_with_local_mcp_tools(client: SupportsChatGetResponse)
pass
async def test_agent_tool_receives_thread_in_kwargs(chat_client_base: Any) -> None:
"""Verify tool execution receives 'thread' inside **kwargs when function is called by client."""
async def test_agent_tool_receives_session_in_kwargs(chat_client_base: Any) -> None:
"""Verify tool execution receives 'session' inside **kwargs when function is called by client."""
captured: dict[str, Any] = {}
@tool(name="echo_thread_info", approval_mode="never_require")
def echo_thread_info(text: str, **kwargs: Any) -> str: # type: ignore[reportUnknownParameterType]
thread = kwargs.get("thread")
captured["has_thread"] = thread is not None
captured["has_message_store"] = thread.message_store is not None if isinstance(thread, AgentThread) else False
@tool(name="echo_session_info", approval_mode="never_require")
def echo_session_info(text: str, **kwargs: Any) -> str: # type: ignore[reportUnknownParameterType]
session = kwargs.get("session")
captured["has_session"] = session is not None
captured["has_state"] = session.state is not None if isinstance(session, AgentSession) else False
return f"echo: {text}"
# Make the base client emit a function call for our tool
@@ -580,21 +568,21 @@ async def test_agent_tool_receives_thread_in_kwargs(chat_client_base: Any) -> No
messages=Message(
role="assistant",
contents=[
Content.from_function_call(call_id="1", name="echo_thread_info", arguments='{"text": "hello"}')
Content.from_function_call(call_id="1", name="echo_session_info", arguments='{"text": "hello"}')
],
)
),
ChatResponse(messages=Message(role="assistant", text="done")),
]
agent = Agent(client=chat_client_base, tools=[echo_thread_info], chat_message_store_factory=ChatMessageStore)
thread = agent.get_new_thread()
agent = Agent(client=chat_client_base, tools=[echo_session_info])
session = agent.create_session()
result = await agent.run("hello", thread=thread, options={"additional_function_arguments": {"thread": thread}})
result = await agent.run("hello", session=session, options={"additional_function_arguments": {"session": session}})
assert result.text == "done"
assert captured.get("has_thread") is True
assert captured.get("has_message_store") is True
assert captured.get("has_session") is True
assert captured.get("has_state") is True
async def test_chat_agent_tool_choice_run_level_overrides_agent_level(chat_client_base: Any, tool_tool: Any) -> None:
@@ -801,73 +789,67 @@ def test_sanitize_agent_name_replaces_invalid_chars():
# endregion
# region Test SupportsAgentRun.get_new_thread and deserialize_thread
# region Test SupportsAgentRun.create_session
@pytest.mark.asyncio
async def test_agent_get_new_thread(chat_client_base: SupportsChatGetResponse, tool_tool: FunctionTool):
"""Test that get_new_thread returns a new AgentThread."""
async def test_agent_create_session(chat_client_base: SupportsChatGetResponse, tool_tool: FunctionTool):
"""Test that create_session returns a new AgentSession."""
agent = Agent(client=chat_client_base, tools=[tool_tool])
thread = agent.get_new_thread()
session = agent.create_session()
assert thread is not None
assert isinstance(thread, AgentThread)
assert session is not None
assert isinstance(session, AgentSession)
@pytest.mark.asyncio
async def test_agent_get_new_thread_with_context_provider(
async def test_agent_create_session_with_context_providers(
chat_client_base: SupportsChatGetResponse, tool_tool: FunctionTool
):
"""Test that get_new_thread passes context_provider to the thread."""
"""Test that create_session works when context_providers are set on the agent."""
class TestContextProvider(ContextProvider):
async def invoking(self, messages, **kwargs):
return Context()
class TestContextProvider(BaseContextProvider):
def __init__(self):
super().__init__(source_id="test")
provider = TestContextProvider()
agent = Agent(client=chat_client_base, tools=[tool_tool], context_provider=provider)
agent = Agent(client=chat_client_base, tools=[tool_tool], context_providers=[provider])
thread = agent.get_new_thread()
session = agent.create_session()
assert thread is not None
assert thread.context_provider is provider
assert session is not None
assert agent.context_providers[0] is provider
@pytest.mark.asyncio
async def test_agent_get_new_thread_with_service_thread_id(
async def test_agent_get_session_with_service_session_id(
chat_client_base: SupportsChatGetResponse, tool_tool: FunctionTool
):
"""Test that get_new_thread passes kwargs like service_thread_id to the thread."""
"""Test that get_session creates a session with service_session_id."""
agent = Agent(client=chat_client_base, tools=[tool_tool])
thread = agent.get_new_thread(service_thread_id="test-thread-123")
session = agent.get_session(service_session_id="test-thread-123")
assert thread is not None
assert thread.service_thread_id == "test-thread-123"
assert session is not None
assert session.service_session_id == "test-thread-123"
@pytest.mark.asyncio
async def test_agent_deserialize_thread(chat_client_base: SupportsChatGetResponse, tool_tool: FunctionTool):
"""Test deserialize_thread restores a thread from serialized state."""
agent = Agent(client=chat_client_base, tools=[tool_tool])
# Create serialized thread state with messages
def test_agent_session_from_dict(chat_client_base: SupportsChatGetResponse, tool_tool: FunctionTool):
"""Test AgentSession.from_dict restores a session from serialized state."""
# Create serialized session state
serialized_state = {
"service_thread_id": None,
"chat_message_store_state": {
"messages": [{"role": "user", "text": "Hello"}],
},
"type": "session",
"session_id": "test-session",
"service_session_id": None,
"state": {},
}
thread = await agent.deserialize_thread(serialized_state)
session = AgentSession.from_dict(serialized_state)
assert thread is not None
assert isinstance(thread, AgentThread)
assert thread.message_store is not None
messages = await thread.message_store.list_messages()
assert len(messages) == 1
assert messages[0].text == "Hello"
assert session is not None
assert isinstance(session, AgentSession)
assert session.session_id == "test-session"
# endregion
@@ -876,20 +858,6 @@ async def test_agent_deserialize_thread(chat_client_base: SupportsChatGetRespons
# region Test Agent initialization edge cases
@pytest.mark.asyncio
async def test_chat_agent_raises_with_both_conversation_id_and_store():
"""Test Agent raises error with both conversation_id and chat_message_store_factory."""
mock_client = MagicMock()
mock_store_factory = MagicMock()
with pytest.raises(AgentInitializationError, match="Cannot specify both"):
Agent(
client=mock_client,
default_options={"conversation_id": "test_id"},
chat_message_store_factory=mock_store_factory,
)
def test_chat_agent_calls_update_agent_name_on_client():
"""Test that Agent calls _update_agent_name_and_description on client if available."""
mock_client = MagicMock()
@@ -914,19 +882,22 @@ async def test_chat_agent_context_provider_adds_tools_when_agent_has_none(chat_c
"""A tool provided by context."""
return text
class ToolContextProvider(ContextProvider):
async def invoking(self, messages, **kwargs):
return Context(tools=[context_tool])
class ToolContextProvider(BaseContextProvider):
def __init__(self):
super().__init__(source_id="tool-context")
async def before_run(self, *, agent, session, context, state):
context.extend_tools("tool-context", [context_tool])
provider = ToolContextProvider()
agent = Agent(client=chat_client_base, context_provider=provider)
agent = Agent(client=chat_client_base, context_providers=[provider])
# Agent starts with empty tools list
assert agent.default_options.get("tools") == []
# Run the agent and verify context tools are added
_, options, _ = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=None, input_messages=[Message(role="user", text="Hello")]
_, options = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
session=None, input_messages=[Message(role="user", text="Hello")]
)
# The context tools should now be in the options
@@ -940,40 +911,76 @@ async def test_chat_agent_context_provider_adds_instructions_when_agent_has_none
):
"""Test that context provider instructions are used when agent has no default instructions."""
class InstructionContextProvider(ContextProvider):
async def invoking(self, messages, **kwargs):
return Context(instructions="Context-provided instructions")
class InstructionContextProvider(BaseContextProvider):
def __init__(self):
super().__init__(source_id="instruction-context")
async def before_run(self, *, agent, session, context, state):
context.extend_instructions("instruction-context", "Context-provided instructions")
provider = InstructionContextProvider()
agent = Agent(client=chat_client_base, context_provider=provider)
agent = Agent(client=chat_client_base, context_providers=[provider])
# Verify agent has no default instructions
assert agent.default_options.get("instructions") is None
# Run the agent and verify context instructions are available
_, options, _ = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=None, input_messages=[Message(role="user", text="Hello")]
_, options = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
session=None, input_messages=[Message(role="user", text="Hello")]
)
# The context instructions should now be in the options
assert options.get("instructions") == "Context-provided instructions"
@pytest.mark.asyncio
async def test_chat_agent_raises_on_conversation_id_mismatch(chat_client_base: SupportsChatGetResponse):
"""Test that Agent raises when thread and agent have different conversation IDs."""
agent = Agent(
client=chat_client_base,
default_options={"conversation_id": "agent-conversation-id"},
)
# region STORES_BY_DEFAULT tests
# Create a thread with a different service_thread_id
thread = AgentThread(service_thread_id="different-thread-id")
with pytest.raises(AgentExecutionException, match="conversation_id set on the agent is different"):
await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=thread, input_messages=[Message(role="user", text="Hello")]
)
async def test_stores_by_default_skips_inmemory_injection(client: SupportsChatGetResponse) -> None:
"""Client with STORES_BY_DEFAULT=True should not auto-inject InMemoryHistoryProvider."""
from agent_framework._sessions import InMemoryHistoryProvider
# Simulate a client that stores by default
client.STORES_BY_DEFAULT = True # type: ignore[attr-defined]
agent = Agent(client=client)
session = agent.create_session()
await agent.run("Hello", session=session)
# No InMemoryHistoryProvider should have been injected
assert not any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers)
async def test_stores_by_default_false_injects_inmemory(client: SupportsChatGetResponse) -> None:
"""Client with STORES_BY_DEFAULT=False (default) should auto-inject InMemoryHistoryProvider."""
from agent_framework._sessions import InMemoryHistoryProvider
agent = Agent(client=client)
session = agent.create_session()
await agent.run("Hello", session=session)
# InMemoryHistoryProvider should have been injected
assert any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers)
async def test_stores_by_default_with_store_false_injects_inmemory(client: SupportsChatGetResponse) -> None:
"""Client with STORES_BY_DEFAULT=True but store=False should still inject InMemoryHistoryProvider."""
from agent_framework._sessions import InMemoryHistoryProvider
client.STORES_BY_DEFAULT = True # type: ignore[attr-defined]
agent = Agent(client=client)
session = agent.create_session()
await agent.run("Hello", session=session, options={"store": False})
# User explicitly disabled server storage, so InMemoryHistoryProvider should be injected
assert any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers)
# endregion
# endregion
python/packages/core/tests/core/test_function_invocation_logic.py
+4 -4
View File
@@ -168,8 +168,8 @@ async def test_function_invocation_inside_aiohttp_server(chat_client_base: Suppo
agent = Agent(client=chat_client_base, tools=[ai_func])
async def handler(request: web.Request) -> web.Response:
thread = agent.get_new_thread()
result = await agent.run("Fix issue", thread=thread)
session = agent.create_session()
result = await agent.run("Fix issue", session=session)
return web.Response(text=result.text or "")
app = web.Application()
@@ -230,8 +230,8 @@ async def test_function_invocation_in_threaded_aiohttp_app(chat_client_base: Sup
async def init_app() -> web.Application:
async def handler(request: web.Request) -> web.Response:
thread = agent.get_new_thread()
result = await agent.run("Fix issue", thread=thread)
session = agent.create_session()
result = await agent.run("Fix issue", session=session)
return web.Response(text=result.text or "")
app = web.Application()
python/packages/core/tests/core/test_mcp.py
+3 -5
View File
@@ -25,8 +25,8 @@ from agent_framework._mcp import (
_get_input_model_from_mcp_tool,
_normalize_mcp_name,
_parse_content_from_mcp,
_parse_tool_result_from_mcp,
_parse_message_from_mcp,
_parse_tool_result_from_mcp,
_prepare_content_for_mcp,
_prepare_message_for_mcp,
logger,
@@ -97,9 +97,7 @@ def test_parse_tool_result_from_mcp():
def test_parse_tool_result_from_mcp_single_text():
"""Test conversion from MCP tool result with a single text item."""
mcp_result = types.CallToolResult(
content=[types.TextContent(type="text", text="Simple result")]
)
mcp_result = types.CallToolResult(content=[types.TextContent(type="text", text="Simple result")])
result = _parse_tool_result_from_mcp(mcp_result)
# Single text item returns just the text
@@ -2590,7 +2588,7 @@ async def test_mcp_tool_filters_framework_kwargs():
chat_options={"some": "option"}, # Should be filtered
tools=[Mock()], # Should be filtered
tool_choice="auto", # Should be filtered
thread=Mock(), # Should be filtered
session=Mock(), # Should be filtered
conversation_id="conv-123", # Should be filtered
options={"metadata": "value"}, # Should be filtered
)
python/packages/core/tests/core/test_memory.py
-136
View File
@@ -1,136 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import sys
from collections.abc import MutableSequence
from typing import Any
from agent_framework import Message
from agent_framework._memory import Context, ContextProvider
class MockContextProvider(ContextProvider):
"""Mock ContextProvider for testing."""
def __init__(self, messages: list[Message] | None = None) -> None:
self.context_messages = messages
self.thread_created_called = False
self.invoked_called = False
self.invoking_called = False
self.thread_created_thread_id = None
self.new_messages = None
self.model_invoking_messages = None
async def thread_created(self, thread_id: str | None) -> None:
"""Track thread_created calls."""
self.thread_created_called = True
self.thread_created_thread_id = thread_id
async def invoked(
self,
request_messages: Any,
response_messages: Any | None = None,
invoke_exception: Exception | None = None,
**kwargs: Any,
) -> None:
"""Track invoked calls."""
self.invoked_called = True
self.new_messages = request_messages
async def invoking(self, messages: Message | MutableSequence[Message], **kwargs: Any) -> Context:
"""Track invoking calls and return context."""
self.invoking_called = True
self.model_invoking_messages = messages
context = Context()
context.messages = self.context_messages
return context
class MinimalContextProvider(ContextProvider):
"""Minimal ContextProvider that only implements the required abstract method.
Used to test the base class default implementations of thread_created,
invoked, __aenter__, and __aexit__.
"""
async def invoking(self, messages: Message | MutableSequence[Message], **kwargs: Any) -> Context:
"""Return empty context."""
return Context()
class TestContext:
"""Tests for Context class."""
def test_context_default_values(self) -> None:
"""Test Context has correct default values."""
context = Context()
assert context.instructions is None
assert context.messages == []
assert context.tools == []
def test_context_with_values(self) -> None:
"""Test Context can be initialized with values."""
messages = [Message(role="user", text="Test message")]
context = Context(instructions="Test instructions", messages=messages)
assert context.instructions == "Test instructions"
assert len(context.messages) == 1
assert context.messages[0].text == "Test message"
class TestContextProvider:
"""Tests for ContextProvider class."""
async def test_thread_created(self) -> None:
"""Test thread_created is called."""
provider = MockContextProvider()
await provider.thread_created("test-thread-id")
assert provider.thread_created_called
assert provider.thread_created_thread_id == "test-thread-id"
async def test_invoked(self) -> None:
"""Test invoked is called."""
provider = MockContextProvider()
message = Message(role="user", text="Test message")
await provider.invoked(message)
assert provider.invoked_called
assert provider.new_messages == message
async def test_invoking(self) -> None:
"""Test invoking is called and returns context."""
provider = MockContextProvider(messages=[Message(role="user", text="Context message")])
message = Message(role="user", text="Test message")
context = await provider.invoking(message)
assert provider.invoking_called
assert provider.model_invoking_messages == message
assert context.messages is not None
assert len(context.messages) == 1
assert context.messages[0].text == "Context message"
async def test_base_thread_created_does_nothing(self) -> None:
"""Test that base ContextProvider.thread_created does nothing by default."""
provider = MinimalContextProvider()
await provider.thread_created("some-thread-id")
await provider.thread_created(None)
async def test_base_invoked_does_nothing(self) -> None:
"""Test that base ContextProvider.invoked does nothing by default."""
provider = MinimalContextProvider()
message = Message(role="user", text="Test")
await provider.invoked(message)
await provider.invoked(message, response_messages=message)
await provider.invoked(message, invoke_exception=Exception("test"))
async def test_base_aenter_returns_self(self) -> None:
"""Test that base ContextProvider.__aenter__ returns self."""
provider = MinimalContextProvider()
async with provider as p:
assert p is provider
async def test_base_aexit_does_nothing(self) -> None:
"""Test that base ContextProvider.__aexit__ handles exceptions gracefully."""
provider = MinimalContextProvider()
await provider.__aexit__(None, None, None)
try:
raise ValueError("test error")
except ValueError:
exc_info = sys.exc_info()
await provider.__aexit__(exc_info[0], exc_info[1], exc_info[2])
python/packages/core/tests/core/test_middleware.py
+26 -26
View File
@@ -56,17 +56,17 @@ class TestAgentContext:
assert context.stream is True
assert context.metadata == metadata
def test_init_with_thread(self, mock_agent: SupportsAgentRun) -> None:
"""Test AgentContext initialization with thread parameter."""
from agent_framework import AgentThread
def test_init_with_session(self, mock_agent: SupportsAgentRun) -> None:
"""Test AgentContext initialization with session parameter."""
from agent_framework import AgentSession
messages = [Message(role="user", text="test")]
thread = AgentThread()
context = AgentContext(agent=mock_agent, messages=messages, thread=thread)
session = AgentSession()
context = AgentContext(agent=mock_agent, messages=messages, session=session)
assert context.agent is mock_agent
assert context.messages == messages
assert context.thread is thread
assert context.session is session
assert context.stream is False
assert context.metadata == {}
@@ -356,23 +356,23 @@ class TestAgentMiddlewarePipeline:
assert updates[1].text == "chunk2"
assert execution_order == ["handler_start", "handler_end"]
async def test_execute_with_thread_in_context(self, mock_agent: SupportsAgentRun) -> None:
"""Test pipeline execution properly passes thread to middleware."""
from agent_framework import AgentThread
async def test_execute_with_session_in_context(self, mock_agent: SupportsAgentRun) -> None:
"""Test pipeline execution properly passes session to middleware."""
from agent_framework import AgentSession
captured_thread = None
captured_session = None
class ThreadCapturingMiddleware(AgentMiddleware):
class SessionCapturingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
nonlocal captured_thread
captured_thread = context.thread
nonlocal captured_session
captured_session = context.session
await call_next()
middleware = ThreadCapturingMiddleware()
middleware = SessionCapturingMiddleware()
pipeline = AgentMiddlewarePipeline(middleware)
messages = [Message(role="user", text="test")]
thread = AgentThread()
context = AgentContext(agent=mock_agent, messages=messages, thread=thread)
session = AgentSession()
context = AgentContext(agent=mock_agent, messages=messages, session=session)
expected_response = AgentResponse(messages=[Message(role="assistant", text="response")])
@@ -381,22 +381,22 @@ class TestAgentMiddlewarePipeline:
result = await pipeline.execute(context, final_handler)
assert result == expected_response
assert captured_thread is thread
assert captured_session is session
async def test_execute_with_no_thread_in_context(self, mock_agent: SupportsAgentRun) -> None:
"""Test pipeline execution when no thread is provided."""
captured_thread = "not_none" # Use string to distinguish from None
async def test_execute_with_no_session_in_context(self, mock_agent: SupportsAgentRun) -> None:
"""Test pipeline execution when no session is provided."""
captured_session = "not_none" # Use string to distinguish from None
class ThreadCapturingMiddleware(AgentMiddleware):
class SessionCapturingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
nonlocal captured_thread
captured_thread = context.thread
nonlocal captured_session
captured_session = context.session
await call_next()
middleware = ThreadCapturingMiddleware()
middleware = SessionCapturingMiddleware()
pipeline = AgentMiddlewarePipeline(middleware)
messages = [Message(role="user", text="test")]
context = AgentContext(agent=mock_agent, messages=messages, thread=None)
context = AgentContext(agent=mock_agent, messages=messages, session=None)
expected_response = AgentResponse(messages=[Message(role="assistant", text="response")])
@@ -405,7 +405,7 @@ class TestAgentMiddlewarePipeline:
result = await pipeline.execute(context, final_handler)
assert result == expected_response
assert captured_thread is None
assert captured_session is None
class TestFunctionMiddlewarePipeline:
python/packages/core/tests/core/test_middleware_with_agent.py
+16 -19
View File
@@ -1405,19 +1405,19 @@ class TestMiddlewareDecoratorLogic:
assert test_function_middleware._middleware_type == MiddlewareType.FUNCTION # type: ignore[attr-defined]
class TestChatAgentThreadBehavior:
"""Test cases for thread behavior in AgentContext across multiple runs."""
class TestChatAgentSessionBehavior:
"""Test cases for session behavior in AgentContext across multiple runs."""
async def test_agent_context_thread_behavior_across_multiple_runs(self, client: "MockChatClient") -> None:
"""Test that AgentContext.thread property behaves correctly across multiple agent runs."""
async def test_agent_context_session_behavior_across_multiple_runs(self, client: "MockChatClient") -> None:
"""Test that AgentContext.session property behaves correctly across multiple agent runs."""
thread_states: list[dict[str, Any]] = []
class ThreadTrackingMiddleware(AgentMiddleware):
class SessionTrackingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
# Capture state before next() call
thread_messages = []
if context.thread and context.thread.message_store:
thread_messages = await context.thread.message_store.list_messages()
if context.session and context.session.state.get("memory"):
thread_messages = context.session.state.get("memory", {}).get("messages", [])
before_state = {
"before_next": True,
@@ -1432,8 +1432,8 @@ class TestChatAgentThreadBehavior:
# Capture state after next() call
thread_messages_after = []
if context.thread and context.thread.message_store:
thread_messages_after = await context.thread.message_store.list_messages()
if context.session and context.session.state.get("memory"):
thread_messages_after = context.session.state.get("memory", {}).get("messages", [])
after_state = {
"before_next": False,
@@ -1444,19 +1444,16 @@ class TestChatAgentThreadBehavior:
}
thread_states.append(after_state)
# Import the ChatMessageStore to configure the agent with a message store factory
from agent_framework import ChatMessageStore
# Create Agent with session tracking middleware
middleware = SessionTrackingMiddleware()
agent = Agent(client=client, middleware=[middleware])
# Create Agent with thread tracking middleware and a message store factory
middleware = ThreadTrackingMiddleware()
agent = Agent(client=client, middleware=[middleware], chat_message_store_factory=ChatMessageStore)
# Create a thread that will persist messages between runs
thread = agent.get_new_thread()
# Create a session that will persist messages between runs
session = agent.create_session()
# First run
first_messages = [Message(role="user", text="first message")]
first_response = await agent.run(first_messages, thread=thread)
first_response = await agent.run(first_messages, session=session)
# Verify first response
assert first_response is not None
@@ -1464,7 +1461,7 @@ class TestChatAgentThreadBehavior:
# Second run - use the same thread
second_messages = [Message(role="user", text="second message")]
second_response = await agent.run(second_messages, thread=thread)
second_response = await agent.run(second_messages, session=session)
# Verify second response
assert second_response is not None
python/packages/core/tests/core/test_observability.py
+18 -18
View File
@@ -441,19 +441,19 @@ def mock_chat_agent():
self.description = "Test agent description"
self.default_options: dict[str, Any] = {"model_id": "TestModel"}
def run(self, messages=None, *, thread=None, stream=False, **kwargs):
def run(self, messages=None, *, session=None, stream=False, **kwargs):
if stream:
return self._run_stream_impl(messages=messages, **kwargs)
return self._run_impl(messages=messages, **kwargs)
async def _run_impl(self, messages=None, *, thread=None, **kwargs):
async def _run_impl(self, messages=None, *, session=None, **kwargs):
return AgentResponse(
messages=[Message("assistant", ["Agent response"])],
usage_details=UsageDetails(input_token_count=15, output_token_count=25),
response_id="test_response_id",
)
async def _run_stream_impl(self, messages=None, *, thread=None, **kwargs):
async def _run_stream_impl(self, messages=None, *, session=None, **kwargs):
from agent_framework import AgentResponse, AgentResponseUpdate, ResponseStream
async def _stream():
@@ -1572,12 +1572,12 @@ async def test_agent_observability(span_exporter: InMemorySpanExporter, enable_s
messages=None,
*,
stream: bool = False,
thread=None,
session=None,
**kwargs,
):
if stream:
return ResponseStream(
self._run_stream(messages=messages, thread=thread),
self._run_stream(messages=messages, session=session),
finalizer=lambda x: AgentResponse.from_updates(x),
)
return AgentResponse(messages=[Message("assistant", ["Test response"])])
@@ -1586,7 +1586,7 @@ async def test_agent_observability(span_exporter: InMemorySpanExporter, enable_s
self,
messages=None,
*,
thread=None,
session=None,
**kwargs,
):
from agent_framework import AgentResponseUpdate
@@ -1635,7 +1635,7 @@ async def test_agent_observability_with_exception(span_exporter: InMemorySpanExp
def default_options(self):
return self._default_options
async def run(self, messages=None, *, stream: bool = False, thread=None, **kwargs):
async def run(self, messages=None, *, stream: bool = False, session=None, **kwargs):
raise RuntimeError("Agent failed")
class FailingAgent(AgentTelemetryLayer, _FailingAgent):
@@ -1685,15 +1685,15 @@ async def test_agent_streaming_observability(span_exporter: InMemorySpanExporter
def default_options(self):
return self._default_options
def run(self, messages=None, *, stream=False, thread=None, **kwargs):
def run(self, messages=None, *, stream=False, session=None, **kwargs):
if stream:
return self._run_stream_impl(messages=messages, **kwargs)
return self._run_impl(messages=messages, **kwargs)
async def _run_impl(self, messages=None, *, thread=None, **kwargs):
async def _run_impl(self, messages=None, *, session=None, **kwargs):
return AgentResponse(messages=[Message("assistant", ["Test"])])
def _run_stream_impl(self, messages=None, *, thread=None, **kwargs):
def _run_stream_impl(self, messages=None, *, session=None, **kwargs):
async def _stream():
yield AgentResponseUpdate(contents=[Content.from_text("Hello ")], role="assistant")
yield AgentResponseUpdate(contents=[Content.from_text("World")], role="assistant")
@@ -1822,15 +1822,15 @@ async def test_agent_streaming_exception(span_exporter: InMemorySpanExporter, en
def default_options(self):
return self._default_options
def run(self, messages=None, *, stream=False, thread=None, **kwargs):
def run(self, messages=None, *, stream=False, session=None, **kwargs):
if stream:
return self._run_stream_impl(messages=messages, **kwargs)
return self._run_impl(messages=messages, **kwargs)
async def _run_impl(self, messages=None, *, thread=None, **kwargs):
async def _run_impl(self, messages=None, *, session=None, **kwargs):
return AgentResponse(messages=[])
def _run_stream_impl(self, messages=None, *, thread=None, **kwargs):
def _run_stream_impl(self, messages=None, *, session=None, **kwargs):
async def _stream():
yield AgentResponseUpdate(contents=[Content.from_text("Starting")], role="assistant")
raise RuntimeError("Stream failed")
@@ -1919,7 +1919,7 @@ async def test_agent_when_disabled(span_exporter: InMemorySpanExporter):
def default_options(self):
return self._default_options
async def run(self, messages=None, *, stream: bool = False, thread=None, **kwargs):
async def run(self, messages=None, *, stream: bool = False, session=None, **kwargs):
if stream:
return ResponseStream(
self._run_stream(messages=messages, **kwargs),
@@ -1927,7 +1927,7 @@ async def test_agent_when_disabled(span_exporter: InMemorySpanExporter):
)
return AgentResponse(messages=[])
async def _run_stream(self, messages=None, *, thread=None, **kwargs):
async def _run_stream(self, messages=None, *, session=None, **kwargs):
from agent_framework import AgentResponseUpdate
yield AgentResponseUpdate(contents=[Content.from_text("test")], role="assistant")
@@ -1974,15 +1974,15 @@ async def test_agent_streaming_when_disabled(span_exporter: InMemorySpanExporter
def default_options(self):
return self._default_options
def run(self, messages=None, *, stream=False, thread=None, **kwargs):
def run(self, messages=None, *, stream=False, session=None, **kwargs):
if stream:
return self._run_stream(messages=messages, **kwargs)
return self._run(messages=messages, **kwargs)
async def _run(self, messages=None, *, thread=None, **kwargs):
async def _run(self, messages=None, *, session=None, **kwargs):
return AgentResponse(messages=[])
async def _run_stream(self, messages=None, *, thread=None, **kwargs):
async def _run_stream(self, messages=None, *, session=None, **kwargs):
yield AgentResponseUpdate(contents=[Content.from_text("test")], role="assistant")
class TestAgent(AgentTelemetryLayer, _TestAgent):
python/packages/core/tests/core/test_settings.py
+92
View File
@@ -28,6 +28,12 @@ class SecretSettings(TypedDict, total=False):
username: str | None
class ExclusiveSettings(TypedDict, total=False):
source_a: str | None
source_b: str | None
other: str | None
class TestLoadSettingsBasic:
"""Test basic load_settings functionality."""
@@ -236,3 +242,89 @@ class TestOverrideTypeValidation:
assert isinstance(settings["api_key"], SecretString)
assert settings["api_key"] == "plain-string"
class TestMutuallyExclusive:
"""Test mutually exclusive field validation via tuple entries in required_fields."""
def test_exactly_one_set_passes(self) -> None:
settings = load_settings(
ExclusiveSettings,
env_prefix="TEST_",
required_fields=[("source_a", "source_b")],
source_a="value-a",
)
assert settings["source_a"] == "value-a"
assert settings["source_b"] is None
def test_none_set_raises(self) -> None:
from agent_framework.exceptions import SettingNotFoundError
with pytest.raises(SettingNotFoundError, match="none was set"):
load_settings(
ExclusiveSettings,
env_prefix="TEST_",
required_fields=[("source_a", "source_b")],
)
def test_both_set_raises(self) -> None:
from agent_framework.exceptions import SettingNotFoundError
with pytest.raises(SettingNotFoundError, match="multiple were set"):
load_settings(
ExclusiveSettings,
env_prefix="TEST_",
required_fields=[("source_a", "source_b")],
source_a="a",
source_b="b",
)
def test_env_var_counts_as_set(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("TEST_SOURCE_B", "env-b")
settings = load_settings(
ExclusiveSettings,
env_prefix="TEST_",
required_fields=[("source_a", "source_b")],
)
assert settings["source_b"] == "env-b"
def test_env_var_and_override_both_set_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
from agent_framework.exceptions import SettingNotFoundError
monkeypatch.setenv("TEST_SOURCE_B", "env-b")
with pytest.raises(SettingNotFoundError, match="multiple were set"):
load_settings(
ExclusiveSettings,
env_prefix="TEST_",
required_fields=[("source_a", "source_b")],
source_a="a",
)
def test_other_fields_unaffected(self) -> None:
settings = load_settings(
ExclusiveSettings,
env_prefix="TEST_",
required_fields=[("source_a", "source_b")],
source_a="a",
other="extra",
)
assert settings["source_a"] == "a"
assert settings["other"] == "extra"
def test_mixed_required_and_exclusive(self) -> None:
settings = load_settings(
ExclusiveSettings,
env_prefix="TEST_",
required_fields=["other", ("source_a", "source_b")],
source_b="b",
other="required-val",
)
assert settings["other"] == "required-val"
assert settings["source_b"] == "b"
assert settings["source_a"] is None
python/packages/core/tests/core/test_threads.py
-600
View File
@@ -1,600 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import Sequence
from typing import Any
import pytest
from agent_framework import AgentThread, ChatMessageStore, Message
from agent_framework._threads import AgentThreadState, ChatMessageStoreState
from agent_framework.exceptions import AgentThreadException
class MockChatMessageStore:
"""Mock implementation of ChatMessageStoreProtocol for testing."""
def __init__(self, messages: list[Message] | None = None) -> None:
self._messages = messages or []
self._serialize_calls = 0
self._deserialize_calls = 0
async def list_messages(self) -> list[Message]:
return self._messages
async def add_messages(self, messages: Sequence[Message]) -> None:
self._messages.extend(messages)
async def serialize(self, **kwargs: Any) -> Any:
self._serialize_calls += 1
return {"messages": [msg.__dict__ for msg in self._messages], "kwargs": kwargs}
async def update_from_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
self._deserialize_calls += 1
if serialized_store_state and "messages" in serialized_store_state:
self._messages = serialized_store_state["messages"]
@classmethod
async def deserialize(cls, serialized_store_state: Any, **kwargs: Any) -> "MockChatMessageStore":
instance = cls()
await instance.update_from_state(serialized_store_state, **kwargs)
return instance
@pytest.fixture
def sample_messages() -> list[Message]:
"""Fixture providing sample chat messages for testing."""
return [
Message(role="user", text="Hello", message_id="msg1"),
Message(role="assistant", text="Hi there!", message_id="msg2"),
Message(role="user", text="How are you?", message_id="msg3"),
]
@pytest.fixture
def sample_message() -> Message:
"""Fixture providing a single sample chat message for testing."""
return Message(role="user", text="Test message", message_id="test1")
class TestAgentThread:
"""Test cases for AgentThread class."""
def test_init_with_no_parameters(self) -> None:
"""Test AgentThread initialization with no parameters."""
thread = AgentThread()
assert thread.service_thread_id is None
assert thread.message_store is None
def test_init_with_service_thread_id(self) -> None:
"""Test AgentThread initialization with service_thread_id."""
service_thread_id = "test-conversation-123"
thread = AgentThread(service_thread_id=service_thread_id)
assert thread.service_thread_id == service_thread_id
assert thread.message_store is None
def test_init_with_message_store(self) -> None:
"""Test AgentThread initialization with message_store."""
store = ChatMessageStore()
thread = AgentThread(message_store=store)
assert thread.service_thread_id is None
assert thread.message_store is store
def test_service_thread_id_property_setter(self) -> None:
"""Test service_thread_id property setter."""
thread = AgentThread()
service_thread_id = "test-conversation-456"
thread.service_thread_id = service_thread_id
assert thread.service_thread_id == service_thread_id
def test_service_thread_id_setter_with_existing_message_store_raises_error(self) -> None:
"""Test that setting service_thread_id when message_store exists raises AgentThreadException."""
store = ChatMessageStore()
thread = AgentThread(message_store=store)
with pytest.raises(AgentThreadException, match="Only the service_thread_id or message_store may be set"):
thread.service_thread_id = "test-conversation-789"
def test_service_thread_id_setter_with_none_values(self) -> None:
"""Test service_thread_id setter with None values does nothing."""
thread = AgentThread()
thread.service_thread_id = None # Should not raise error
assert thread.service_thread_id is None
def test_message_store_property_setter(self) -> None:
"""Test message_store property setter."""
thread = AgentThread()
store = ChatMessageStore()
thread.message_store = store
assert thread.message_store is store
def test_message_store_setter_with_existing_service_thread_id_raises_error(self) -> None:
"""Test that setting message_store when service_thread_id exists raises AgentThreadException."""
service_thread_id = "test-conversation-999"
thread = AgentThread(service_thread_id=service_thread_id)
store = ChatMessageStore()
with pytest.raises(AgentThreadException, match="Only the service_thread_id or message_store may be set"):
thread.message_store = store
def test_message_store_setter_with_none_values(self) -> None:
"""Test message_store setter with None values does nothing."""
thread = AgentThread()
thread.message_store = None # Should not raise error
assert thread.message_store is None
async def test_get_messages_with_message_store(self, sample_messages: list[Message]) -> None:
"""Test get_messages when message_store is set."""
store = ChatMessageStore(sample_messages)
thread = AgentThread(message_store=store)
assert thread.message_store is not None
messages: list[Message] = await thread.message_store.list_messages()
assert messages is not None
assert len(messages) == 3
assert messages[0].text == "Hello"
assert messages[1].text == "Hi there!"
assert messages[2].text == "How are you?"
async def test_get_messages_with_no_message_store(self) -> None:
"""Test get_messages when no message_store is set."""
thread = AgentThread()
assert thread.message_store is None
async def test_on_new_messages_with_service_thread_id(self, sample_message: Message) -> None:
"""Test _on_new_messages when service_thread_id is set (should do nothing)."""
thread = AgentThread(service_thread_id="test-conv")
await thread.on_new_messages(sample_message)
# Should not create a message store
assert thread.message_store is None
async def test_on_new_messages_single_message_creates_store(self, sample_message: Message) -> None:
"""Test _on_new_messages with single message creates ChatMessageStore."""
thread = AgentThread()
await thread.on_new_messages(sample_message)
assert thread.message_store is not None
assert isinstance(thread.message_store, ChatMessageStore)
messages = await thread.message_store.list_messages()
assert len(messages) == 1
assert messages[0].text == "Test message"
async def test_on_new_messages_multiple_messages(self, sample_messages: list[Message]) -> None:
"""Test _on_new_messages with multiple messages."""
thread = AgentThread()
await thread.on_new_messages(sample_messages)
assert thread.message_store is not None
messages = await thread.message_store.list_messages()
assert len(messages) == 3
async def test_on_new_messages_with_existing_store(self, sample_message: Message) -> None:
"""Test _on_new_messages adds to existing message store."""
initial_messages = [Message(role="user", text="Initial", message_id="init1")]
store = ChatMessageStore(initial_messages)
thread = AgentThread(message_store=store)
await thread.on_new_messages(sample_message)
assert thread.message_store is not None
messages = await thread.message_store.list_messages()
assert len(messages) == 2
assert messages[0].text == "Initial"
assert messages[1].text == "Test message"
async def test_deserialize_with_service_thread_id(self) -> None:
"""Test _deserialize with service_thread_id."""
serialized_data = {"service_thread_id": "test-conv-123", "chat_message_store_state": None}
thread = await AgentThread.deserialize(serialized_data)
assert thread.service_thread_id == "test-conv-123"
assert thread.message_store is None
async def test_deserialize_with_store_state(self, sample_messages: list[Message]) -> None:
"""Test _deserialize with chat_message_store_state."""
store_state = {"messages": sample_messages}
serialized_data = {"service_thread_id": None, "chat_message_store_state": store_state}
thread = await AgentThread.deserialize(serialized_data)
assert thread.service_thread_id is None
assert thread.message_store is not None
assert isinstance(thread.message_store, ChatMessageStore)
async def test_deserialize_with_no_state(self) -> None:
"""Test _deserialize with no state."""
thread = AgentThread()
serialized_data = {"service_thread_id": None, "chat_message_store_state": None}
await thread.deserialize(serialized_data)
assert thread.service_thread_id is None
assert thread.message_store is None
async def test_deserialize_with_existing_store(self) -> None:
"""Test _deserialize with existing message store."""
store = MockChatMessageStore()
thread = AgentThread(message_store=store)
serialized_data: dict[str, Any] = {
"service_thread_id": None,
"chat_message_store_state": {"messages": [Message(role="user", text="test")]},
}
await thread.update_from_thread_state(serialized_data)
assert store._messages
assert store._messages[0].text == "test"
async def test_serialize_with_service_thread_id(self) -> None:
"""Test serialize with service_thread_id."""
thread = AgentThread(service_thread_id="test-conv-456")
result = await thread.serialize()
assert result["service_thread_id"] == "test-conv-456"
assert result["chat_message_store_state"] is None
async def test_serialize_with_message_store(self) -> None:
"""Test serialize with message_store."""
store = MockChatMessageStore()
thread = AgentThread(message_store=store)
result = await thread.serialize()
assert result["service_thread_id"] is None
assert result["chat_message_store_state"] is not None
assert store._serialize_calls == 1 # pyright: ignore[reportPrivateUsage]
async def test_serialize_with_no_state(self) -> None:
"""Test serialize with no state."""
thread = AgentThread()
result = await thread.serialize()
assert result["service_thread_id"] is None
assert result["chat_message_store_state"] is None
async def test_serialize_with_kwargs(self) -> None:
"""Test serialize passes kwargs to message store."""
store = MockChatMessageStore()
thread = AgentThread(message_store=store)
await thread.serialize(custom_param="test_value")
assert store._serialize_calls == 1 # pyright: ignore[reportPrivateUsage]
async def test_serialize_round_trip_messages(self, sample_messages: list[Message]) -> None:
"""Test a roundtrip of the serialization."""
store = ChatMessageStore(sample_messages)
thread = AgentThread(message_store=store)
new_thread = await AgentThread.deserialize(await thread.serialize())
assert new_thread.message_store is not None
new_messages = await new_thread.message_store.list_messages()
assert len(new_messages) == len(sample_messages)
assert {new.text for new in new_messages} == {orig.text for orig in sample_messages}
async def test_serialize_round_trip_thread_id(self) -> None:
"""Test a roundtrip of the serialization."""
thread = AgentThread(service_thread_id="test-1234")
new_thread = await AgentThread.deserialize(await thread.serialize())
assert new_thread.message_store is None
assert new_thread.service_thread_id == "test-1234"
class TestChatMessageList:
"""Test cases for ChatMessageStore class."""
def test_init_empty(self) -> None:
"""Test ChatMessageStore initialization with no messages."""
store = ChatMessageStore()
assert len(store.messages) == 0
def test_init_with_messages(self, sample_messages: list[Message]) -> None:
"""Test ChatMessageStore initialization with messages."""
store = ChatMessageStore(sample_messages)
assert len(store.messages) == 3
async def test_add_messages(self, sample_messages: list[Message]) -> None:
"""Test adding messages to the store."""
store = ChatMessageStore()
await store.add_messages(sample_messages)
assert len(store.messages) == 3
messages = await store.list_messages()
assert messages[0].text == "Hello"
async def test_get_messages(self, sample_messages: list[Message]) -> None:
"""Test getting messages from the store."""
store = ChatMessageStore(sample_messages)
messages = await store.list_messages()
assert len(messages) == 3
assert messages[0].message_id == "msg1"
async def test_serialize_state(self, sample_messages: list[Message]) -> None:
"""Test serializing store state."""
store = ChatMessageStore(sample_messages)
result = await store.serialize()
assert "messages" in result
assert len(result["messages"]) == 3
async def test_serialize_state_empty(self) -> None:
"""Test serializing empty store state."""
store = ChatMessageStore()
result = await store.serialize()
assert "messages" in result
assert len(result["messages"]) == 0
async def test_deserialize_state(self, sample_messages: list[Message]) -> None:
"""Test deserializing store state."""
store = ChatMessageStore()
state_data = {"messages": sample_messages}
await store.update_from_state(state_data)
messages = await store.list_messages()
assert len(messages) == 3
assert messages[0].text == "Hello"
async def test_deserialize_state_none(self) -> None:
"""Test deserializing None state."""
store = ChatMessageStore()
await store.update_from_state(None)
assert len(store.messages) == 0
async def test_deserialize_state_empty(self) -> None:
"""Test deserializing empty state."""
store = ChatMessageStore()
await store.update_from_state({})
assert len(store.messages) == 0
class TestStoreState:
"""Test cases for ChatMessageStoreState class."""
def test_init(self, sample_messages: list[Message]) -> None:
"""Test ChatMessageStoreState initialization."""
state = ChatMessageStoreState(messages=sample_messages)
assert len(state.messages) == 3
assert state.messages[0].text == "Hello"
def test_init_empty(self) -> None:
"""Test ChatMessageStoreState initialization with empty messages."""
state = ChatMessageStoreState(messages=[])
assert len(state.messages) == 0
def test_init_none(self) -> None:
"""Test ChatMessageStoreState initialization with None messages."""
state = ChatMessageStoreState(messages=None)
assert len(state.messages) == 0
def test_init_no_messages_arg(self) -> None:
"""Test ChatMessageStoreState initialization without messages argument."""
state = ChatMessageStoreState()
assert len(state.messages) == 0
class TestThreadState:
"""Test cases for AgentThreadState class."""
def test_init_with_service_thread_id(self) -> None:
"""Test AgentThreadState initialization with service_thread_id."""
state = AgentThreadState(service_thread_id="test-conv-123")
assert state.service_thread_id == "test-conv-123"
assert state.chat_message_store_state is None
def test_init_with_chat_message_store_state(self) -> None:
"""Test AgentThreadState initialization with chat_message_store_state."""
store_data: dict[str, Any] = {"messages": []}
state = AgentThreadState.from_dict({"chat_message_store_state": store_data})
assert state.service_thread_id is None
assert state.chat_message_store_state.messages == []
def test_init_with_both(self) -> None:
"""Test AgentThreadState initialization with both parameters."""
store_data: dict[str, Any] = {"messages": []}
with pytest.raises(AgentThreadException):
AgentThreadState(service_thread_id="test-conv-123", chat_message_store_state=store_data)
def test_init_defaults(self) -> None:
"""Test AgentThreadState initialization with defaults."""
state = AgentThreadState()
assert state.service_thread_id is None
assert state.chat_message_store_state is None
def test_init_with_chat_message_store_state_no_messages(self) -> None:
"""Test AgentThreadState initialization with chat_message_store_state without messages field.
This tests the scenario where a custom ChatMessageStore (like RedisChatMessageStore)
serializes its state without a 'messages' field, containing only configuration data
like thread_id, redis_url, etc.
"""
store_data: dict[str, Any] = {
"type": "redis_store_state",
"thread_id": "test_thread_123",
"redis_url": "redis://localhost:6379",
"key_prefix": "chat_messages",
}
state = AgentThreadState.from_dict({"chat_message_store_state": store_data})
assert state.service_thread_id is None
assert state.chat_message_store_state is not None
assert state.chat_message_store_state.messages == []
def test_init_with_chat_message_store_state_object(self) -> None:
"""Test AgentThreadState initialization with ChatMessageStoreState object."""
store_state = ChatMessageStoreState(messages=[Message(role="user", text="test")])
state = AgentThreadState(chat_message_store_state=store_state)
assert state.service_thread_id is None
assert state.chat_message_store_state is store_state
assert len(state.chat_message_store_state.messages) == 1
def test_init_with_invalid_chat_message_store_state_type(self) -> None:
"""Test AgentThreadState initialization with invalid chat_message_store_state type."""
with pytest.raises(TypeError, match="Could not parse ChatMessageStoreState"):
AgentThreadState(chat_message_store_state="invalid_type") # type: ignore[arg-type]
class TestChatMessageStoreStateEdgeCases:
"""Additional edge case tests for ChatMessageStoreState."""
def test_init_with_invalid_messages_type(self) -> None:
"""Test ChatMessageStoreState initialization with invalid messages type."""
with pytest.raises(TypeError, match="Messages should be a list"):
ChatMessageStoreState(messages="invalid") # type: ignore[arg-type]
def test_init_with_dict_messages(self) -> None:
"""Test ChatMessageStoreState initialization with dict messages."""
messages = [
{"role": "user", "text": "Hello"},
{"role": "assistant", "text": "Hi there!"},
]
state = ChatMessageStoreState(messages=messages)
assert len(state.messages) == 2
assert isinstance(state.messages[0], Message)
assert state.messages[0].text == "Hello"
class TestChatMessageStoreEdgeCases:
"""Additional edge case tests for ChatMessageStore."""
async def test_deserialize_class_method(self) -> None:
"""Test ChatMessageStore.deserialize class method."""
serialized_data = {
"messages": [
{"role": "user", "text": "Hello", "message_id": "msg1"},
]
}
store = await ChatMessageStore.deserialize(serialized_data)
assert isinstance(store, ChatMessageStore)
messages = await store.list_messages()
assert len(messages) == 1
assert messages[0].text == "Hello"
async def test_deserialize_empty_state(self) -> None:
"""Test ChatMessageStore.deserialize with empty state."""
serialized_data: dict[str, Any] = {"messages": []}
store = await ChatMessageStore.deserialize(serialized_data)
assert isinstance(store, ChatMessageStore)
messages = await store.list_messages()
assert len(messages) == 0
class TestAgentThreadEdgeCases:
"""Additional edge case tests for AgentThread."""
def test_is_initialized_with_service_thread_id(self) -> None:
"""Test is_initialized property when service_thread_id is set."""
thread = AgentThread(service_thread_id="test-123")
assert thread.is_initialized is True
def test_is_initialized_with_message_store(self) -> None:
"""Test is_initialized property when message_store is set."""
store = ChatMessageStore()
thread = AgentThread(message_store=store)
assert thread.is_initialized is True
def test_is_initialized_with_nothing(self) -> None:
"""Test is_initialized property when nothing is set."""
thread = AgentThread()
assert thread.is_initialized is False
async def test_deserialize_with_custom_message_store(self) -> None:
"""Test deserialize using a custom message store."""
serialized_data = {
"service_thread_id": None,
"chat_message_store_state": {
"messages": [{"role": "user", "text": "Hello"}],
},
}
custom_store = MockChatMessageStore()
thread = await AgentThread.deserialize(serialized_data, message_store=custom_store)
assert thread.message_store is custom_store
messages = await custom_store.list_messages()
assert len(messages) == 1
async def test_deserialize_with_failing_message_store_raises(self) -> None:
"""Test deserialize raises AgentThreadException when message store fails."""
class FailingStore:
async def add_messages(self, messages: Sequence[Message], **kwargs: Any) -> None:
raise RuntimeError("Store failed")
serialized_data = {
"service_thread_id": None,
"chat_message_store_state": {
"messages": [{"role": "user", "text": "Hello"}],
},
}
failing_store = FailingStore()
with pytest.raises(AgentThreadException, match="Failed to deserialize"):
await AgentThread.deserialize(serialized_data, message_store=failing_store)
async def test_update_from_thread_state_with_service_thread_id(self) -> None:
"""Test update_from_thread_state sets service_thread_id."""
thread = AgentThread()
serialized_data = {"service_thread_id": "new-thread-id"}
await thread.update_from_thread_state(serialized_data)
assert thread.service_thread_id == "new-thread-id"
async def test_update_from_thread_state_with_empty_chat_state(self) -> None:
"""Test update_from_thread_state with empty chat_message_store_state."""
thread = AgentThread()
serialized_data = {"service_thread_id": None, "chat_message_store_state": None}
await thread.update_from_thread_state(serialized_data)
assert thread.message_store is None
async def test_update_from_thread_state_creates_message_store(self) -> None:
"""Test update_from_thread_state creates message store if not existing."""
thread = AgentThread()
serialized_data = {
"service_thread_id": None,
"chat_message_store_state": {
"messages": [{"role": "user", "text": "Hello"}],
},
}
await thread.update_from_thread_state(serialized_data)
assert thread.message_store is not None
messages = await thread.message_store.list_messages()
assert len(messages) == 1
python/packages/core/tests/openai/test_openai_assistants_client.py
+24 -24
View File
@@ -14,7 +14,7 @@ from agent_framework import (
Agent,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
ChatResponse,
ChatResponseUpdate,
Content,
@@ -1264,70 +1264,70 @@ async def test_openai_assistants_agent_basic_run_streaming():
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_agent_thread_persistence():
"""Test Agent thread persistence across runs with OpenAIAssistantsClient."""
async def test_openai_assistants_agent_session_persistence():
"""Test Agent session persistence across runs with OpenAIAssistantsClient."""
async with Agent(
client=OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# Create a new session that will be reused
session = agent.create_session()
# First message - establish context
first_response = await agent.run(
"Remember this number: 42. What number did I just tell you to remember?", thread=thread
"Remember this number: 42. What number did I just tell you to remember?", session=session
)
assert isinstance(first_response, AgentResponse)
assert "42" in first_response.text
# Second message - test conversation memory
second_response = await agent.run(
"What number did I tell you to remember in my previous message?", thread=thread
"What number did I tell you to remember in my previous message?", session=session
)
assert isinstance(second_response, AgentResponse)
assert "42" in second_response.text
# Verify thread has been populated with conversation ID
assert thread.service_thread_id is not None
# Verify session has been populated with conversation ID
assert session.service_session_id is not None
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_assistants_agent_existing_thread_id():
"""Test Agent with existing thread ID to continue conversations across agent instances."""
# First, create a conversation and capture the thread ID
existing_thread_id = None
async def test_openai_assistants_agent_existing_session_id():
"""Test Agent with existing session ID to continue conversations across agent instances."""
# First, create a conversation and capture the session ID
existing_session_id = None
async with Agent(
client=OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Start a conversation and get the thread ID
thread = agent.get_new_thread()
response1 = await agent.run("What's the weather in Paris?", thread=thread)
# Start a conversation and get the session ID
session = agent.create_session()
response1 = await agent.run("What's the weather in Paris?", session=session)
# Validate first response
assert isinstance(response1, AgentResponse)
assert response1.text is not None
assert any(word in response1.text.lower() for word in ["weather", "paris"])
# The thread ID is set after the first response
existing_thread_id = thread.service_thread_id
assert existing_thread_id is not None
# The session ID is set after the first response
existing_session_id = session.service_session_id
assert existing_session_id is not None
# Now continue with the same thread ID in a new agent instance
# Now continue with the same session ID in a new agent instance
async with Agent(
client=OpenAIAssistantsClient(thread_id=existing_thread_id),
client=OpenAIAssistantsClient(thread_id=existing_session_id),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
# Create a thread with the existing ID
thread = AgentThread(service_thread_id=existing_thread_id)
# Create a session with the existing ID
session = AgentSession(service_session_id=existing_session_id)
# Ask about the previous conversation
response2 = await agent.run("What was the last city I asked about?", thread=thread)
response2 = await agent.run("What was the last city I asked about?", session=session)
# Validate that the agent remembers the previous conversation
assert isinstance(response2, AgentResponse)
python/packages/core/tests/workflow/__init__.py
View File
python/packages/core/tests/workflow/test_agent_executor.py
+43 -65
View File
@@ -7,9 +7,8 @@ from agent_framework import (
AgentExecutor,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
ChatMessageStore,
Content,
Message,
ResponseStream,
@@ -21,7 +20,7 @@ from agent_framework.orchestrations import SequentialBuilder
class _CountingAgent(BaseAgent):
"""Agent that echoes messages with a counter to verify thread state persistence."""
"""Agent that echoes messages with a counter to verify session state persistence."""
def __init__(self, **kwargs: Any):
super().__init__(**kwargs)
@@ -32,7 +31,7 @@ class _CountingAgent(BaseAgent):
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
self.call_count += 1
@@ -52,22 +51,22 @@ class _CountingAgent(BaseAgent):
async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
"""Test that workflow checkpoint stores AgentExecutor's cache and thread states and restores them correctly."""
"""Test that workflow checkpoint stores AgentExecutor's cache and session states and restores them correctly."""
storage = InMemoryCheckpointStorage()
# Create initial agent with a custom thread that has a message store
# Create initial agent with a custom session
initial_agent = _CountingAgent(id="test_agent", name="TestAgent")
initial_thread = AgentThread(message_store=ChatMessageStore())
initial_session = AgentSession()
# Add some initial messages to the thread to verify thread state persistence
# Add some initial messages to the session state to verify session state persistence
initial_messages = [
Message(role="user", text="Initial message 1"),
Message(role="assistant", text="Initial response 1"),
]
await initial_thread.on_new_messages(initial_messages)
initial_session.state["history"] = {"messages": initial_messages}
# Create AgentExecutor with the thread
executor = AgentExecutor(initial_agent, agent_thread=initial_thread)
# Create AgentExecutor with the session
executor = AgentExecutor(initial_agent, session=initial_session)
# Build workflow with checkpointing enabled
wf = SequentialBuilder(participants=[executor], checkpoint_storage=storage).build()
@@ -95,7 +94,7 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
checkpoints.sort(key=lambda cp: cp.timestamp)
restore_checkpoint = checkpoints[1]
# Verify checkpoint contains executor state with both cache and thread
# Verify checkpoint contains executor state with both cache and session
assert "_executor_state" in restore_checkpoint.state
executor_states = restore_checkpoint.state["_executor_state"]
assert isinstance(executor_states, dict)
@@ -103,13 +102,12 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
executor_state = executor_states[executor.id] # type: ignore[index]
assert "cache" in executor_state, "Checkpoint should store executor cache state"
assert "agent_thread" in executor_state, "Checkpoint should store executor thread state"
assert "agent_session" in executor_state, "Checkpoint should store executor session state"
# Verify thread state includes message store
thread_state = executor_state["agent_thread"] # type: ignore[index]
assert "chat_message_store_state" in thread_state, "Thread state should include message store"
chat_store_state = thread_state["chat_message_store_state"] # type: ignore[index]
assert "messages" in chat_store_state, "Message store state should include messages"
# Verify session state structure
session_state = executor_state["agent_session"] # type: ignore[index]
assert "session_id" in session_state, "Session state should include session_id"
assert "state" in session_state, "Session state should include state dict"
# Verify checkpoint contains pending requests from agents and responses to be sent
assert "pending_agent_requests" in executor_state
@@ -118,8 +116,8 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
# Create a new agent and executor for restoration
# This simulates starting from a fresh state and restoring from checkpoint
restored_agent = _CountingAgent(id="test_agent", name="TestAgent")
restored_thread = AgentThread(message_store=ChatMessageStore())
restored_executor = AgentExecutor(restored_agent, agent_thread=restored_thread)
restored_session = AgentSession()
restored_executor = AgentExecutor(restored_agent, session=restored_session)
# Verify the restored agent starts with a fresh state
assert restored_agent.call_count == 0
@@ -140,39 +138,27 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
assert resumed_output is not None
# Verify the restored executor's state matches the original
# The cache should be restored (though it may be cleared after processing)
# The thread should have all messages including those from the initial state
message_store = restored_executor._agent_thread.message_store # type: ignore[reportPrivateUsage]
assert message_store is not None
thread_messages = await message_store.list_messages()
# Thread should contain:
# 1. Initial messages from before the checkpoint (2 messages)
# 2. User message from first run (1 message)
# 3. Assistant response from first run (1 message)
assert len(thread_messages) >= 2, "Thread should preserve initial messages from before checkpoint"
# Verify initial messages are preserved
assert thread_messages[0].text == "Initial message 1"
assert thread_messages[1].text == "Initial response 1"
# Verify the restored executor's session state was restored
restored_session_obj = restored_executor._session # type: ignore[reportPrivateUsage]
assert restored_session_obj is not None
assert restored_session_obj.session_id == initial_session.session_id
async def test_agent_executor_save_and_restore_state_directly() -> None:
"""Test AgentExecutor's on_checkpoint_save and on_checkpoint_restore methods directly."""
# Create agent with thread containing messages
# Create agent with session containing state
agent = _CountingAgent(id="direct_test_agent", name="DirectTestAgent")
thread = AgentThread(message_store=ChatMessageStore())
session = AgentSession()
# Add messages to thread
thread_messages = [
Message(role="user", text="Message in thread 1"),
Message(role="assistant", text="Thread response 1"),
Message(role="user", text="Message in thread 2"),
# Add messages to session state
session_messages = [
Message(role="user", text="Message in session 1"),
Message(role="assistant", text="Session response 1"),
Message(role="user", text="Message in session 2"),
]
await thread.on_new_messages(thread_messages)
session.state["history"] = {"messages": session_messages}
executor = AgentExecutor(agent, agent_thread=thread)
executor = AgentExecutor(agent, session=session)
# Add messages to executor cache
cache_messages = [
@@ -184,26 +170,23 @@ async def test_agent_executor_save_and_restore_state_directly() -> None:
# Snapshot the state
state = await executor.on_checkpoint_save()
# Verify snapshot contains both cache and thread
# Verify snapshot contains both cache and session
assert "cache" in state
assert "agent_thread" in state
assert "agent_session" in state
# Verify thread state structure
thread_state = state["agent_thread"] # type: ignore[index]
assert "chat_message_store_state" in thread_state
assert "messages" in thread_state["chat_message_store_state"]
# Verify session state structure
session_state = state["agent_session"] # type: ignore[index]
assert "session_id" in session_state
assert "state" in session_state
# Create new executor to restore into
new_agent = _CountingAgent(id="direct_test_agent", name="DirectTestAgent")
new_thread = AgentThread(message_store=ChatMessageStore())
new_executor = AgentExecutor(new_agent, agent_thread=new_thread)
new_session = AgentSession()
new_executor = AgentExecutor(new_agent, session=new_session)
# Verify new executor starts empty
assert len(new_executor._cache) == 0 # type: ignore[reportPrivateUsage]
initial_message_store = new_thread.message_store
assert initial_message_store is not None
initial_thread_msgs = await initial_message_store.list_messages()
assert len(initial_thread_msgs) == 0
assert len(new_session.state) == 0
# Restore state
await new_executor.on_checkpoint_restore(state)
@@ -214,11 +197,6 @@ async def test_agent_executor_save_and_restore_state_directly() -> None:
assert restored_cache[0].text == "Cached user message"
assert restored_cache[1].text == "Cached assistant response"
# Verify thread messages are restored
restored_message_store = new_executor._agent_thread.message_store # type: ignore[reportPrivateUsage]
assert restored_message_store is not None
restored_thread_msgs = await restored_message_store.list_messages()
assert len(restored_thread_msgs) == len(thread_messages)
assert restored_thread_msgs[0].text == "Message in thread 1"
assert restored_thread_msgs[1].text == "Thread response 1"
assert restored_thread_msgs[2].text == "Message in thread 2"
# Verify session was restored with correct session_id
restored_session = new_executor._session # type: ignore[reportPrivateUsage]
assert restored_session.session_id == session.session_id
python/packages/core/tests/workflow/test_agent_executor_tool_calls.py
+2 -2
View File
@@ -13,7 +13,7 @@ from agent_framework import (
AgentExecutorResponse,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
ChatResponse,
ChatResponseUpdate,
@@ -42,7 +42,7 @@ class _ToolCallingAgent(BaseAgent):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
if stream:
python/packages/core/tests/workflow/test_agent_utils.py
+4 -4
View File
@@ -3,7 +3,7 @@
from collections.abc import AsyncIterable
from typing import Any
from agent_framework import AgentResponse, AgentResponseUpdate, AgentThread, Message
from agent_framework import AgentResponse, AgentResponseUpdate, AgentSession, Message
from agent_framework._workflows._agent_utils import resolve_agent_id
@@ -37,12 +37,12 @@ class MockAgent:
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> AgentResponse | AsyncIterable[AgentResponseUpdate]: ...
def get_new_thread(self, **kwargs: Any) -> AgentThread:
"""Creates a new conversation thread for the agent."""
def create_session(self, **kwargs: Any) -> AgentSession:
"""Creates a new conversation session for the agent."""
...
python/packages/core/tests/workflow/test_full_conversation.py
+3 -3
View File
@@ -12,7 +12,7 @@ from agent_framework import (
AgentExecutorResponse,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
Content,
Executor,
@@ -38,7 +38,7 @@ class _SimpleAgent(BaseAgent):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
if stream:
@@ -108,7 +108,7 @@ class _CaptureAgent(BaseAgent):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
# Normalize and record messages for verification
python/packages/core/tests/workflow/test_workflow.py
+2 -2
View File
@@ -13,7 +13,7 @@ from agent_framework import (
AgentExecutor,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
Content,
Executor,
@@ -838,7 +838,7 @@ class _StreamingTestAgent(BaseAgent):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
if stream:
python/packages/core/tests/workflow/test_workflow_agent.py
+44 -72
View File
@@ -11,8 +11,7 @@ from agent_framework import (
AgentExecutorRequest,
AgentResponse,
AgentResponseUpdate,
AgentThread,
ChatMessageStore,
AgentSession,
Content,
Executor,
Message,
@@ -511,80 +510,53 @@ class TestWorkflowAgent:
texts = [message.text for message in result.messages]
assert texts == ["first message", "second message", "third fourth"]
async def test_thread_conversation_history_included_in_workflow_run(self) -> None:
"""Test that conversation history from thread is included when running WorkflowAgent.
This verifies that when a thread with existing messages is provided to agent.run(),
the workflow receives the complete conversation history (thread history + new messages).
"""
async def test_session_conversation_history_included_in_workflow_run(self) -> None:
"""Test that messages provided to agent.run() are passed through to the workflow."""
# Create an executor that captures all received messages
capturing_executor = ConversationHistoryCapturingExecutor(id="capturing", streaming=False)
workflow = WorkflowBuilder(start_executor=capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Thread History Test Agent")
agent = WorkflowAgent(workflow=workflow, name="Session History Test Agent")
# Create a thread with existing conversation history
history_messages = [
Message(role="user", text="Previous user message"),
Message(role="assistant", text="Previous assistant response"),
]
message_store = ChatMessageStore(messages=history_messages)
thread = AgentThread(message_store=message_store)
# Create a session
session = AgentSession()
# Run the agent with the thread and a new message
# Run the agent with the session and a new message
new_message = "New user question"
await agent.run(new_message, thread=thread)
await agent.run(new_message, session=session)
# Verify the executor received both history AND new message
assert len(capturing_executor.received_messages) == 3
# Verify the executor received the message
assert len(capturing_executor.received_messages) == 1
assert capturing_executor.received_messages[0].text == "New user question"
# Verify the order: history first, then new message
assert capturing_executor.received_messages[0].text == "Previous user message"
assert capturing_executor.received_messages[1].text == "Previous assistant response"
assert capturing_executor.received_messages[2].text == "New user question"
async def test_thread_conversation_history_included_in_workflow_stream(self) -> None:
"""Test that conversation history from thread is included when streaming WorkflowAgent.
This verifies that stream=True also includes thread history.
"""
async def test_session_conversation_history_included_in_workflow_stream(self) -> None:
"""Test that messages provided to agent.run() are passed through when streaming WorkflowAgent."""
# Create an executor that captures all received messages
capturing_executor = ConversationHistoryCapturingExecutor(id="capturing_stream")
workflow = WorkflowBuilder(start_executor=capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Thread Stream Test Agent")
agent = WorkflowAgent(workflow=workflow, name="Session Stream Test Agent")
# Create a thread with existing conversation history
history_messages = [
Message(role="system", text="You are a helpful assistant"),
Message(role="user", text="Hello"),
Message("assistant", ["Hi there!"]),
]
message_store = ChatMessageStore(messages=history_messages)
thread = AgentThread(message_store=message_store)
# Create a session
session = AgentSession()
# Stream from the agent with the thread and a new message
async for _ in agent.run("How are you?", stream=True, thread=thread):
# Stream from the agent with the session and a new message
async for _ in agent.run("How are you?", stream=True, session=session):
pass
# Verify the executor received all messages (3 from history + 1 new)
assert len(capturing_executor.received_messages) == 4
# Verify the executor received the message
assert len(capturing_executor.received_messages) == 1
assert capturing_executor.received_messages[0].text == "How are you?"
# Verify the order
assert capturing_executor.received_messages[0].text == "You are a helpful assistant"
assert capturing_executor.received_messages[1].text == "Hello"
assert capturing_executor.received_messages[2].text == "Hi there!"
assert capturing_executor.received_messages[3].text == "How are you?"
async def test_empty_thread_works_correctly(self) -> None:
"""Test that an empty thread (no message store) works correctly."""
capturing_executor = ConversationHistoryCapturingExecutor(id="empty_thread_test")
async def test_empty_session_works_correctly(self) -> None:
"""Test that an empty session (no message store) works correctly."""
capturing_executor = ConversationHistoryCapturingExecutor(id="empty_session_test")
workflow = WorkflowBuilder(start_executor=capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Empty Thread Test Agent")
agent = WorkflowAgent(workflow=workflow, name="Empty Session Test Agent")
# Create an empty thread
thread = AgentThread()
# Create an empty session
session = AgentSession()
# Run with the empty thread
await agent.run("Just a new message", thread=thread)
# Run with the empty session
await agent.run("Just a new message", session=session)
# Should only receive the new message
assert len(capturing_executor.received_messages) == 1
@@ -622,27 +594,27 @@ class TestWorkflowAgent:
self.description: str | None = None
self._response_text = response_text
def get_new_thread(self, **kwargs: Any) -> AgentThread:
return AgentThread()
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
if stream:
return self._run_stream(messages=messages, thread=thread, **kwargs)
return self._run(messages=messages, thread=thread, **kwargs)
return self._run_stream(messages=messages, session=session, **kwargs)
return self._run(messages=messages, session=session, **kwargs)
async def _run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> AgentResponse:
@@ -654,7 +626,7 @@ class TestWorkflowAgent:
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
async def _iter():
@@ -710,27 +682,27 @@ class TestWorkflowAgent:
self.description: str | None = None
self._response_text = response_text
def get_new_thread(self, **kwargs: Any) -> AgentThread:
return AgentThread()
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
if stream:
return self._run_stream(messages=messages, thread=thread, **kwargs)
return self._run(messages=messages, thread=thread, **kwargs)
return self._run_stream(messages=messages, session=session, **kwargs)
return self._run(messages=messages, session=session, **kwargs)
async def _run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> AgentResponse:
@@ -742,7 +714,7 @@ class TestWorkflowAgent:
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
async def _iter():
@@ -1037,7 +1009,7 @@ class TestWorkflowAgentMergeUpdates:
def test_merge_updates_function_result_ordering_github_2977(self):
"""Test that FunctionResultContent updates are placed after their FunctionCallContent.
This test reproduces GitHub issue #2977: When using a thread with WorkflowAgent,
This test reproduces GitHub issue #2977: When using a session with WorkflowAgent,
FunctionResultContent updates without response_id were being added to global_dangling
and placed at the end of messages. This caused OpenAI to reject the conversation because
"An assistant message with 'tool_calls' must be followed by tool messages responding
python/packages/core/tests/workflow/test_workflow_builder.py
+2 -2
View File
@@ -9,7 +9,7 @@ from agent_framework import (
AgentExecutor,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
Executor,
Message,
@@ -21,7 +21,7 @@ from agent_framework import (
class DummyAgent(BaseAgent):
def run(self, messages=None, *, stream: bool = False, thread: AgentThread | None = None, **kwargs): # type: ignore[override]
def run(self, messages=None, *, stream: bool = False, session: AgentSession | None = None, **kwargs): # type: ignore[override]
if stream:
return self._run_stream_impl()
return self._run_impl(messages)
python/packages/core/tests/workflow/test_workflow_kwargs.py
+3 -3
View File
@@ -8,7 +8,7 @@ import pytest
from agent_framework import (
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
Content,
Message,
@@ -55,7 +55,7 @@ class _KwargsCapturingAgent(BaseAgent):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
self.captured_kwargs.append(dict(kwargs))
@@ -88,7 +88,7 @@ class _OptionsAwareAgent(BaseAgent):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: dict[str, Any] | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
python/packages/devui/agent_framework_devui/_conversations.py
+128 -131
View File
@@ -3,7 +3,7 @@
"""Conversation storage abstraction for OpenAI Conversations API.
This module provides a clean abstraction layer for managing conversations
while wrapping AgentFramework's AgentThread underneath.
with in-memory message storage.
"""
from __future__ import annotations
@@ -13,7 +13,7 @@ import uuid
from abc import ABC, abstractmethod
from typing import Any, Literal, cast
from agent_framework import AgentThread, Message
from agent_framework import AgentSession, Message
from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
from openai.types.conversations import Conversation, ConversationDeletedResource
from openai.types.conversations.conversation_item import ConversationItem
@@ -38,14 +38,14 @@ class ConversationStore(ABC):
"""Abstract base class for conversation storage.
Provides OpenAI Conversations API interface while managing
AgentThread instances underneath.
message storage internally.
"""
@abstractmethod
def create_conversation(
self, metadata: dict[str, str] | None = None, conversation_id: str | None = None
) -> Conversation:
"""Create a new conversation (wraps AgentThread creation).
"""Create a new conversation.
Args:
metadata: Optional metadata dict (e.g., {"agent_id": "weather_agent"})
@@ -86,7 +86,7 @@ class ConversationStore(ABC):
@abstractmethod
def delete_conversation(self, conversation_id: str) -> ConversationDeletedResource:
"""Delete conversation (including AgentThread).
"""Delete conversation.
Args:
conversation_id: Conversation ID
@@ -101,7 +101,7 @@ class ConversationStore(ABC):
@abstractmethod
async def add_items(self, conversation_id: str, items: list[dict[str, Any]]) -> list[ConversationItem]:
"""Add items to conversation (syncs to AgentThread.message_store).
"""Add items to conversation.
Args:
conversation_id: Conversation ID
@@ -119,7 +119,7 @@ class ConversationStore(ABC):
async def list_items(
self, conversation_id: str, limit: int = 100, after: str | None = None, order: str = "asc"
) -> tuple[list[ConversationItem], bool]:
"""List conversation items from AgentThread.message_store.
"""List conversation items.
Args:
conversation_id: Conversation ID
@@ -152,17 +152,17 @@ class ConversationStore(ABC):
pass
@abstractmethod
def get_thread(self, conversation_id: str) -> AgentThread | None:
"""Get underlying AgentThread for execution (internal use).
def get_session(self, conversation_id: str) -> AgentSession | None:
"""Get AgentSession for agent execution.
This is the critical method that allows the executor to get the
AgentThread for running agents with conversation context.
AgentSession for running agents with conversation context.
Args:
conversation_id: Conversation ID
Returns:
AgentThread object or None if not found
AgentSession object or None if not found
"""
pass
@@ -183,7 +183,7 @@ class ConversationStore(ABC):
"""Add a trace event to the conversation for context inspection.
Traces capture execution metadata like token usage, timing, and LLM context
that isn't stored in the AgentThread but is useful for debugging.
that is useful for debugging.
Args:
conversation_id: Conversation ID
@@ -205,17 +205,17 @@ class ConversationStore(ABC):
class InMemoryConversationStore(ConversationStore):
"""In-memory conversation storage wrapping AgentThread.
"""In-memory conversation storage.
This implementation stores conversations in memory with their
underlying AgentThread instances for execution.
underlying message lists and AgentSession instances for execution.
"""
def __init__(self) -> None:
"""Initialize in-memory conversation storage.
Storage structure maps conversation IDs to conversation data including
the underlying AgentThread, metadata, and cached ConversationItems.
messages, metadata, and cached ConversationItems.
"""
self._conversations: dict[str, dict[str, Any]] = {}
@@ -225,20 +225,22 @@ class InMemoryConversationStore(ConversationStore):
def create_conversation(
self, metadata: dict[str, str] | None = None, conversation_id: str | None = None
) -> Conversation:
"""Create a new conversation with underlying AgentThread and checkpoint storage."""
"""Create a new conversation with message storage and checkpoint storage."""
conv_id = conversation_id or f"conv_{uuid.uuid4().hex}"
created_at = int(time.time())
# Create AgentThread with default ChatMessageStore
thread = AgentThread()
# Create message list for internal storage and AgentSession for execution
messages: list[Message] = []
session = AgentSession(session_id=conv_id)
# Create session-scoped checkpoint storage (one per conversation)
checkpoint_storage = InMemoryCheckpointStorage()
self._conversations[conv_id] = {
"id": conv_id,
"thread": thread,
"checkpoint_storage": checkpoint_storage, # Stored alongside thread
"messages": messages,
"session": session,
"checkpoint_storage": checkpoint_storage,
"metadata": metadata or {},
"created_at": created_at,
"items": [],
@@ -279,7 +281,7 @@ class InMemoryConversationStore(ConversationStore):
)
def delete_conversation(self, conversation_id: str) -> ConversationDeletedResource:
"""Delete conversation and its AgentThread."""
"""Delete conversation."""
if conversation_id not in self._conversations:
raise ValueError(f"Conversation {conversation_id} not found")
@@ -290,14 +292,14 @@ class InMemoryConversationStore(ConversationStore):
return ConversationDeletedResource(id=conversation_id, object="conversation.deleted", deleted=True)
async def add_items(self, conversation_id: str, items: list[dict[str, Any]]) -> list[ConversationItem]:
"""Add items to conversation and sync to AgentThread."""
"""Add items to conversation."""
conv_data = self._conversations.get(conversation_id)
if not conv_data:
raise ValueError(f"Conversation {conversation_id} not found")
thread: AgentThread = conv_data["thread"]
stored_messages: list[Message] = conv_data["messages"]
# Convert items to ChatMessages and add to thread
# Convert items to Messages and add to storage
chat_messages = []
for item in items:
# Simple conversion - assume text content for now
@@ -308,8 +310,8 @@ class InMemoryConversationStore(ConversationStore):
chat_msg = Message(role=role, text=text) # type: ignore[arg-type]
chat_messages.append(chat_msg)
# Add messages to AgentThread
await thread.on_new_messages(chat_messages)
# Add messages to internal storage
stored_messages.extend(chat_messages)
# Create Message objects (ConversationItem is a Union - use concrete Message type)
conv_items: list[ConversationItem] = []
@@ -354,9 +356,9 @@ class InMemoryConversationStore(ConversationStore):
async def list_items(
self, conversation_id: str, limit: int = 100, after: str | None = None, order: str = "asc"
) -> tuple[list[ConversationItem], bool]:
"""List conversation items from AgentThread message store.
"""List conversation items.
Converts AgentFramework ChatMessages to proper OpenAI ConversationItem types:
Converts stored Messages to proper OpenAI ConversationItem types:
- Messages with text/images/files → Message
- Function calls → ResponseFunctionToolCallItem
- Function results → ResponseFunctionToolCallOutputItem
@@ -365,119 +367,114 @@ class InMemoryConversationStore(ConversationStore):
if not conv_data:
raise ValueError(f"Conversation {conversation_id} not found")
thread: AgentThread = conv_data["thread"]
stored_messages: list[Message] = conv_data["messages"]
# Get messages from thread's message store
# Convert stored messages to ConversationItem types
items: list[ConversationItem] = []
if thread.message_store:
af_messages = await thread.message_store.list_messages()
af_messages = stored_messages
# Convert each AgentFramework Message to appropriate ConversationItem type(s)
for i, msg in enumerate(af_messages):
item_id = f"item_{i}"
role_str = msg.role if hasattr(msg.role, "value") else str(msg.role)
role = cast(MessageRole, role_str) # Safe: Agent Framework roles match OpenAI roles
# Convert each AgentFramework Message to appropriate ConversationItem type(s)
for i, msg in enumerate(af_messages):
item_id = f"item_{i}"
role_str = msg.role if hasattr(msg.role, "value") else str(msg.role)
role = cast(MessageRole, role_str) # Safe: Agent Framework roles match OpenAI roles
# Process each content item in the message
# A single Message may produce multiple ConversationItems
# (e.g., a message with both text and a function call)
message_contents: list[TextContent | ResponseInputImage | ResponseInputFile] = []
function_calls = []
function_results = []
# Process each content item in the message
# A single Message may produce multiple ConversationItems
# (e.g., a message with both text and a function call)
message_contents: list[TextContent | ResponseInputImage | ResponseInputFile] = []
function_calls = []
function_results = []
for content in msg.contents:
content_type = getattr(content, "type", None)
for content in msg.contents:
content_type = getattr(content, "type", None)
if content_type == "text":
# Text content for Message
text_value = getattr(content, "text", "")
message_contents.append(TextContent(type="text", text=text_value))
if content_type == "text":
# Text content for Message
text_value = getattr(content, "text", "")
message_contents.append(TextContent(type="text", text=text_value))
elif content_type == "data":
# Data content (images, files, PDFs)
uri = getattr(content, "uri", "")
media_type = getattr(content, "media_type", None)
elif content_type == "data":
# Data content (images, files, PDFs)
uri = getattr(content, "uri", "")
media_type = getattr(content, "media_type", None)
if media_type and media_type.startswith("image/"):
# Convert to ResponseInputImage
message_contents.append(
ResponseInputImage(type="input_image", image_url=uri, detail="auto")
if media_type and media_type.startswith("image/"):
# Convert to ResponseInputImage
message_contents.append(ResponseInputImage(type="input_image", image_url=uri, detail="auto"))
else:
# Convert to ResponseInputFile
# Extract filename from URI if possible
filename = None
if media_type == "application/pdf":
filename = "document.pdf"
message_contents.append(ResponseInputFile(type="input_file", file_url=uri, filename=filename))
elif content_type == "function_call":
# Function call - create separate ConversationItem
call_id = getattr(content, "call_id", None)
name = getattr(content, "name", "")
arguments = getattr(content, "arguments", "")
if call_id and name:
function_calls.append(
ResponseFunctionToolCallItem(
id=f"{item_id}_call_{call_id}",
call_id=call_id,
name=name,
arguments=arguments,
type="function_call",
status="completed",
)
else:
# Convert to ResponseInputFile
# Extract filename from URI if possible
filename = None
if media_type == "application/pdf":
filename = "document.pdf"
)
message_contents.append(
ResponseInputFile(type="input_file", file_url=uri, filename=filename)
elif content_type == "function_result":
# Function result - create separate ConversationItem
call_id = getattr(content, "call_id", None)
# Output is stored in the 'result' field of FunctionResultContent
result_value = getattr(content, "result", None)
# Convert result to string (it could be dict, list, or other types)
if result_value is None:
output = ""
elif isinstance(result_value, str):
output = result_value
else:
import json
try:
output = json.dumps(result_value)
except (TypeError, ValueError):
output = str(result_value)
if call_id:
function_results.append(
ResponseFunctionToolCallOutputItem(
id=f"{item_id}_result_{call_id}",
call_id=call_id,
output=output,
type="function_call_output",
status="completed",
)
)
elif content_type == "function_call":
# Function call - create separate ConversationItem
call_id = getattr(content, "call_id", None)
name = getattr(content, "name", "")
arguments = getattr(content, "arguments", "")
# Create ConversationItems based on what we found
# If message has text/images/files, create a Message item
if message_contents:
message = OpenAIMessage(
id=item_id,
type="message",
role=role, # type: ignore
content=message_contents, # type: ignore
status="completed",
)
items.append(message)
if call_id and name:
function_calls.append(
ResponseFunctionToolCallItem(
id=f"{item_id}_call_{call_id}",
call_id=call_id,
name=name,
arguments=arguments,
type="function_call",
status="completed",
)
)
# Add function call items
items.extend(function_calls)
elif content_type == "function_result":
# Function result - create separate ConversationItem
call_id = getattr(content, "call_id", None)
# Output is stored in the 'result' field of FunctionResultContent
result_value = getattr(content, "result", None)
# Convert result to string (it could be dict, list, or other types)
if result_value is None:
output = ""
elif isinstance(result_value, str):
output = result_value
else:
import json
try:
output = json.dumps(result_value)
except (TypeError, ValueError):
output = str(result_value)
if call_id:
function_results.append(
ResponseFunctionToolCallOutputItem(
id=f"{item_id}_result_{call_id}",
call_id=call_id,
output=output,
type="function_call_output",
status="completed",
)
)
# Create ConversationItems based on what we found
# If message has text/images/files, create a Message item
if message_contents:
message = OpenAIMessage(
id=item_id,
type="message",
role=role, # type: ignore
content=message_contents, # type: ignore
status="completed",
)
items.append(message)
# Add function call items
items.extend(function_calls)
# Add function result items
items.extend(function_results)
# Add function result items
items.extend(function_results)
# Include checkpoints from checkpoint storage as conversation items
checkpoint_storage = conv_data.get("checkpoint_storage")
@@ -589,16 +586,16 @@ class InMemoryConversationStore(ConversationStore):
return None
def get_thread(self, conversation_id: str) -> AgentThread | None:
"""Get AgentThread for execution - CRITICAL for agent.run()."""
def get_session(self, conversation_id: str) -> AgentSession | None:
"""Get AgentSession for execution - CRITICAL for agent.run()."""
conv_data = self._conversations.get(conversation_id)
return conv_data["thread"] if conv_data else None
return conv_data["session"] if conv_data else None
def add_trace(self, conversation_id: str, trace_event: dict[str, Any]) -> None:
"""Add a trace event to the conversation for context inspection.
Traces capture execution metadata like token usage, timing, and LLM context
that isn't stored in the AgentThread but is useful for debugging.
that is useful for debugging.
Args:
conversation_id: Conversation ID
python/packages/devui/agent_framework_devui/_executor.py
+7 -7
View File
@@ -308,15 +308,15 @@ class AgentFrameworkExecutor:
# Convert input to proper Message or string
user_message = self._convert_input_to_chat_message(request.input)
# Get thread from conversation parameter (OpenAI standard!)
thread = None
# Get session from conversation parameter (OpenAI standard!)
session = None
conversation_id = request._get_conversation_id()
if conversation_id:
thread = self.conversation_store.get_thread(conversation_id)
if thread:
session = self.conversation_store.get_session(conversation_id)
if session:
logger.debug(f"Using existing conversation: {conversation_id}")
else:
logger.warning(f"Conversation {conversation_id} not found, proceeding without thread")
logger.warning(f"Conversation {conversation_id} not found, proceeding without session")
if isinstance(user_message, str):
logger.debug(f"Executing agent with text input: {user_message[:100]}...")
@@ -331,8 +331,8 @@ class AgentFrameworkExecutor:
# Agent must have run() method - use stream=True for streaming
if hasattr(agent, "run") and callable(agent.run):
# Use Agent Framework's run() with stream=True for streaming
if thread:
async for update in agent.run(user_message, stream=True, thread=thread):
if session:
async for update in agent.run(user_message, stream=True, session=session):
for trace_event in trace_collector.get_pending_events():
yield trace_event
python/packages/devui/tests/devui/conftest.py
+11 -11
View File
@@ -20,7 +20,7 @@ from agent_framework import (
Agent,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
BaseChatClient,
ChatResponse,
@@ -162,19 +162,19 @@ class MockAgent(BaseAgent):
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
self.call_count += 1
if stream:
return self._run_stream(messages=messages, thread=thread, **kwargs)
return self._run(messages=messages, thread=thread, **kwargs)
return self._run_stream(messages=messages, session=session, **kwargs)
return self._run(messages=messages, session=session, **kwargs)
async def _run(
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> AgentResponse:
self.call_count += 1
@@ -184,7 +184,7 @@ class MockAgent(BaseAgent):
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
self.call_count += 1
@@ -208,19 +208,19 @@ class MockToolCallingAgent(BaseAgent):
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
self.call_count += 1
if stream:
return self._run_stream(messages=messages, thread=thread, **kwargs)
return self._run(messages=messages, thread=thread, **kwargs)
return self._run_stream(messages=messages, session=session, **kwargs)
return self._run(messages=messages, session=session, **kwargs)
async def _run(
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> AgentResponse:
return AgentResponse(messages=[Message("assistant", ["done"])])
@@ -229,7 +229,7 @@ class MockToolCallingAgent(BaseAgent):
self,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
async def _iter() -> AsyncIterable[AgentResponseUpdate]:
python/packages/devui/tests/devui/test_conversations.py
+17 -31
View File
@@ -83,29 +83,29 @@ async def test_delete_conversation():
@pytest.mark.asyncio
async def test_get_thread():
"""Test getting underlying AgentThread."""
async def test_get_session():
"""Test getting AgentSession for execution."""
store = InMemoryConversationStore()
# Create conversation
conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
# Get thread
thread = store.get_thread(conversation.id)
# Get session
session = store.get_session(conversation.id)
assert thread is not None
# AgentThread should have message_store
assert hasattr(thread, "message_store")
assert session is not None
# AgentSession should have session_id
assert hasattr(session, "session_id")
@pytest.mark.asyncio
async def test_get_thread_not_found():
"""Test getting thread for non-existent conversation."""
async def test_get_session_not_found():
"""Test getting session for non-existent conversation."""
store = InMemoryConversationStore()
thread = store.get_thread("conv_nonexistent")
session = store.get_session("conv_nonexistent")
assert thread is None
assert session is None
@pytest.mark.asyncio
@@ -199,21 +199,13 @@ async def test_list_items_pagination():
@pytest.mark.asyncio
async def test_list_items_converts_function_calls():
"""Test that list_items properly converts function calls to ResponseFunctionToolCallItem."""
from agent_framework import ChatMessageStore, Message
from agent_framework import Message
store = InMemoryConversationStore()
# Create conversation
conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
# Get the underlying thread and set up message store
thread = store.get_thread(conversation.id)
assert thread is not None
# Initialize message store if not present
if thread.message_store is None:
thread.message_store = ChatMessageStore()
# Simulate messages from agent execution with function calls
messages = [
Message(role="user", contents=[{"type": "text", "text": "What's the weather in SF?"}]),
@@ -241,8 +233,8 @@ async def test_list_items_converts_function_calls():
Message(role="assistant", contents=[{"type": "text", "text": "The weather is sunny, 65°F"}]),
]
# Add messages to thread
await thread.on_new_messages(messages)
# Add messages to internal storage
store._conversations[conversation.id]["messages"].extend(messages)
# List conversation items
items, has_more = await store.list_items(conversation.id)
@@ -284,20 +276,13 @@ async def test_list_items_converts_function_calls():
@pytest.mark.asyncio
async def test_list_items_handles_images_and_files():
"""Test that list_items properly converts data content (images/files) to OpenAI types."""
from agent_framework import ChatMessageStore, Message
from agent_framework import Message
store = InMemoryConversationStore()
# Create conversation
conversation = store.create_conversation(metadata={"agent_id": "test_agent"})
# Get the underlying thread
thread = store.get_thread(conversation.id)
assert thread is not None
if thread.message_store is None:
thread.message_store = ChatMessageStore()
# Simulate message with image and file
messages = [
Message(
@@ -310,7 +295,8 @@ async def test_list_items_handles_images_and_files():
),
]
await thread.on_new_messages(messages)
# Add messages to internal storage
store._conversations[conversation.id]["messages"].extend(messages)
# List items
items, has_more = await store.list_items(conversation.id)
python/packages/devui/tests/devui/test_discovery.py
+9 -9
View File
@@ -74,14 +74,14 @@ async def test_discovery_accepts_agents_with_only_run():
init_file = agent_dir / "__init__.py"
init_file.write_text("""
from agent_framework import AgentResponse, AgentThread, Message, Role, Content
from agent_framework import AgentResponse, AgentSession, Message, Role, Content
class NonStreamingAgent:
id = "non_streaming"
name = "Non-Streaming Agent"
description = "Agent with run() method"
async def run(self, messages=None, *, thread=None, **kwargs):
async def run(self, messages=None, *, session=None, **kwargs):
return AgentResponse(
messages=[Message(
role="assistant",
@@ -90,8 +90,8 @@ class NonStreamingAgent:
response_id="test"
)
def get_new_thread(self, **kwargs):
return AgentThread()
def create_session(self, **kwargs):
return AgentSession()
agent = NonStreamingAgent()
""")
@@ -188,19 +188,19 @@ workflow = WorkflowBuilder(start_executor=executor).build()
agent_dir = temp_path / "my_agent"
agent_dir.mkdir()
(agent_dir / "agent.py").write_text("""
from agent_framework import AgentResponse, AgentThread, Message, Role, TextContent
from agent_framework import AgentResponse, AgentSession, Message, Role, TextContent
class TestAgent:
name = "Test Agent"
async def run(self, messages=None, *, thread=None, **kwargs):
async def run(self, messages=None, *, session=None, **kwargs):
return AgentResponse(
messages=[Message(role="assistant", contents=[Content.from_text(text="test")])],
response_id="test"
)
def get_new_thread(self, **kwargs):
return AgentThread()
def create_session(self, **kwargs):
return AgentSession()
agent = TestAgent()
""")
@@ -320,7 +320,7 @@ class WeatherAgent:
name = "Weather Agent"
description = "Gets weather information"
def run(self, input_str, *, stream: bool = False, thread=None, **kwargs):
def run(self, input_str, *, stream: bool = False, session=None, **kwargs):
return f"Weather in {input_str}"
""")
python/packages/devui/tests/devui/test_execution.py
+5 -5
View File
@@ -538,7 +538,7 @@ def test_extract_workflow_hil_responses_handles_stringified_json():
async def test_executor_handles_streaming_agent():
"""Test executor handles agents with run(stream=True) method."""
from agent_framework import AgentResponse, AgentResponseUpdate, AgentThread, Content, Message
from agent_framework import AgentResponse, AgentResponseUpdate, AgentSession, Content, Message
class StreamingAgent:
"""Agent with run() method supporting stream parameter."""
@@ -547,7 +547,7 @@ async def test_executor_handles_streaming_agent():
name = "Streaming Test Agent"
description = "Test agent with run(stream=True)"
def run(self, messages=None, *, stream=False, thread=None, **kwargs):
def run(self, messages=None, *, stream=False, session=None, **kwargs):
if stream:
# Return an async generator for streaming
return self._stream_impl(messages)
@@ -566,8 +566,8 @@ async def test_executor_handles_streaming_agent():
role="assistant",
)
def get_new_thread(self, **kwargs):
return AgentThread()
def create_session(self, **kwargs):
return AgentSession()
# Create executor and register agent
discovery = EntityDiscovery(None)
@@ -754,7 +754,7 @@ class StreamingAgent:
name = "Streaming Test Agent"
description = "Test agent for streaming"
async def run(self, input_str, *, stream: bool = False, thread=None, **kwargs):
async def run(self, input_str, *, stream: bool = False, session=None, **kwargs):
if stream:
async def _stream():
for i, word in enumerate(f"Processing {input_str}".split()):
python/packages/durabletask/AGENTS.md
+1 -1
View File
@@ -18,7 +18,7 @@ Durable execution support for long-running agent workflows using Azure Durable F
### State Management
- **`DurableAgentState`** - State container for durable agents
- **`DurableAgentThread`** - Thread management for durable agents
- **`DurableAgentSession`** - Session management for durable agents
- **`DurableAIAgentOrchestrationContext`** - Orchestration context
### Callbacks
python/packages/durabletask/agent_framework_durabletask/__init__.py
+2 -2
View File
@@ -45,7 +45,7 @@ from ._durable_agent_state import (
)
from ._entities import AgentEntity, AgentEntityStateProviderMixin
from ._executors import DurableAgentExecutor
from ._models import AgentSessionId, DurableAgentThread, RunRequest
from ._models import AgentSessionId, DurableAgentSession, RunRequest
from ._orchestration_context import DurableAIAgentOrchestrationContext
from ._response_utils import ensure_response_format, load_agent_response
from ._shim import DurableAIAgent
@@ -79,6 +79,7 @@ __all__ = [
"DurableAIAgentOrchestrationContext",
"DurableAIAgentWorker",
"DurableAgentExecutor",
"DurableAgentSession",
"DurableAgentState",
"DurableAgentStateContent",
"DurableAgentStateData",
@@ -99,7 +100,6 @@ __all__ = [
"DurableAgentStateUriContent",
"DurableAgentStateUsage",
"DurableAgentStateUsageContent",
"DurableAgentThread",
"DurableStateFields",
"RunRequest",
"__version__",
python/packages/durabletask/agent_framework_durabletask/_executors.py
+19 -19
View File
@@ -16,7 +16,7 @@ from abc import ABC, abstractmethod
from datetime import datetime, timezone
from typing import Any, Generic, TypeVar
from agent_framework import AgentResponse, AgentThread, Content, Message, get_logger
from agent_framework import AgentResponse, AgentSession, Content, Message, get_logger
from durabletask.client import TaskHubGrpcClient
from durabletask.entities import EntityInstanceId
from durabletask.task import CompletableTask, CompositeTask, OrchestrationContext, Task
@@ -24,7 +24,7 @@ from pydantic import BaseModel
from ._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
from ._durable_agent_state import DurableAgentState
from ._models import AgentSessionId, DurableAgentThread, RunRequest
from ._models import AgentSessionId, DurableAgentSession, RunRequest
from ._response_utils import ensure_response_format, load_agent_response
logger = get_logger("agent_framework.durabletask.executors")
@@ -114,7 +114,7 @@ class DurableAgentExecutor(ABC, Generic[TaskT]):
self,
agent_name: str,
run_request: RunRequest,
thread: AgentThread | None = None,
session: AgentSession | None = None,
) -> TaskT:
"""Execute the durable agent.
@@ -123,20 +123,20 @@ class DurableAgentExecutor(ABC, Generic[TaskT]):
"""
raise NotImplementedError
def get_new_thread(self, agent_name: str, **kwargs: Any) -> DurableAgentThread:
"""Create a new DurableAgentThread with random session ID."""
def get_new_session(self, agent_name: str, **kwargs: Any) -> DurableAgentSession:
"""Create a new DurableAgentSession with random session ID."""
session_id = self._create_session_id(agent_name)
return DurableAgentThread.from_session_id(session_id, **kwargs)
return DurableAgentSession.from_session_id(session_id, **kwargs)
def _create_session_id(
self,
agent_name: str,
thread: AgentThread | None = None,
session: AgentSession | None = None,
) -> AgentSessionId:
"""Create the AgentSessionId for the execution."""
if isinstance(thread, DurableAgentThread) and thread.session_id is not None:
return thread.session_id
# Create new session ID - either no thread provided or it's a regular AgentThread
if isinstance(session, DurableAgentSession) and session.durable_session_id is not None:
return session.durable_session_id
# Create new session ID - either no session provided or it's a regular AgentSession
key = self.generate_unique_id()
return AgentSessionId(name=agent_name, key=key)
@@ -217,7 +217,7 @@ class ClientAgentExecutor(DurableAgentExecutor[AgentResponse]):
self,
agent_name: str,
run_request: RunRequest,
thread: AgentThread | None = None,
session: AgentSession | None = None,
) -> AgentResponse:
"""Execute the agent via the durabletask client.
@@ -231,14 +231,14 @@ class ClientAgentExecutor(DurableAgentExecutor[AgentResponse]):
Args:
agent_name: Name of the agent to execute
run_request: The run request containing message and optional response format
thread: Optional conversation thread (creates new if not provided)
session: Optional conversation session (creates new if not provided)
Returns:
AgentResponse: The agent's response after execution completes, or an immediate
acknowledgement if wait_for_response is False
"""
# Signal the entity with the request
entity_id = self._signal_agent_entity(agent_name, run_request, thread)
entity_id = self._signal_agent_entity(agent_name, run_request, session)
# If fire-and-forget mode, return immediately without polling
if not run_request.wait_for_response:
@@ -258,20 +258,20 @@ class ClientAgentExecutor(DurableAgentExecutor[AgentResponse]):
self,
agent_name: str,
run_request: RunRequest,
thread: AgentThread | None,
session: AgentSession | None,
) -> EntityInstanceId:
"""Signal the agent entity with a run request.
Args:
agent_name: Name of the agent to execute
run_request: The run request containing message and optional response format
thread: Optional conversation thread
session: Optional conversation session
Returns:
entity_id
"""
# Get or create session ID
session_id = self._create_session_id(agent_name, thread)
session_id = self._create_session_id(agent_name, session)
# Create the entity ID
entity_id = EntityInstanceId(
@@ -460,7 +460,7 @@ class OrchestrationAgentExecutor(DurableAgentExecutor[DurableAgentTask]):
self,
agent_name: str,
run_request: RunRequest,
thread: AgentThread | None = None,
session: AgentSession | None = None,
) -> DurableAgentTask:
"""Execute the agent via orchestration context.
@@ -470,13 +470,13 @@ class OrchestrationAgentExecutor(DurableAgentExecutor[DurableAgentTask]):
Args:
agent_name: Name of the agent to execute
run_request: The run request containing message and optional response format
thread: Optional conversation thread (creates new if not provided)
session: Optional conversation session (creates new if not provided)
Returns:
DurableAgentTask: A task wrapping the entity call that yields AgentResponse
"""
# Resolve session
session_id = self._create_session_id(agent_name, thread)
session_id = self._create_session_id(agent_name, session)
# Create the entity ID
entity_id = EntityInstanceId(
python/packages/durabletask/agent_framework_durabletask/_models.py
+32 -41
View File
@@ -10,13 +10,12 @@ from __future__ import annotations
import inspect
import json
import uuid
from collections.abc import MutableMapping
from dataclasses import dataclass, field
from datetime import datetime, timezone
from importlib import import_module
from typing import TYPE_CHECKING, Any, cast
from agent_framework import AgentThread
from agent_framework import AgentSession
from ._constants import REQUEST_RESPONSE_FORMAT_TEXT
@@ -274,65 +273,57 @@ class AgentSessionId:
raise ValueError(f"Invalid agent session ID format: {session_id_string}")
class DurableAgentThread(AgentThread):
"""Durable agent thread that tracks the owning :class:`AgentSessionId`."""
class DurableAgentSession(AgentSession):
"""Durable agent session that tracks the owning :class:`AgentSessionId`."""
_SERIALIZED_SESSION_ID_KEY = "durable_session_id"
def __init__(
self,
*,
session_id: AgentSessionId | None = None,
durable_session_id: AgentSessionId | None = None,
session_id: str | None = None,
service_session_id: str | None = None,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
self._session_id: AgentSessionId | None = session_id
super().__init__(session_id=session_id, service_session_id=service_session_id, **kwargs)
self._session_id_value: AgentSessionId | None = durable_session_id
@property
def session_id(self) -> AgentSessionId | None:
return self._session_id
def durable_session_id(self) -> AgentSessionId | None:
return self._session_id_value
@session_id.setter
def session_id(self, value: AgentSessionId | None) -> None:
self._session_id = value
@durable_session_id.setter
def durable_session_id(self, value: AgentSessionId | None) -> None:
self._session_id_value = value
@classmethod
def from_session_id(
cls,
session_id: AgentSessionId,
**kwargs: Any,
) -> DurableAgentThread:
return cls(session_id=session_id, **kwargs)
) -> DurableAgentSession:
return cls(durable_session_id=session_id, **kwargs)
async def serialize(self, **kwargs: Any) -> dict[str, Any]:
state = await super().serialize(**kwargs)
if self._session_id is not None:
state[self._SERIALIZED_SESSION_ID_KEY] = str(self._session_id)
def to_dict(self) -> dict[str, Any]:
state = super().to_dict()
if self._session_id_value is not None:
state[self._SERIALIZED_SESSION_ID_KEY] = str(self._session_id_value)
return state
@classmethod
async def deserialize(
cls,
serialized_thread_state: MutableMapping[str, Any],
*,
message_store: Any = None,
**kwargs: Any,
) -> DurableAgentThread:
state_payload = dict(serialized_thread_state)
def from_dict(cls, data: dict[str, Any]) -> DurableAgentSession:
state_payload = dict(data)
session_id_value = state_payload.pop(cls._SERIALIZED_SESSION_ID_KEY, None)
thread = await super().deserialize(
state_payload,
message_store=message_store,
**kwargs,
session = super().from_dict(state_payload)
# We need to create a DurableAgentSession from the base AgentSession
durable_session = cls(
session_id=session.session_id,
service_session_id=session.service_session_id,
)
if not isinstance(thread, DurableAgentThread):
raise TypeError("Deserialized thread is not a DurableAgentThread instance")
if session_id_value is None:
return thread
if not isinstance(session_id_value, str):
raise ValueError("durable_session_id must be a string when present in serialized state")
thread.session_id = AgentSessionId.parse(session_id_value)
return thread
durable_session.state.update(session.state)
if session_id_value is not None:
if not isinstance(session_id_value, str):
raise ValueError("durable_session_id must be a string when present in serialized state")
durable_session._session_id_value = AgentSessionId.parse(session_id_value)
return durable_session
python/packages/durabletask/agent_framework_durabletask/_shim.py
+15 -8
View File
@@ -12,10 +12,10 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any, Generic, Literal, TypeVar
from agent_framework import AgentThread, Message, SupportsAgentRun
from agent_framework import AgentSession, Message, SupportsAgentRun
from ._executors import DurableAgentExecutor
from ._models import DurableAgentThread
from ._models import DurableAgentSession
# TypeVar for the task type returned by executors
# Covariant because TaskT only appears in return positions (output)
@@ -89,7 +89,7 @@ class DurableAIAgent(SupportsAgentRun, Generic[TaskT]):
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: Literal[False] = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: dict[str, Any] | None = None,
) -> TaskT:
"""Execute the agent via the injected provider.
@@ -98,7 +98,7 @@ class DurableAIAgent(SupportsAgentRun, Generic[TaskT]):
messages: The message(s) to send to the agent
stream: Whether to use streaming for the response (must be False)
DurableAgents do not support streaming mode.
thread: Optional agent thread for conversation context
session: Optional agent session for conversation context
options: Optional options dictionary. Supported keys include
``response_format``, ``enable_tool_calls``, and ``wait_for_response``.
Additional keys are forwarded to the agent execution.
@@ -129,12 +129,19 @@ class DurableAIAgent(SupportsAgentRun, Generic[TaskT]):
return self._executor.run_durable_agent(
agent_name=self.name,
run_request=run_request,
thread=thread,
session=session,
)
def get_new_thread(self, **kwargs: Any) -> DurableAgentThread:
"""Create a new agent thread via the provider."""
return self._executor.get_new_thread(self.name, **kwargs)
def create_session(self, **kwargs: Any) -> DurableAgentSession:
"""Create a new agent session via the provider."""
return self._executor.get_new_session(self.name, **kwargs)
def get_session(self, **kwargs: Any) -> AgentSession:
"""Retrieve an existing session via the provider.
For durable agents, sessions do not use `service_session_id` so this is not used.
"""
return self._executor.get_new_session(self.name, **kwargs)
def _normalize_messages(self, messages: str | Message | list[str] | list[Message] | None) -> str:
"""Convert supported message inputs to a single string.
python/packages/durabletask/tests/integration_tests/test_01_dt_single_agent.py
+14 -14
View File
@@ -39,9 +39,9 @@ class TestSingleAgent:
def test_single_interaction(self):
"""Test a single interaction with the agent."""
agent = self.agent_client.get_agent("Joker")
thread = agent.get_new_thread()
session = agent.create_session()
response = agent.run("Tell me a short joke about programming.", thread=thread)
response = agent.run("Tell me a short joke about programming.", session=session)
assert response is not None
assert response.text is not None
@@ -50,33 +50,33 @@ class TestSingleAgent:
def test_conversation_continuity(self):
"""Test that conversation context is maintained across turns."""
agent = self.agent_client.get_agent("Joker")
thread = agent.get_new_thread()
session = agent.create_session()
# First turn: Ask for a joke about a specific topic
response1 = agent.run("Tell me a joke about cats.", thread=thread)
response1 = agent.run("Tell me a joke about cats.", session=session)
assert response1 is not None
assert len(response1.text) > 0
# Second turn: Ask a follow-up that requires context
response2 = agent.run("Can you make it funnier?", thread=thread)
response2 = agent.run("Can you make it funnier?", session=session)
assert response2 is not None
assert len(response2.text) > 0
# The agent should understand "it" refers to the previous joke
def test_multiple_threads(self):
"""Test that different threads maintain separate contexts."""
def test_multiple_sessions(self):
"""Test that different sessions maintain separate contexts."""
agent = self.agent_client.get_agent("Joker")
# Create two separate threads
thread1 = agent.get_new_thread()
thread2 = agent.get_new_thread()
# Create two separate sessions
session1 = agent.create_session()
session2 = agent.create_session()
assert thread1.session_id != thread2.session_id
assert session1.durable_session_id != session2.durable_session_id
# Send different messages to each thread
response1 = agent.run("Tell me a joke about dogs.", thread=thread1)
response2 = agent.run("Tell me a joke about birds.", thread=thread2)
# Send different messages to each session
response1 = agent.run("Tell me a joke about dogs.", session=session1)
response2 = agent.run("Tell me a joke about birds.", session=session2)
assert response1 is not None
assert response2 is not None
python/packages/durabletask/tests/integration_tests/test_02_dt_multi_agent.py
+7 -7
View File
@@ -47,9 +47,9 @@ class TestMultiAgent:
def test_weather_agent_with_tool(self):
"""Test weather agent with weather tool execution."""
agent = self.agent_client.get_agent(WEATHER_AGENT_NAME)
thread = agent.get_new_thread()
session = agent.create_session()
response = agent.run("What's the weather in Seattle?", thread=thread)
response = agent.run("What's the weather in Seattle?", session=session)
assert response is not None
assert response.text is not None
@@ -66,9 +66,9 @@ class TestMultiAgent:
def test_math_agent_with_tool(self):
"""Test math agent with calculation tool execution."""
agent = self.agent_client.get_agent(MATH_AGENT_NAME)
thread = agent.get_new_thread()
session = agent.create_session()
response = agent.run("Calculate a 20% tip on a $50 bill.", thread=thread)
response = agent.run("Calculate a 20% tip on a $50 bill.", session=session)
assert response is not None
assert response.text is not None
@@ -85,11 +85,11 @@ class TestMultiAgent:
def test_multiple_calls_to_same_agent(self):
"""Test multiple sequential calls to the same agent."""
agent = self.agent_client.get_agent(WEATHER_AGENT_NAME)
thread = agent.get_new_thread()
session = agent.create_session()
# Multiple weather queries
response1 = agent.run("What's the weather in Chicago?", thread=thread)
response2 = agent.run("And what about Los Angeles?", thread=thread)
response1 = agent.run("What's the weather in Chicago?", session=session)
response2 = agent.run("And what about Los Angeles?", session=session)
assert response1 is not None
assert response2 is not None
python/packages/durabletask/tests/integration_tests/test_03_dt_single_agent_streaming.py
+18 -18
View File
@@ -70,7 +70,7 @@ class TestSampleReliableStreaming:
async def _stream_from_redis(
self,
thread_id: str,
session_key: str,
cursor: str | None = None,
timeout: float = 30.0,
) -> tuple[str, bool, str]:
@@ -78,7 +78,7 @@ class TestSampleReliableStreaming:
Stream responses from Redis using the sample's RedisStreamResponseHandler.
Args:
thread_id: The conversation/thread ID to stream from
session_key: The conversation/thread ID to stream from
cursor: Optional cursor to resume from
timeout: Maximum time to wait for stream completion
@@ -92,7 +92,7 @@ class TestSampleReliableStreaming:
async with await self._get_stream_handler() as stream_handler: # type: ignore[reportUnknownMemberType]
try:
async for chunk in stream_handler.read_stream(thread_id, cursor): # type: ignore[reportUnknownMemberType]
async for chunk in stream_handler.read_stream(session_key, cursor): # type: ignore[reportUnknownMemberType]
if time.time() - start_time > timeout:
break
@@ -124,15 +124,15 @@ class TestSampleReliableStreaming:
assert travel_planner is not None
assert travel_planner.name == "TravelPlanner"
# Create a new thread
thread = travel_planner.get_new_thread()
assert thread.session_id is not None
assert thread.session_id.key is not None
thread_id = str(thread.session_id.key)
# Create a new session
session = travel_planner.create_session()
assert session.durable_session_id is not None
assert session.durable_session_id.key is not None
session_key = str(session.durable_session_id.key)
# Start agent run with wait_for_response=False for non-blocking execution
travel_planner.run(
"Plan a 1-day trip to Seattle in 1 sentence", thread=thread, options={"wait_for_response": False}
"Plan a 1-day trip to Seattle in 1 sentence", session=session, options={"wait_for_response": False}
)
# Poll Redis stream with retries to handle race conditions
@@ -146,7 +146,7 @@ class TestSampleReliableStreaming:
while retry_count < max_retries and not is_complete:
text, is_complete, last_cursor = asyncio.run(
self._stream_from_redis(thread_id, cursor=cursor, timeout=10.0)
self._stream_from_redis(session_key, cursor=cursor, timeout=10.0)
)
accumulated_text += text
cursor = last_cursor # Resume from last position on next read
@@ -166,7 +166,7 @@ class TestSampleReliableStreaming:
# Verify we got content
assert len(accumulated_text) > 0, (
f"Expected text content but got empty string for thread_id: {thread_id} after {retry_count} retries"
f"Expected text content but got empty string for session_key: {session_key} after {retry_count} retries"
)
assert "seattle" in accumulated_text.lower(), f"Expected 'seattle' in response but got: {accumulated_text}"
assert is_complete, "Expected stream to be complete"
@@ -175,13 +175,13 @@ class TestSampleReliableStreaming:
"""Test streaming with cursor-based resumption."""
# Get the TravelPlanner agent
travel_planner = self.agent_client.get_agent("TravelPlanner")
thread = travel_planner.get_new_thread()
assert thread.session_id is not None
assert thread.session_id.key is not None
thread_id = str(thread.session_id.key)
session = travel_planner.create_session()
assert session.durable_session_id is not None
assert session.durable_session_id.key is not None
session_key = str(session.durable_session_id.key)
# Start agent run
travel_planner.run("What's the weather like?", thread=thread, options={"wait_for_response": False})
travel_planner.run("What's the weather like?", session=session, options={"wait_for_response": False})
# Wait for agent to start writing
time.sleep(3)
@@ -194,7 +194,7 @@ class TestSampleReliableStreaming:
chunk_count = 0
# Read just first 2 chunks
async for chunk in stream_handler.read_stream(thread_id): # type: ignore[reportUnknownMemberType]
async for chunk in stream_handler.read_stream(session_key): # type: ignore[reportUnknownMemberType]
last_entry_id = chunk.entry_id # type: ignore[reportUnknownMemberType]
if chunk.text: # type: ignore[reportUnknownMemberType]
accumulated_text += chunk.text # type: ignore[reportUnknownMemberType]
@@ -207,7 +207,7 @@ class TestSampleReliableStreaming:
partial_text, cursor = asyncio.run(get_partial_stream())
# Resume from cursor
remaining_text, _, _ = asyncio.run(self._stream_from_redis(thread_id, cursor=cursor))
remaining_text, _, _ = asyncio.run(self._stream_from_redis(session_key, cursor=cursor))
# Verify we got some initial content
assert len(partial_text) > 0
python/packages/durabletask/tests/test_agent_session_id.py
+108 -100
View File
@@ -1,11 +1,11 @@
# Copyright (c) Microsoft. All rights reserved.
"""Unit tests for AgentSessionId and DurableAgentThread."""
"""Unit tests for AgentSessionId and DurableAgentSession."""
import pytest
from agent_framework import AgentThread
from agent_framework import AgentSession
from agent_framework_durabletask._models import AgentSessionId, DurableAgentThread
from agent_framework_durabletask._models import AgentSessionId, DurableAgentSession
class TestAgentSessionId:
@@ -121,154 +121,162 @@ class TestAgentSessionId:
assert "Invalid agent session ID format" in str(exc_info.value)
class TestDurableAgentThread:
"""Test suite for DurableAgentThread."""
class TestDurableAgentSession:
"""Test suite for DurableAgentSession."""
def test_init_with_session_id(self) -> None:
"""Test DurableAgentThread initialization with session ID."""
def test_init_with_durable_session_id(self) -> None:
"""Test DurableAgentSession initialization with durable session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread = DurableAgentThread(session_id=session_id)
session = DurableAgentSession(durable_session_id=session_id)
assert thread.session_id is not None
assert thread.session_id == session_id
assert session.durable_session_id is not None
assert session.durable_session_id == session_id
def test_init_without_session_id(self) -> None:
"""Test DurableAgentThread initialization without session ID."""
thread = DurableAgentThread()
def test_init_without_durable_session_id(self) -> None:
"""Test DurableAgentSession initialization without durable session ID."""
session = DurableAgentSession()
assert thread.session_id is None
assert session.durable_session_id is None
def test_session_id_setter(self) -> None:
"""Test setting a session ID to an existing thread."""
thread = DurableAgentThread()
assert thread.session_id is None
def test_durable_session_id_setter(self) -> None:
"""Test setting a durable session ID to an existing session."""
session = DurableAgentSession()
assert session.durable_session_id is None
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread.session_id = session_id
session.durable_session_id = session_id
assert thread.session_id is not None
assert thread.session_id == session_id
assert thread.session_id.name == "TestAgent"
assert session.durable_session_id is not None
assert session.durable_session_id == session_id
assert session.durable_session_id.name == "TestAgent"
def test_from_session_id(self) -> None:
"""Test creating DurableAgentThread from session ID."""
"""Test creating DurableAgentSession from session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread = DurableAgentThread.from_session_id(session_id)
session = DurableAgentSession.from_session_id(session_id)
assert isinstance(thread, DurableAgentThread)
assert thread.session_id is not None
assert thread.session_id == session_id
assert thread.session_id.name == "TestAgent"
assert thread.session_id.key == "test-key"
assert isinstance(session, DurableAgentSession)
assert session.durable_session_id is not None
assert session.durable_session_id == session_id
assert session.durable_session_id.name == "TestAgent"
assert session.durable_session_id.key == "test-key"
def test_from_session_id_with_service_thread_id(self) -> None:
"""Test creating DurableAgentThread with service thread ID."""
def test_from_session_id_with_service_session_id(self) -> None:
"""Test creating DurableAgentSession with service session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread = DurableAgentThread.from_session_id(session_id, service_thread_id="service-123")
session = DurableAgentSession.from_session_id(session_id, service_session_id="service-123")
assert thread.session_id is not None
assert thread.session_id == session_id
assert thread.service_thread_id == "service-123"
assert session.durable_session_id is not None
assert session.durable_session_id == session_id
assert session.service_session_id == "service-123"
async def test_serialize_with_session_id(self) -> None:
"""Test serialization includes session ID."""
def test_to_dict_with_durable_session_id(self) -> None:
"""Test serialization includes durable session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key")
thread = DurableAgentThread(session_id=session_id)
session = DurableAgentSession(durable_session_id=session_id)
serialized = await thread.serialize()
serialized = session.to_dict()
assert isinstance(serialized, dict)
assert "durable_session_id" in serialized
assert serialized["durable_session_id"] == "@TestAgent@test-key"
async def test_serialize_without_session_id(self) -> None:
"""Test serialization without session ID."""
thread = DurableAgentThread()
def test_to_dict_without_durable_session_id(self) -> None:
"""Test serialization without durable session ID."""
session = DurableAgentSession()
serialized = await thread.serialize()
serialized = session.to_dict()
assert isinstance(serialized, dict)
assert "durable_session_id" not in serialized
async def test_deserialize_with_session_id(self) -> None:
"""Test deserialization restores session ID."""
def test_from_dict_with_durable_session_id(self) -> None:
"""Test deserialization restores durable session ID."""
serialized = {
"service_thread_id": "thread-123",
"type": "session",
"session_id": "session-123",
"service_session_id": "service-123",
"state": {},
"durable_session_id": "@TestAgent@test-key",
}
thread = await DurableAgentThread.deserialize(serialized)
session = DurableAgentSession.from_dict(serialized)
assert isinstance(thread, DurableAgentThread)
assert thread.session_id is not None
assert thread.session_id.name == "TestAgent"
assert thread.session_id.key == "test-key"
assert thread.service_thread_id == "thread-123"
assert isinstance(session, DurableAgentSession)
assert session.durable_session_id is not None
assert session.durable_session_id.name == "TestAgent"
assert session.durable_session_id.key == "test-key"
assert session.service_session_id == "service-123"
async def test_deserialize_without_session_id(self) -> None:
"""Test deserialization without session ID."""
def test_from_dict_without_durable_session_id(self) -> None:
"""Test deserialization without durable session ID."""
serialized = {
"service_thread_id": "thread-456",
"type": "session",
"session_id": "session-456",
"service_session_id": "service-456",
"state": {},
}
thread = await DurableAgentThread.deserialize(serialized)
session = DurableAgentSession.from_dict(serialized)
assert isinstance(thread, DurableAgentThread)
assert thread.session_id is None
assert thread.service_thread_id == "thread-456"
assert isinstance(session, DurableAgentSession)
assert session.durable_session_id is None
assert session.session_id == "session-456"
async def test_round_trip_serialization(self) -> None:
"""Test round-trip serialization preserves session ID."""
def test_round_trip_serialization(self) -> None:
"""Test round-trip serialization preserves durable session ID."""
session_id = AgentSessionId(name="TestAgent", key="test-key-789")
original = DurableAgentThread(session_id=session_id)
original = DurableAgentSession(durable_session_id=session_id)
serialized = await original.serialize()
restored = await DurableAgentThread.deserialize(serialized)
serialized = original.to_dict()
restored = DurableAgentSession.from_dict(serialized)
assert isinstance(restored, DurableAgentThread)
assert restored.session_id is not None
assert restored.session_id.name == session_id.name
assert restored.session_id.key == session_id.key
assert isinstance(restored, DurableAgentSession)
assert restored.durable_session_id is not None
assert restored.durable_session_id.name == session_id.name
assert restored.durable_session_id.key == session_id.key
async def test_deserialize_invalid_session_id_type(self) -> None:
"""Test deserialization with invalid session ID type raises error."""
def test_from_dict_invalid_durable_session_id_type(self) -> None:
"""Test deserialization with invalid durable session ID type raises error."""
serialized = {
"service_thread_id": "thread-123",
"type": "session",
"session_id": "session-123",
"state": {},
"durable_session_id": 12345, # Invalid type
}
with pytest.raises(ValueError, match="durable_session_id must be a string"):
await DurableAgentThread.deserialize(serialized)
DurableAgentSession.from_dict(serialized)
class TestAgentThreadCompatibility:
"""Test suite for compatibility between AgentThread and DurableAgentThread."""
class TestAgentSessionCompatibility:
"""Test suite for compatibility between AgentSession and DurableAgentSession."""
async def test_agent_thread_serialize(self) -> None:
"""Test that base AgentThread can be serialized."""
thread = AgentThread()
def test_agent_session_to_dict(self) -> None:
"""Test that base AgentSession can be serialized."""
session = AgentSession()
serialized = await thread.serialize()
serialized = session.to_dict()
assert isinstance(serialized, dict)
assert "service_thread_id" in serialized
assert "session_id" in serialized
async def test_agent_thread_deserialize(self) -> None:
"""Test that base AgentThread can be deserialized."""
thread = AgentThread()
serialized = await thread.serialize()
def test_agent_session_from_dict(self) -> None:
"""Test that base AgentSession can be deserialized."""
session = AgentSession()
serialized = session.to_dict()
restored = await AgentThread.deserialize(serialized)
restored = AgentSession.from_dict(serialized)
assert isinstance(restored, AgentThread)
assert restored.service_thread_id == thread.service_thread_id
assert isinstance(restored, AgentSession)
assert restored.session_id == session.session_id
async def test_durable_thread_is_agent_thread(self) -> None:
"""Test that DurableAgentThread is an AgentThread."""
thread = DurableAgentThread()
def test_durable_session_is_agent_session(self) -> None:
"""Test that DurableAgentSession is an AgentSession."""
session = DurableAgentSession()
assert isinstance(thread, AgentThread)
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, AgentSession)
assert isinstance(session, DurableAgentSession)
class TestModelIntegration:
@@ -281,19 +289,19 @@ class TestModelIntegration:
assert session_id_str.startswith("@AgentEntity@")
async def test_thread_with_session_preserves_on_serialization(self) -> None:
"""Test that thread with session ID preserves it through serialization."""
def test_session_with_durable_id_preserves_on_serialization(self) -> None:
"""Test that session with durable session ID preserves it through serialization."""
session_id = AgentSessionId(name="TestAgent", key="preserved-key")
thread = DurableAgentThread.from_session_id(session_id)
session = DurableAgentSession.from_session_id(session_id)
# Serialize and deserialize
serialized = await thread.serialize()
restored = await DurableAgentThread.deserialize(serialized)
serialized = session.to_dict()
restored = DurableAgentSession.from_dict(serialized)
# Session ID should be preserved
assert restored.session_id is not None
assert restored.session_id.name == "TestAgent"
assert restored.session_id.key == "preserved-key"
# Durable session ID should be preserved
assert restored.durable_session_id is not None
assert restored.durable_session_id.name == "TestAgent"
assert restored.durable_session_id.key == "preserved-key"
if __name__ == "__main__":
python/packages/durabletask/tests/test_client.py
+10 -10
View File
@@ -11,7 +11,7 @@ from unittest.mock import Mock
import pytest
from agent_framework import SupportsAgentRun
from agent_framework_durabletask import DurableAgentThread, DurableAIAgentClient
from agent_framework_durabletask import DurableAgentSession, DurableAIAgentClient
from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
from agent_framework_durabletask._shim import DurableAIAgent
@@ -80,22 +80,22 @@ class TestDurableAIAgentClientIntegration:
assert hasattr(agent, "run")
assert callable(agent.run)
def test_client_agent_can_create_threads(self, agent_client: DurableAIAgentClient) -> None:
"""Verify agent from client can create DurableAgentThread instances."""
def test_client_agent_can_create_sessions(self, agent_client: DurableAIAgentClient) -> None:
"""Verify agent from client can create DurableAgentSession instances."""
agent = agent_client.get_agent("assistant")
thread = agent.get_new_thread()
session = agent.create_session()
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, DurableAgentSession)
def test_client_agent_thread_with_parameters(self, agent_client: DurableAIAgentClient) -> None:
"""Verify agent can create threads with custom parameters."""
def test_client_agent_session_with_parameters(self, agent_client: DurableAIAgentClient) -> None:
"""Verify agent can create sessions with custom parameters."""
agent = agent_client.get_agent("assistant")
thread = agent.get_new_thread(service_thread_id="client-session-123")
session = agent.create_session(service_session_id="client-session-123")
assert isinstance(thread, DurableAgentThread)
assert thread.service_thread_id == "client-session-123"
assert isinstance(session, DurableAgentSession)
assert session.service_session_id == "client-session-123"
class TestDurableAIAgentClientPollingConfiguration:
python/packages/durabletask/tests/test_executors.py
+26 -26
View File
@@ -16,7 +16,7 @@ from durabletask.entities import EntityInstanceId
from durabletask.task import Task
from pydantic import BaseModel
from agent_framework_durabletask import DurableAgentThread
from agent_framework_durabletask import DurableAgentSession
from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
from agent_framework_durabletask._executors import (
ClientAgentExecutor,
@@ -106,42 +106,42 @@ def configure_failed_entity_task(mock_entity_task: Mock) -> Any:
return _configure
class TestExecutorThreadCreation:
"""Test that executors properly create DurableAgentThread with parameters."""
class TestExecutorSessionCreation:
"""Test that executors properly create DurableAgentSession with parameters."""
def test_client_executor_creates_durable_thread(self, mock_client: Mock) -> None:
"""Verify ClientAgentExecutor creates DurableAgentThread instances."""
def test_client_executor_creates_durable_session(self, mock_client: Mock) -> None:
"""Verify ClientAgentExecutor creates DurableAgentSession instances."""
executor = ClientAgentExecutor(mock_client)
thread = executor.get_new_thread("test_agent")
session = executor.get_new_session("test_agent")
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, DurableAgentSession)
def test_client_executor_forwards_kwargs_to_thread(self, mock_client: Mock) -> None:
"""Verify ClientAgentExecutor forwards kwargs to DurableAgentThread creation."""
def test_client_executor_forwards_kwargs_to_session(self, mock_client: Mock) -> None:
"""Verify ClientAgentExecutor forwards kwargs to DurableAgentSession creation."""
executor = ClientAgentExecutor(mock_client)
thread = executor.get_new_thread("test_agent", service_thread_id="client-123")
session = executor.get_new_session("test_agent", service_session_id="client-123")
assert isinstance(thread, DurableAgentThread)
assert thread.service_thread_id == "client-123"
assert isinstance(session, DurableAgentSession)
assert session.service_session_id == "client-123"
def test_orchestration_executor_creates_durable_thread(
def test_orchestration_executor_creates_durable_session(
self, orchestration_executor: OrchestrationAgentExecutor
) -> None:
"""Verify OrchestrationAgentExecutor creates DurableAgentThread instances."""
thread = orchestration_executor.get_new_thread("test_agent")
"""Verify OrchestrationAgentExecutor creates DurableAgentSession instances."""
session = orchestration_executor.get_new_session("test_agent")
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, DurableAgentSession)
def test_orchestration_executor_forwards_kwargs_to_thread(
def test_orchestration_executor_forwards_kwargs_to_session(
self, orchestration_executor: OrchestrationAgentExecutor
) -> None:
"""Verify OrchestrationAgentExecutor forwards kwargs to DurableAgentThread creation."""
thread = orchestration_executor.get_new_thread("test_agent", service_thread_id="orch-456")
"""Verify OrchestrationAgentExecutor forwards kwargs to DurableAgentSession creation."""
session = orchestration_executor.get_new_session("test_agent", service_session_id="orch-456")
assert isinstance(thread, DurableAgentThread)
assert thread.service_thread_id == "orch-456"
assert isinstance(session, DurableAgentSession)
assert session.service_session_id == "orch-456"
class TestClientAgentExecutorRun:
@@ -353,18 +353,18 @@ class TestOrchestrationAgentExecutorRun:
# Verify request dict
assert request_dict_arg == sample_run_request.to_dict()
def test_orchestration_executor_uses_thread_session_id(
def test_orchestration_executor_uses_session_durable_id(
self,
mock_orchestration_context: Mock,
orchestration_executor: OrchestrationAgentExecutor,
sample_run_request: RunRequest,
) -> None:
"""Verify executor uses thread's session ID when provided."""
# Create thread with specific session ID
"""Verify executor uses session's durable session ID when provided."""
# Create session with specific durable session ID
session_id = AgentSessionId(name="test_agent", key="specific-key-123")
thread = DurableAgentThread.from_session_id(session_id)
session = DurableAgentSession.from_session_id(session_id)
result = orchestration_executor.run_durable_agent("test_agent", sample_run_request, thread=thread)
result = orchestration_executor.run_durable_agent("test_agent", sample_run_request, session=session)
# Verify call_entity was called with the specific key
call_args = mock_orchestration_context.call_entity.call_args
python/packages/durabletask/tests/test_orchestration_context.py
+10 -10
View File
@@ -11,7 +11,7 @@ from unittest.mock import Mock
import pytest
from agent_framework import SupportsAgentRun
from agent_framework_durabletask import DurableAgentThread
from agent_framework_durabletask import DurableAgentSession
from agent_framework_durabletask._orchestration_context import DurableAIAgentOrchestrationContext
from agent_framework_durabletask._shim import DurableAIAgent
@@ -74,24 +74,24 @@ class TestDurableAIAgentOrchestrationContextIntegration:
assert hasattr(agent, "run")
assert callable(agent.run)
def test_orchestration_agent_can_create_threads(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
"""Verify agent from context can create DurableAgentThread instances."""
def test_orchestration_agent_can_create_sessions(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
"""Verify agent from context can create DurableAgentSession instances."""
agent = agent_context.get_agent("assistant")
thread = agent.get_new_thread()
session = agent.create_session()
assert isinstance(thread, DurableAgentThread)
assert isinstance(session, DurableAgentSession)
def test_orchestration_agent_thread_with_parameters(
def test_orchestration_agent_session_with_parameters(
self, agent_context: DurableAIAgentOrchestrationContext
) -> None:
"""Verify agent can create threads with custom parameters."""
"""Verify agent can create sessions with custom parameters."""
agent = agent_context.get_agent("assistant")
thread = agent.get_new_thread(service_thread_id="orch-session-456")
session = agent.create_session(service_session_id="orch-session-456")
assert isinstance(thread, DurableAgentThread)
assert thread.service_thread_id == "orch-session-456"
assert isinstance(session, DurableAgentSession)
assert session.service_session_id == "orch-session-456"
if __name__ == "__main__":
python/packages/durabletask/tests/test_shim.py
+24 -24
View File
@@ -13,7 +13,7 @@ import pytest
from agent_framework import Message, SupportsAgentRun
from pydantic import BaseModel
from agent_framework_durabletask import DurableAgentThread
from agent_framework_durabletask import DurableAgentSession
from agent_framework_durabletask._executors import DurableAgentExecutor
from agent_framework_durabletask._models import RunRequest
from agent_framework_durabletask._shim import DurableAgentProvider, DurableAIAgent
@@ -30,7 +30,7 @@ def mock_executor() -> Mock:
"""Create a mock executor for testing."""
mock = Mock(spec=DurableAgentExecutor)
mock.run_durable_agent = Mock(return_value=None)
mock.get_new_thread = Mock(return_value=DurableAgentThread())
mock.get_new_session = Mock(return_value=DurableAgentSession())
# Mock get_run_request to create actual RunRequest objects
def create_run_request(
@@ -124,14 +124,14 @@ class TestDurableAIAgentMessageNormalization:
class TestDurableAIAgentParameterFlow:
"""Test that parameters flow correctly through the shim to executor."""
def test_run_forwards_thread_parameter(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify run forwards thread parameter to executor."""
thread = DurableAgentThread(service_thread_id="test-thread")
test_agent.run("message", thread=thread)
def test_run_forwards_session_parameter(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify run forwards session parameter to executor."""
session = DurableAgentSession(service_session_id="test-session")
test_agent.run("message", session=session)
mock_executor.run_durable_agent.assert_called_once()
_, kwargs = mock_executor.run_durable_agent.call_args
assert kwargs["thread"] == thread
assert kwargs["session"] == session
def test_run_forwards_response_format(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify run forwards response_format parameter to executor."""
@@ -171,29 +171,29 @@ class TestDurableAISupportsAgentRunCompliance:
assert agent.name == "my_agent"
class TestDurableAIAgentThreadManagement:
"""Test thread creation and management."""
class TestDurableAIAgentSessionManagement:
"""Test session creation and management."""
def test_get_new_thread_delegates_to_executor(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify get_new_thread delegates to executor."""
mock_thread = DurableAgentThread()
mock_executor.get_new_thread.return_value = mock_thread
def test_create_session_delegates_to_executor(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify create_session delegates to executor."""
mock_session = DurableAgentSession()
mock_executor.get_new_session.return_value = mock_session
thread = test_agent.get_new_thread()
session = test_agent.create_session()
mock_executor.get_new_thread.assert_called_once_with("test_agent")
assert thread == mock_thread
mock_executor.get_new_session.assert_called_once_with("test_agent")
assert session == mock_session
def test_get_new_thread_forwards_kwargs(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify get_new_thread forwards kwargs to executor."""
mock_thread = DurableAgentThread(service_thread_id="thread-123")
mock_executor.get_new_thread.return_value = mock_thread
def test_create_session_forwards_kwargs(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
"""Verify create_session forwards kwargs to executor."""
mock_session = DurableAgentSession(service_session_id="session-123")
mock_executor.get_new_session.return_value = mock_session
test_agent.get_new_thread(service_thread_id="thread-123")
test_agent.create_session(service_session_id="session-123")
mock_executor.get_new_thread.assert_called_once()
_, kwargs = mock_executor.get_new_thread.call_args
assert kwargs["service_thread_id"] == "thread-123"
mock_executor.get_new_session.assert_called_once()
_, kwargs = mock_executor.get_new_session.call_args
assert kwargs["service_session_id"] == "session-123"
class TestDurableAgentProviderInterface:
python/packages/github_copilot/agent_framework_github_copilot/_agent.py
+29 -29
View File
@@ -13,10 +13,10 @@ from agent_framework import (
AgentMiddlewareTypes,
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
BaseAgent,
BaseContextProvider,
Content,
ContextProvider,
Message,
ResponseStream,
normalize_messages,
@@ -149,7 +149,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
id: str | None = None,
name: str | None = None,
description: str | None = None,
context_provider: ContextProvider | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[AgentMiddlewareTypes] | None = None,
tools: FunctionTool
| Callable[..., Any]
@@ -171,7 +171,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
id: ID of the GitHubCopilotAgent.
name: Name of the GitHubCopilotAgent.
description: Description of the GitHubCopilotAgent.
context_provider: Context Provider, to be used by the agent.
context_providers: Context Providers, to be used by the agent.
middleware: Agent middleware used by the agent.
tools: Tools to use for the agent. Can be functions
or tool definition dicts. These are converted to Copilot SDK tools internally.
@@ -187,7 +187,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
id=id,
name=name,
description=description,
context_provider=context_provider,
context_providers=context_providers,
middleware=list(middleware) if middleware else None,
)
@@ -280,7 +280,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: OptionsT | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse]: ...
@@ -291,7 +291,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: OptionsT | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ...
@@ -301,7 +301,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: OptionsT | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
@@ -316,7 +316,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
Keyword Args:
stream: Whether to stream the response. Defaults to False.
thread: The conversation thread associated with the message(s).
session: The conversation session associated with the message(s).
options: Runtime options (model, timeout, etc.).
kwargs: Additional keyword arguments.
@@ -333,16 +333,16 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
return AgentResponse.from_updates(updates)
return ResponseStream(
self._stream_updates(messages=messages, thread=thread, options=options, **kwargs),
self._stream_updates(messages=messages, session=session, options=options, **kwargs),
finalizer=_finalize,
)
return self._run_impl(messages=messages, thread=thread, options=options, **kwargs)
return self._run_impl(messages=messages, session=session, options=options, **kwargs)
async def _run_impl(
self,
messages: str | Message | Sequence[str | Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: OptionsT | None = None,
**kwargs: Any,
) -> AgentResponse:
@@ -350,18 +350,18 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
if not self._started:
await self.start()
if not thread:
thread = self.get_new_thread()
if not session:
session = self.create_session()
opts: dict[str, Any] = dict(options) if options else {}
timeout = opts.pop("timeout", None) or self._settings["timeout"] or DEFAULT_TIMEOUT_SECONDS
session = await self._get_or_create_session(thread, streaming=False, runtime_options=opts)
copilot_session = await self._get_or_create_session(session, streaming=False, runtime_options=opts)
input_messages = normalize_messages(messages)
prompt = "\n".join([message.text for message in input_messages])
try:
response_event = await session.send_and_wait({"prompt": prompt}, timeout=timeout)
response_event = await copilot_session.send_and_wait({"prompt": prompt}, timeout=timeout)
except Exception as ex:
raise ServiceException(f"GitHub Copilot request failed: {ex}") from ex
@@ -390,7 +390,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
self,
messages: str | Message | Sequence[str | Message] | None = None,
*,
thread: AgentThread | None = None,
session: AgentSession | None = None,
options: OptionsT | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
@@ -400,7 +400,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
messages: The message(s) to send to the agent.
Keyword Args:
thread: The conversation thread associated with the message(s).
session: The conversation session associated with the message(s).
options: Runtime options (model, timeout, etc.).
kwargs: Additional keyword arguments.
@@ -413,12 +413,12 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
if not self._started:
await self.start()
if not thread:
thread = self.get_new_thread()
if not session:
session = self.create_session()
opts: dict[str, Any] = dict(options) if options else {}
session = await self._get_or_create_session(thread, streaming=True, runtime_options=opts)
copilot_session = await self._get_or_create_session(session, streaming=True, runtime_options=opts)
input_messages = normalize_messages(messages)
prompt = "\n".join([message.text for message in input_messages])
@@ -441,10 +441,10 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
error_msg = event.data.message or "Unknown error"
queue.put_nowait(ServiceException(f"GitHub Copilot session error: {error_msg}"))
unsubscribe = session.on(event_handler)
unsubscribe = copilot_session.on(event_handler)
try:
await session.send({"prompt": prompt})
await copilot_session.send({"prompt": prompt})
while (item := await queue.get()) is not None:
if isinstance(item, Exception):
@@ -530,14 +530,14 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
async def _get_or_create_session(
self,
thread: AgentThread,
agent_session: AgentSession,
streaming: bool = False,
runtime_options: dict[str, Any] | None = None,
) -> CopilotSession:
"""Get an existing session or create a new one for the thread.
"""Get an existing session or create a new one for the session.
Args:
thread: The conversation thread.
agent_session: The conversation session.
streaming: Whether to enable streaming for the session.
runtime_options: Runtime options from run that take precedence.
@@ -551,11 +551,11 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]):
raise ServiceException("GitHub Copilot client not initialized. Call start() first.")
try:
if thread.service_thread_id:
return await self._resume_session(thread.service_thread_id, streaming)
if agent_session.service_session_id:
return await self._resume_session(agent_session.service_session_id, streaming)
session = await self._create_session(streaming, runtime_options)
thread.service_thread_id = session.session_id
agent_session.service_session_id = session.session_id
return session
except Exception as ex:
raise ServiceException(f"Failed to create GitHub Copilot session: {ex}") from ex
python/packages/github_copilot/tests/test_github_copilot_agent.py
+40 -40
View File
@@ -10,7 +10,7 @@ import pytest
from agent_framework import (
AgentResponse,
AgentResponseUpdate,
AgentThread,
AgentSession,
Content,
Message,
)
@@ -300,21 +300,21 @@ class TestGitHubCopilotAgentRun:
assert isinstance(response, AgentResponse)
assert len(response.messages) == 1
async def test_run_with_thread(
async def test_run_with_session(
self,
mock_client: MagicMock,
mock_session: MagicMock,
assistant_message_event: SessionEvent,
) -> None:
"""Test run method with existing thread."""
"""Test run method with existing session."""
mock_session.send_and_wait.return_value = assistant_message_event
agent = GitHubCopilotAgent(client=mock_client)
thread = AgentThread()
response = await agent.run("Hello", thread=thread)
session = AgentSession()
response = await agent.run("Hello", session=session)
assert isinstance(response, AgentResponse)
assert thread.service_thread_id == mock_session.session_id
assert session.service_session_id == mock_session.session_id
async def test_run_with_runtime_options(
self,
@@ -392,13 +392,13 @@ class TestGitHubCopilotAgentRunStreaming:
assert responses[0].role == "assistant"
assert responses[0].contents[0].text == "Hello"
async def test_run_streaming_with_thread(
async def test_run_streaming_with_session(
self,
mock_client: MagicMock,
mock_session: MagicMock,
session_idle_event: SessionEvent,
) -> None:
"""Test streaming with existing thread."""
"""Test streaming with existing session."""
def mock_on(handler: Any) -> Any:
handler(session_idle_event)
@@ -407,12 +407,12 @@ class TestGitHubCopilotAgentRunStreaming:
mock_session.on = mock_on
agent = GitHubCopilotAgent(client=mock_client)
thread = AgentThread()
session = AgentSession()
async for _ in agent.run("Hello", thread=thread, stream=True):
async for _ in agent.run("Hello", session=session, stream=True):
pass
assert thread.service_thread_id == mock_session.session_id
assert session.service_session_id == mock_session.session_id
async def test_run_streaming_error(
self,
@@ -461,20 +461,20 @@ class TestGitHubCopilotAgentRunStreaming:
class TestGitHubCopilotAgentSessionManagement:
"""Test cases for session management."""
async def test_session_resumed_for_same_thread(
async def test_session_resumed_for_same_session(
self,
mock_client: MagicMock,
mock_session: MagicMock,
assistant_message_event: SessionEvent,
) -> None:
"""Test that subsequent calls on the same thread resume the session."""
"""Test that subsequent calls on the same session resume the session."""
mock_session.send_and_wait.return_value = assistant_message_event
agent = GitHubCopilotAgent(client=mock_client)
thread = AgentThread()
session = AgentSession()
await agent.run("Hello", thread=thread)
await agent.run("World", thread=thread)
await agent.run("Hello", session=session)
await agent.run("World", session=session)
mock_client.create_session.assert_called_once()
mock_client.resume_session.assert_called_once_with(mock_session.session_id, unittest.mock.ANY)
@@ -490,7 +490,7 @@ class TestGitHubCopilotAgentSessionManagement:
)
await agent.start()
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
@@ -508,7 +508,7 @@ class TestGitHubCopilotAgentSessionManagement:
)
await agent.start()
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
@@ -531,7 +531,7 @@ class TestGitHubCopilotAgentSessionManagement:
"system_message": {"mode": "replace", "content": "Runtime instructions"}
}
await agent._get_or_create_session( # type: ignore
AgentThread(),
AgentSession(),
runtime_options=runtime_options,
)
@@ -549,25 +549,25 @@ class TestGitHubCopilotAgentSessionManagement:
agent = GitHubCopilotAgent(client=mock_client)
await agent.start()
await agent._get_or_create_session(AgentThread(), streaming=True) # type: ignore
await agent._get_or_create_session(AgentSession(), streaming=True) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
assert config["streaming"] is True
async def test_resume_session_with_existing_service_thread_id(
async def test_resume_session_with_existing_service_session_id(
self,
mock_client: MagicMock,
mock_session: MagicMock,
) -> None:
"""Test that session is resumed when thread has a service_thread_id."""
"""Test that session is resumed when session has a service_session_id."""
agent = GitHubCopilotAgent(client=mock_client)
await agent.start()
thread = AgentThread()
thread.service_thread_id = "existing-session-id"
session = AgentSession()
session.service_session_id = "existing-session-id"
await agent._get_or_create_session(thread) # type: ignore
await agent._get_or_create_session(session) # type: ignore
mock_client.create_session.assert_not_called()
mock_client.resume_session.assert_called_once()
@@ -596,10 +596,10 @@ class TestGitHubCopilotAgentSessionManagement:
)
await agent.start()
thread = AgentThread()
thread.service_thread_id = "existing-session-id"
session = AgentSession()
session.service_session_id = "existing-session-id"
await agent._get_or_create_session(thread) # type: ignore
await agent._get_or_create_session(session) # type: ignore
mock_client.resume_session.assert_called_once()
call_args = mock_client.resume_session.call_args
@@ -639,7 +639,7 @@ class TestGitHubCopilotAgentMCPServers:
)
await agent.start()
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
@@ -672,10 +672,10 @@ class TestGitHubCopilotAgentMCPServers:
)
await agent.start()
thread = AgentThread()
thread.service_thread_id = "existing-session-id"
session = AgentSession()
session.service_session_id = "existing-session-id"
await agent._get_or_create_session(thread) # type: ignore
await agent._get_or_create_session(session) # type: ignore
mock_client.resume_session.assert_called_once()
call_args = mock_client.resume_session.call_args
@@ -692,7 +692,7 @@ class TestGitHubCopilotAgentMCPServers:
agent = GitHubCopilotAgent(client=mock_client)
await agent.start()
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
@@ -716,7 +716,7 @@ class TestGitHubCopilotAgentToolConversion:
agent = GitHubCopilotAgent(client=mock_client, tools=[my_tool])
await agent.start()
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
@@ -739,7 +739,7 @@ class TestGitHubCopilotAgentToolConversion:
agent = GitHubCopilotAgent(client=mock_client, tools=[my_tool])
await agent.start()
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
@@ -764,7 +764,7 @@ class TestGitHubCopilotAgentToolConversion:
agent = GitHubCopilotAgent(client=mock_client, tools=[failing_tool])
await agent.start()
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
@@ -867,7 +867,7 @@ class TestGitHubCopilotAgentErrorHandling:
await agent.start()
with pytest.raises(ServiceException, match="Failed to create GitHub Copilot session"):
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
async def test_get_or_create_session_raises_when_client_not_initialized(self) -> None:
"""Test that _get_or_create_session raises ServiceException when client is not initialized."""
@@ -875,7 +875,7 @@ class TestGitHubCopilotAgentErrorHandling:
# Don't call start() - client remains None
with pytest.raises(ServiceException, match="GitHub Copilot client not initialized"):
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
class TestGitHubCopilotAgentPermissions:
@@ -919,7 +919,7 @@ class TestGitHubCopilotAgentPermissions:
)
await agent.start()
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
@@ -935,7 +935,7 @@ class TestGitHubCopilotAgentPermissions:
agent = GitHubCopilotAgent(client=mock_client)
await agent.start()
await agent._get_or_create_session(AgentThread()) # type: ignore
await agent._get_or_create_session(AgentSession()) # type: ignore
call_args = mock_client.create_session.call_args
config = call_args[0][0]
python/packages/lab/tau2/agent_framework_lab_tau2/_sliding_window.py
+31 -40
View File
@@ -1,63 +1,57 @@
# Copyright (c) Microsoft. All rights reserved.
import json
from collections.abc import Sequence
from typing import Any
import tiktoken
from agent_framework import ChatMessageStore, Message
from agent_framework import InMemoryHistoryProvider, Message
from loguru import logger
class SlidingWindowChatMessageStore(ChatMessageStore):
"""A token-aware sliding window implementation of ChatMessageStore.
class SlidingWindowHistoryProvider(InMemoryHistoryProvider):
"""A token-aware sliding window implementation of InMemoryHistoryProvider.
Maintains two message lists: complete history and truncated window.
Automatically removes oldest messages when token limit is exceeded.
Also removes leading tool messages to ensure valid conversation flow.
Stores all messages in session state but returns a truncated window from
``get_messages`` that fits within ``max_tokens``. Automatically removes
oldest messages and leading tool messages to ensure valid conversation flow.
"""
def __init__(
self,
messages: Sequence[Message] | None = None,
source_id: str = "memory",
*,
max_tokens: int = 3800,
system_message: str | None = None,
tool_definitions: Any | None = None,
):
super().__init__(messages=messages)
self.truncated_messages = self.messages.copy()
super().__init__(source_id)
self.max_tokens = max_tokens
self.system_message = system_message # Included in token count
self.tool_definitions = tool_definitions
# An estimation based on a commonly used vocab table
self.encoding = tiktoken.get_encoding("o200k_base")
async def add_messages(self, messages: Sequence[Message]) -> None:
await super().add_messages(messages)
async def get_messages(
self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any
) -> list[Message]:
"""Retrieve messages from session state, truncated to fit within max_tokens."""
all_messages = await super().get_messages(session_id, state=state, **kwargs)
return self._truncate(list(all_messages))
self.truncated_messages = self.messages.copy()
self.truncate_messages()
async def list_messages(self) -> list[Message]:
"""Get the current list of messages, which may be truncated."""
return self.truncated_messages
async def list_all_messages(self) -> list[Message]:
"""Get all messages from the store including the truncated ones."""
return self.messages
def truncate_messages(self) -> None:
while len(self.truncated_messages) > 0 and self.get_token_count() > self.max_tokens:
def _truncate(self, messages: list[Message]) -> list[Message]:
"""Truncate messages to fit within max_tokens and remove leading tool messages."""
while len(messages) > 0 and self._get_token_count(messages) > self.max_tokens:
logger.warning("Messages exceed max tokens. Truncating oldest message.")
self.truncated_messages.pop(0)
messages.pop(0)
# Remove leading tool messages
while len(self.truncated_messages) > 0:
if self.truncated_messages[0].role != "tool":
while len(messages) > 0:
if messages[0].role != "tool":
break
logger.warning("Removing leading tool message because tool result cannot be the first message.")
self.truncated_messages.pop(0)
messages.pop(0)
return messages
def get_token_count(self) -> int:
def _get_token_count(self, messages: list[Message]) -> int:
"""Estimate token count for a list of messages using tiktoken.
Returns:
@@ -70,7 +64,7 @@ class SlidingWindowChatMessageStore(ChatMessageStore):
total_tokens += len(self.encoding.encode(self.system_message))
total_tokens += 4 # Extra tokens for system message formatting
for msg in self.truncated_messages:
for msg in messages:
# Add 4 tokens per message for role, formatting, etc.
total_tokens += 4
@@ -87,7 +81,7 @@ class SlidingWindowChatMessageStore(ChatMessageStore):
"name": content.name,
"arguments": content.arguments,
}
total_tokens += self.estimate_any_object_token_count(func_call_data)
total_tokens += self._estimate_any_object_token_count(func_call_data)
elif content.type == "function_result":
total_tokens += 4
# Serialize function result and count tokens
@@ -95,19 +89,16 @@ class SlidingWindowChatMessageStore(ChatMessageStore):
"call_id": content.call_id,
"result": content.result,
}
total_tokens += self.estimate_any_object_token_count(func_result_data)
total_tokens += self._estimate_any_object_token_count(func_result_data)
else:
# For other content types, serialize the whole content
total_tokens += self.estimate_any_object_token_count(content)
total_tokens += self._estimate_any_object_token_count(content)
else:
# Content without type, treat as text
total_tokens += self.estimate_any_object_token_count(content)
total_tokens += self._estimate_any_object_token_count(content)
elif hasattr(msg, "text") and msg.text:
# Simple text message
total_tokens += self.estimate_any_object_token_count(msg.text)
else:
# Skip it
pass
total_tokens += self._estimate_any_object_token_count(msg.text)
if total_tokens > self.max_tokens / 2:
logger.opt(colors=True).warning(
@@ -122,7 +113,7 @@ class SlidingWindowChatMessageStore(ChatMessageStore):
return total_tokens
def estimate_any_object_token_count(self, obj: Any) -> int:
def _estimate_any_object_token_count(self, obj: Any) -> int:
try:
serialized = json.dumps(obj)
except Exception:
python/packages/lab/tau2/agent_framework_lab_tau2/runner.py
+13 -11
View File
@@ -3,7 +3,7 @@
from __future__ import annotations
import uuid
from typing import cast
from typing import Any
from agent_framework import (
Agent,
@@ -32,7 +32,7 @@ from tau2.user.user_simulator import ( # type: ignore[import-untyped]
from tau2.utils.utils import get_now # type: ignore[import-untyped]
from ._message_utils import flip_messages, log_messages
from ._sliding_window import SlidingWindowChatMessageStore
from ._sliding_window import SlidingWindowHistoryProvider
from ._tau2_utils import convert_agent_framework_messages_to_tau2_messages, convert_tau2_tool_to_function_tool
__all__ = ["ASSISTANT_AGENT_ID", "ORCHESTRATOR_ID", "USER_SIMULATOR_ID", "TaskRunner"]
@@ -201,11 +201,13 @@ class TaskRunner:
instructions=assistant_system_prompt,
tools=tools,
temperature=self.assistant_sampling_temperature,
chat_message_store_factory=lambda: SlidingWindowChatMessageStore(
system_message=assistant_system_prompt,
tool_definitions=[tool.openai_schema for tool in tools],
max_tokens=self.assistant_window_size,
),
context_providers=[
SlidingWindowHistoryProvider(
system_message=assistant_system_prompt,
tool_definitions=[tool.openai_schema for tool in tools],
max_tokens=self.assistant_window_size,
)
],
)
def user_simulator(self, user_simuator_chat_client: SupportsChatGetResponse, task: Task) -> Agent:
@@ -354,11 +356,11 @@ class TaskRunner:
# STEP 5: Ensemble the conversation history needed for evaluation.
# It's coming from three parts:
# 1. The initial greeting
# 2. The assistant's message store (not just the truncated window)
# 2. The assistant's session state (full history, not just the truncated window)
# 3. The final user message (if any)
assistant_executor = cast(AgentExecutor, self._assistant_executor)
message_store = cast(SlidingWindowChatMessageStore, assistant_executor._agent_thread.message_store)
full_conversation = [first_message] + await message_store.list_all_messages()
session_state: dict[str, Any] = self._assistant_executor._session.state # type: ignore
all_messages: list[Message] = list(session_state.get("memory", {}).get("messages", [])) # type: ignore
full_conversation = [first_message, *all_messages]
if self._final_user_message is not None:
full_conversation.extend(self._final_user_message)
python/packages/lab/tau2/tests/test_sliding_window.py
+84 -149
View File
@@ -1,145 +1,120 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for sliding window message list."""
"""Tests for sliding window history provider."""
from unittest.mock import patch
from agent_framework._types import Content, Message
from agent_framework_lab_tau2._sliding_window import SlidingWindowChatMessageStore
from agent_framework_lab_tau2._sliding_window import SlidingWindowHistoryProvider
def test_initialization_empty():
"""Test initializing with no messages."""
sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
assert sliding_window.max_tokens == 1000
assert sliding_window.system_message is None
assert sliding_window.tool_definitions is None
assert len(sliding_window.messages) == 0
assert len(sliding_window.truncated_messages) == 0
def _make_state(provider: SlidingWindowHistoryProvider, messages: list[Message] | None = None) -> dict:
"""Helper to create a session state dict with messages pre-loaded."""
state: dict = {}
if messages:
state[provider.source_id] = {"messages": list(messages)}
return state
def test_initialization_with_parameters():
"""Test initializing with system message and tool definitions."""
system_msg = "You are a helpful assistant"
tool_defs = [{"name": "test_tool", "description": "A test tool"}]
sliding_window = SlidingWindowChatMessageStore(
max_tokens=2000, system_message=system_msg, tool_definitions=tool_defs
def test_initialization():
"""Test initializing with parameters."""
provider = SlidingWindowHistoryProvider(
max_tokens=2000,
system_message="You are a helpful assistant",
tool_definitions=[{"name": "test_tool"}],
)
assert sliding_window.max_tokens == 2000
assert sliding_window.system_message == system_msg
assert sliding_window.tool_definitions == tool_defs
assert provider.max_tokens == 2000
assert provider.system_message == "You are a helpful assistant"
assert provider.tool_definitions == [{"name": "test_tool"}]
assert provider.source_id == "memory"
def test_initialization_with_messages():
"""Test initializing with existing messages."""
messages = [
Message(role="user", contents=[Content.from_text(text="Hello")]),
Message(role="assistant", contents=[Content.from_text(text="Hi there!")]),
]
sliding_window = SlidingWindowChatMessageStore(messages=messages, max_tokens=1000)
assert len(sliding_window.messages) == 2
assert len(sliding_window.truncated_messages) == 2
async def test_get_messages_empty():
"""Test getting messages from empty state."""
provider = SlidingWindowHistoryProvider(max_tokens=1000)
messages = await provider.get_messages(None, state={})
assert messages == []
async def test_add_messages_simple():
"""Test adding messages without truncation."""
sliding_window = SlidingWindowChatMessageStore(max_tokens=10000) # Large limit
new_messages = [
async def test_get_messages_simple():
"""Test getting messages without truncation."""
provider = SlidingWindowHistoryProvider(max_tokens=10000)
msgs = [
Message(role="user", contents=[Content.from_text(text="What's the weather?")]),
Message(role="assistant", contents=[Content.from_text(text="I can help with that.")]),
]
state = _make_state(provider, msgs)
await sliding_window.add_messages(new_messages)
messages = await sliding_window.list_messages()
assert len(messages) == 2
assert messages[0].text == "What's the weather?"
assert messages[1].text == "I can help with that."
result = await provider.get_messages(None, state=state)
assert len(result) == 2
assert result[0].text == "What's the weather?"
assert result[1].text == "I can help with that."
async def test_list_all_messages_vs_list_messages():
"""Test difference between list_all_messages and list_messages."""
sliding_window = SlidingWindowChatMessageStore(max_tokens=50) # Small limit to force truncation
async def test_save_and_get_messages():
"""Test saving then getting messages with truncation."""
provider = SlidingWindowHistoryProvider(max_tokens=50)
state: dict = {}
# Add many messages to trigger truncation
messages = [
# Save many messages
msgs = [
Message(role="user", contents=[Content.from_text(text=f"Message {i} with some content")]) for i in range(10)
]
await provider.save_messages(None, msgs, state=state)
await sliding_window.add_messages(messages)
# get_messages returns truncated
truncated = await provider.get_messages(None, state=state)
# Full history is in session state
all_msgs = state[provider.source_id]["messages"]
truncated_messages = await sliding_window.list_messages()
all_messages = await sliding_window.list_all_messages()
# All messages should contain everything
assert len(all_messages) == 10
# Truncated messages should be fewer due to token limit
assert len(truncated_messages) < len(all_messages)
assert len(all_msgs) == 10
assert len(truncated) < len(all_msgs)
def test_get_token_count_basic():
"""Test basic token counting."""
sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
sliding_window.truncated_messages = [Message(role="user", contents=[Content.from_text(text="Hello")])]
provider = SlidingWindowHistoryProvider(max_tokens=1000)
messages = [Message(role="user", contents=[Content.from_text(text="Hello")])]
token_count = sliding_window.get_token_count()
# Should be more than 0 (exact count depends on encoding)
token_count = provider._get_token_count(messages)
assert token_count > 0
def test_get_token_count_with_system_message():
"""Test token counting includes system message."""
system_msg = "You are a helpful assistant"
sliding_window = SlidingWindowChatMessageStore(max_tokens=1000, system_message=system_msg)
provider = SlidingWindowHistoryProvider(max_tokens=1000, system_message="You are a helpful assistant")
# Without messages
token_count_empty = sliding_window.get_token_count()
count_empty = provider._get_token_count([])
count_with_msg = provider._get_token_count([Message(role="user", contents=[Content.from_text(text="Hello")])])
# Add a message
sliding_window.truncated_messages = [Message(role="user", contents=[Content.from_text(text="Hello")])]
token_count_with_message = sliding_window.get_token_count()
# With message should be more tokens
assert token_count_with_message > token_count_empty
assert token_count_empty > 0 # System message contributes tokens
assert count_with_msg > count_empty
assert count_empty > 0 # System message contributes tokens
def test_get_token_count_function_call():
"""Test token counting with function calls."""
function_call = Content.from_function_call(call_id="call_123", name="test_function", arguments={"param": "value"})
provider = SlidingWindowHistoryProvider(max_tokens=1000)
sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
sliding_window.truncated_messages = [Message(role="assistant", contents=[function_call])]
token_count = sliding_window.get_token_count()
token_count = provider._get_token_count([Message(role="assistant", contents=[function_call])])
assert token_count > 0
def test_get_token_count_function_result():
"""Test token counting with function results."""
function_result = Content.from_function_result(call_id="call_123", result={"success": True, "data": "result"})
provider = SlidingWindowHistoryProvider(max_tokens=1000)
sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
sliding_window.truncated_messages = [Message(role="tool", contents=[function_result])]
token_count = sliding_window.get_token_count()
token_count = provider._get_token_count([Message(role="tool", contents=[function_result])])
assert token_count > 0
@patch("agent_framework_lab_tau2._sliding_window.logger")
def test_truncate_messages_removes_old_messages(mock_logger):
def test_truncate_removes_old_messages(mock_logger):
"""Test that truncation removes old messages when token limit exceeded."""
sliding_window = SlidingWindowChatMessageStore(max_tokens=20) # Very small limit
provider = SlidingWindowHistoryProvider(max_tokens=20)
# Create messages that will exceed the limit
messages = [
Message(
role="user",
@@ -154,80 +129,45 @@ def test_truncate_messages_removes_old_messages(mock_logger):
Message(role="user", contents=[Content.from_text(text="Short msg")]),
]
sliding_window.truncated_messages = messages.copy()
sliding_window.truncate_messages()
# Should have fewer messages after truncation
assert len(sliding_window.truncated_messages) < len(messages)
# Should have logged warnings
result = provider._truncate(list(messages))
assert len(result) < len(messages)
assert mock_logger.warning.called
@patch("agent_framework_lab_tau2._sliding_window.logger")
def test_truncate_messages_removes_leading_tool_messages(mock_logger):
def test_truncate_removes_leading_tool_messages(mock_logger):
"""Test that truncation removes leading tool messages."""
sliding_window = SlidingWindowChatMessageStore(max_tokens=10000) # Large limit
provider = SlidingWindowHistoryProvider(max_tokens=10000)
# Create messages starting with tool message
tool_message = Message(role="tool", contents=[Content.from_function_result(call_id="call_123", result="result")])
user_message = Message(role="user", contents=[Content.from_text(text="Hello")])
sliding_window.truncated_messages = [tool_message, user_message]
sliding_window.truncate_messages()
# Tool message should be removed from the beginning
assert len(sliding_window.truncated_messages) == 1
assert sliding_window.truncated_messages[0].role == "user"
# Should have logged warning about removing tool message
result = provider._truncate([tool_message, user_message])
assert len(result) == 1
assert result[0].role == "user"
mock_logger.warning.assert_called()
def test_estimate_any_object_token_count_dict():
"""Test token counting for dictionary objects."""
sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
def test_estimate_any_object_token_count():
"""Test token counting for various object types."""
provider = SlidingWindowHistoryProvider(max_tokens=1000)
test_dict = {"key": "value", "number": 42}
token_count = sliding_window.estimate_any_object_token_count(test_dict)
assert provider._estimate_any_object_token_count({"key": "value"}) > 0
assert provider._estimate_any_object_token_count("test string") > 0
assert token_count > 0
def test_estimate_any_object_token_count_string():
"""Test token counting for string objects."""
sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
test_string = "This is a test string"
token_count = sliding_window.estimate_any_object_token_count(test_string)
assert token_count > 0
def test_estimate_any_object_token_count_non_serializable():
"""Test token counting for non-JSON-serializable objects."""
sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
# Create an object that can't be JSON serialized
class CustomObject:
# Non-serializable falls back to str()
class Custom:
def __str__(self):
return "CustomObject instance"
return "Custom instance"
custom_obj = CustomObject()
token_count = sliding_window.estimate_any_object_token_count(custom_obj)
# Should fall back to string representation
assert token_count > 0
assert provider._estimate_any_object_token_count(Custom()) > 0
async def test_real_world_scenario():
"""Test a realistic conversation scenario."""
sliding_window = SlidingWindowChatMessageStore(
max_tokens=30,
system_message="You are a helpful assistant", # Moderate limit
)
provider = SlidingWindowHistoryProvider(max_tokens=30, system_message="You are a helpful assistant")
state: dict = {}
# Simulate a conversation
conversation = [
Message(role="user", contents=[Content.from_text(text="Hello, how are you?")]),
Message(
@@ -253,18 +193,13 @@ async def test_real_world_scenario():
),
]
await sliding_window.add_messages(conversation)
await provider.save_messages(None, conversation, state=state)
current_messages = await sliding_window.list_messages()
all_messages = await sliding_window.list_all_messages()
truncated = await provider.get_messages(None, state=state)
all_msgs = state[provider.source_id]["messages"]
# All messages should be preserved
assert len(all_messages) == 6
assert len(all_msgs) == 6
assert len(truncated) <= 6
# Current messages might be truncated
assert len(current_messages) <= 6
# Token count should be within or close to limit
token_count = sliding_window.get_token_count()
# Allow some margin since truncation happens when exceeded
assert token_count <= sliding_window.max_tokens * 1.1
token_count = provider._get_token_count(truncated)
assert token_count <= provider.max_tokens * 1.1
python/packages/mem0/README.md
+1 -1
View File
@@ -27,5 +27,5 @@ Mem0's telemetry is **disabled by default** when using this package. If you want
import os
os.environ["MEM0_TELEMETRY"] = "true"
from agent_framework.mem0 import Mem0Provider
from agent_framework.mem0 import Mem0ContextProvider
```
python/packages/mem0/agent_framework_mem0/__init__.py
+2 -4
View File
@@ -8,8 +8,7 @@ import os
if os.environ.get("MEM0_TELEMETRY") is None:
os.environ["MEM0_TELEMETRY"] = "false"
from ._context_provider import _Mem0ContextProvider
from ._provider import Mem0Provider
from ._context_provider import Mem0ContextProvider
try:
__version__ = importlib.metadata.version(__name__)
@@ -17,7 +16,6 @@ except importlib.metadata.PackageNotFoundError:
__version__ = "0.0.0" # Fallback for development mode
__all__ = [
"Mem0Provider",
"_Mem0ContextProvider",
"Mem0ContextProvider",
"__version__",
]
python/packages/mem0/agent_framework_mem0/_context_provider.py
+14 -14
View File
@@ -2,9 +2,8 @@
"""New-pattern Mem0 context provider using BaseContextProvider.
This module provides ``_Mem0ContextProvider``, a side-by-side implementation of
:class:`Mem0Provider` built on the new :class:`BaseContextProvider` hooks pattern.
It will be renamed to ``Mem0ContextProvider`` in PR2 when the old class is removed.
This module provides ``Mem0ContextProvider``, built on the new
:class:`BaseContextProvider` hooks pattern.
"""
from __future__ import annotations
@@ -35,17 +34,11 @@ class _MemorySearchResponse_v1_1(TypedDict):
_MemorySearchResponse_v2 = list[dict[str, Any]]
class _Mem0ContextProvider(BaseContextProvider):
class Mem0ContextProvider(BaseContextProvider):
"""Mem0 context provider using the new BaseContextProvider hooks pattern.
Integrates Mem0 for persistent semantic memory, searching and storing
memories via the Mem0 API. This is the new-pattern equivalent of
:class:`Mem0Provider`.
Note:
This class uses a temporary ``_`` prefix to coexist with the existing
:class:`Mem0Provider`. It will be renamed to ``Mem0ContextProvider``
in PR2.
memories via the Mem0 API.
"""
DEFAULT_CONTEXT_PROMPT = "## Memories\nConsider the following memories when answering user questions:"
@@ -115,9 +108,16 @@ class _Mem0ContextProvider(BaseContextProvider):
filters = self._build_filters(session_id=context.session_id)
# AsyncMemory (OSS) expects user_id/agent_id/run_id as direct kwargs
# AsyncMemoryClient (Platform) expects them in a filters dict
search_kwargs: dict[str, Any] = {"query": input_text}
if isinstance(self.mem0_client, AsyncMemory):
search_kwargs.update(filters)
else:
search_kwargs["filters"] = filters
search_response: _MemorySearchResponse_v1_1 | _MemorySearchResponse_v2 = await self.mem0_client.search( # type: ignore[misc]
query=input_text,
filters=filters,
**search_kwargs,
)
if isinstance(search_response, list):
@@ -190,4 +190,4 @@ class _Mem0ContextProvider(BaseContextProvider):
return filters
__all__ = ["_Mem0ContextProvider"]
__all__ = ["Mem0ContextProvider"]

Some files were not shown because too many files have changed in this diff Show More