mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: [BREAKING] PR2 — Wire context provider pipeline, remove old types, update all consumers (#3850)
* 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:
committed by
GitHub
Unverified
parent
0c67dbbce5
commit
1e350ea22f
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user