mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: PR1 — New session and context provider types (side-by-side) (#3763)
* PR1: Add core context provider types and tests New types in _sessions.py (no changes to existing code): - SessionContext: per-invocation state with extend_messages/get_messages/ extend_instructions/extend_tools and read-only response property - _ContextProviderBase: base class with before_run/after_run hooks - _HistoryProviderBase: storage base with load/store flags, abstract get_messages/save_messages, default before_run/after_run - AgentSession: lightweight session with state dict, to_dict/from_dict - InMemoryHistoryProvider: built-in provider storing in session.state 35 unit tests covering all classes and configuration flags. * feat: keyword-only params, stateless InMemoryHistoryProvider, deep serialization - Make before_run/after_run parameters keyword-only - InMemoryHistoryProvider stores ChatMessage objects directly (no per-cycle serialization) - Deep serialization via to_dict/from_dict only at session boundary - State type registry for automatic deserialization of registered types - Updated tests for new serialization approach * feat: add new-pattern provider implementations for external packages - _RedisContextProvider(BaseContextProvider) - Redis search/vector context - _RedisHistoryProvider(BaseHistoryProvider) - Redis-backed message storage - _Mem0ContextProvider(BaseContextProvider) - Mem0 semantic memory - _AzureAISearchContextProvider(BaseContextProvider) - Azure AI Search (semantic + agentic) All use temporary _ prefix names for side-by-side coexistence with existing providers. Will be renamed in PR2 when old ContextProvider/ChatMessageStore are removed. * test: add tests for new-pattern provider implementations - 32 tests for _RedisContextProvider and _RedisHistoryProvider - 29 tests for _Mem0ContextProvider - 17 tests for _AzureAISearchContextProvider * fix: address PR review comments and CI failures - Move module docstring before imports in _sessions.py (review comment) - Import TYPE_CHECKING unconditionally in Redis _context_provider.py (NameError on Python <3.12) - Fix Mem0 test_init_auto_creates_client_when_none to patch at class level * feat: add source attribution to extend_messages Set attribution marker in additional_properties for each message added via extend_messages(), matching the tool attribution pattern. Uses setdefault to preserve any existing attribution. * refactor: make attribution value a dict with source_id key * add attribution and use sets for filters * Add source_type to message attribution and copy messages in extend_messages - SessionContext.extend_messages now accepts source as str or object with source_id attribute; when an object is passed, its class name is recorded as source_type in the attribution dict - Messages are shallow-copied before attribution is added so callers' original objects are never mutated - Filter framework-internal keys (attribution) from A2A wire metadata to prevent leaking internal state over the wire * fix: correct mypy type: ignore comment from union-attr to attr-defined * set attribution to _attribution * adjusted naming of bools
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
ccff3d3452
commit
ac0e6b0ee1
@@ -0,0 +1,421 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
|
||||
from agent_framework import ChatMessage
|
||||
from agent_framework._sessions import (
|
||||
AgentSession,
|
||||
BaseContextProvider,
|
||||
BaseHistoryProvider,
|
||||
InMemoryHistoryProvider,
|
||||
SessionContext,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SessionContext tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSessionContext:
|
||||
def test_init_defaults(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
assert ctx.session_id is None
|
||||
assert ctx.service_session_id is None
|
||||
assert ctx.input_messages == []
|
||||
assert ctx.context_messages == {}
|
||||
assert ctx.instructions == []
|
||||
assert ctx.tools == []
|
||||
assert ctx.response is None
|
||||
assert ctx.options == {}
|
||||
assert ctx.metadata == {}
|
||||
|
||||
def test_extend_messages_creates_key(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
msg = ChatMessage(role="user", contents=["hello"])
|
||||
ctx.extend_messages("rag", [msg])
|
||||
assert "rag" in ctx.context_messages
|
||||
assert len(ctx.context_messages["rag"]) == 1
|
||||
assert ctx.context_messages["rag"][0].text == "hello"
|
||||
|
||||
def test_extend_messages_appends_to_existing(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
msg1 = ChatMessage(role="user", contents=["first"])
|
||||
msg2 = ChatMessage(role="user", contents=["second"])
|
||||
ctx.extend_messages("src", [msg1])
|
||||
ctx.extend_messages("src", [msg2])
|
||||
assert len(ctx.context_messages["src"]) == 2
|
||||
|
||||
def test_extend_messages_preserves_source_order(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
ctx.extend_messages("a", [ChatMessage(role="user", contents=["a"])])
|
||||
ctx.extend_messages("b", [ChatMessage(role="user", contents=["b"])])
|
||||
ctx.extend_messages("c", [ChatMessage(role="user", contents=["c"])])
|
||||
assert list(ctx.context_messages.keys()) == ["a", "b", "c"]
|
||||
|
||||
def test_extend_messages_sets_attribution(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
msg = ChatMessage(role="system", contents=["context"])
|
||||
ctx.extend_messages("rag", [msg])
|
||||
stored = ctx.context_messages["rag"][0]
|
||||
assert stored.additional_properties["_attribution"] == {"source_id": "rag"}
|
||||
# Original message is not mutated
|
||||
assert "_attribution" not in msg.additional_properties
|
||||
|
||||
def test_extend_messages_does_not_overwrite_existing_attribution(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
msg = ChatMessage(
|
||||
role="system", contents=["context"], additional_properties={"_attribution": {"source_id": "custom"}}
|
||||
)
|
||||
ctx.extend_messages("rag", [msg])
|
||||
stored = ctx.context_messages["rag"][0]
|
||||
assert stored.additional_properties["_attribution"] == {"source_id": "custom"}
|
||||
|
||||
def test_extend_messages_copies_messages(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
msg = ChatMessage(role="user", contents=["hello"])
|
||||
ctx.extend_messages("src", [msg])
|
||||
stored = ctx.context_messages["src"][0]
|
||||
assert stored is not msg
|
||||
assert stored.text == "hello"
|
||||
# Mutating stored copy does not affect original
|
||||
stored.additional_properties["extra"] = True
|
||||
assert "extra" not in msg.additional_properties
|
||||
|
||||
def test_extend_messages_sender_sets_source_type(self) -> None:
|
||||
class MyProvider:
|
||||
source_id = "rag"
|
||||
|
||||
ctx = SessionContext(input_messages=[])
|
||||
msg = ChatMessage(role="system", contents=["ctx"])
|
||||
ctx.extend_messages(MyProvider(), [msg])
|
||||
stored = ctx.context_messages["rag"][0]
|
||||
assert stored.additional_properties["_attribution"] == {"source_id": "rag", "source_type": "MyProvider"}
|
||||
|
||||
def test_extend_instructions_string(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
ctx.extend_instructions("sys", "Be helpful")
|
||||
assert ctx.instructions == ["Be helpful"]
|
||||
|
||||
def test_extend_instructions_sequence(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
ctx.extend_instructions("sys", ["Be helpful", "Be concise"])
|
||||
assert ctx.instructions == ["Be helpful", "Be concise"]
|
||||
|
||||
def test_get_messages_all(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
ctx.extend_messages("a", [ChatMessage(role="user", contents=["a"])])
|
||||
ctx.extend_messages("b", [ChatMessage(role="user", contents=["b"])])
|
||||
result = ctx.get_messages()
|
||||
assert len(result) == 2
|
||||
assert result[0].text == "a"
|
||||
assert result[1].text == "b"
|
||||
|
||||
def test_get_messages_filter_sources(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
ctx.extend_messages("a", [ChatMessage(role="user", contents=["a"])])
|
||||
ctx.extend_messages("b", [ChatMessage(role="user", contents=["b"])])
|
||||
result = ctx.get_messages(sources=["a"])
|
||||
assert len(result) == 1
|
||||
assert result[0].text == "a"
|
||||
|
||||
def test_get_messages_exclude_sources(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
ctx.extend_messages("a", [ChatMessage(role="user", contents=["a"])])
|
||||
ctx.extend_messages("b", [ChatMessage(role="user", contents=["b"])])
|
||||
result = ctx.get_messages(exclude_sources=["a"])
|
||||
assert len(result) == 1
|
||||
assert result[0].text == "b"
|
||||
|
||||
def test_get_messages_include_input(self) -> None:
|
||||
input_msg = ChatMessage(role="user", contents=["input"])
|
||||
ctx = SessionContext(input_messages=[input_msg])
|
||||
ctx.extend_messages("a", [ChatMessage(role="user", contents=["context"])])
|
||||
result = ctx.get_messages(include_input=True)
|
||||
assert len(result) == 2
|
||||
assert result[1].text == "input"
|
||||
|
||||
def test_get_messages_include_response(self) -> None:
|
||||
from agent_framework import AgentResponse
|
||||
|
||||
ctx = SessionContext(input_messages=[])
|
||||
ctx._response = AgentResponse(messages=[ChatMessage(role="assistant", contents=["reply"])])
|
||||
result = ctx.get_messages(include_response=True)
|
||||
assert len(result) == 1
|
||||
assert result[0].text == "reply"
|
||||
|
||||
def test_response_readonly(self) -> None:
|
||||
ctx = SessionContext(input_messages=[])
|
||||
assert ctx.response is None
|
||||
# Can set via _response internally
|
||||
from agent_framework import AgentResponse
|
||||
|
||||
resp = AgentResponse(messages=[])
|
||||
ctx._response = resp
|
||||
assert ctx.response is resp
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# BaseContextProvider tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestContextProviderBase:
|
||||
def test_source_id_required(self) -> None:
|
||||
provider = BaseContextProvider(source_id="test")
|
||||
assert provider.source_id == "test"
|
||||
|
||||
async def test_before_run_is_noop(self) -> None:
|
||||
provider = BaseContextProvider(source_id="test")
|
||||
session = AgentSession()
|
||||
ctx = SessionContext(input_messages=[])
|
||||
# Should not raise
|
||||
await provider.before_run(agent=None, session=session, context=ctx, state={}) # type: ignore[arg-type]
|
||||
|
||||
async def test_after_run_is_noop(self) -> None:
|
||||
provider = BaseContextProvider(source_id="test")
|
||||
session = AgentSession()
|
||||
ctx = SessionContext(input_messages=[])
|
||||
await provider.after_run(agent=None, session=session, context=ctx, state={}) # type: ignore[arg-type]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# BaseHistoryProvider tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ConcreteHistoryProvider(BaseHistoryProvider):
|
||||
"""Concrete test implementation."""
|
||||
|
||||
def __init__(self, source_id: str, stored_messages: list[ChatMessage] | None = None, **kwargs) -> None:
|
||||
super().__init__(source_id, **kwargs)
|
||||
self.stored: list[ChatMessage] = []
|
||||
self._stored_messages = stored_messages or []
|
||||
|
||||
async def get_messages(self, session_id: str | None, **kwargs) -> list[ChatMessage]:
|
||||
return list(self._stored_messages)
|
||||
|
||||
async def save_messages(self, session_id: str | None, messages: Sequence[ChatMessage], **kwargs) -> None:
|
||||
self.stored.extend(messages)
|
||||
|
||||
|
||||
class TestHistoryProviderBase:
|
||||
def test_default_flags(self) -> None:
|
||||
provider = ConcreteHistoryProvider("mem")
|
||||
assert provider.load_messages is True
|
||||
assert provider.store_outputs is True
|
||||
assert provider.store_inputs is True
|
||||
assert provider.store_context_messages is False
|
||||
assert provider.store_context_from is None
|
||||
|
||||
def test_custom_flags(self) -> None:
|
||||
provider = ConcreteHistoryProvider(
|
||||
"audit",
|
||||
load_messages=False,
|
||||
store_inputs=False,
|
||||
store_context_messages=True,
|
||||
store_context_from={"rag"},
|
||||
)
|
||||
assert provider.load_messages is False
|
||||
assert provider.store_inputs is False
|
||||
assert provider.store_context_messages is True
|
||||
assert provider.store_context_from == {"rag"}
|
||||
|
||||
async def test_before_run_loads_messages(self) -> None:
|
||||
msgs = [ChatMessage(role="user", contents=["history"])]
|
||||
provider = ConcreteHistoryProvider("mem", stored_messages=msgs)
|
||||
session = AgentSession()
|
||||
ctx = SessionContext(session_id="s1", input_messages=[])
|
||||
await provider.before_run(agent=None, session=session, context=ctx, state={}) # type: ignore[arg-type]
|
||||
assert len(ctx.context_messages["mem"]) == 1
|
||||
assert ctx.context_messages["mem"][0].text == "history"
|
||||
|
||||
async def test_after_run_stores_inputs_and_responses(self) -> None:
|
||||
from agent_framework import AgentResponse
|
||||
|
||||
provider = ConcreteHistoryProvider("mem")
|
||||
session = AgentSession()
|
||||
input_msg = ChatMessage(role="user", contents=["hello"])
|
||||
resp_msg = ChatMessage(role="assistant", contents=["hi"])
|
||||
ctx = SessionContext(session_id="s1", input_messages=[input_msg])
|
||||
ctx._response = AgentResponse(messages=[resp_msg])
|
||||
await provider.after_run(agent=None, session=session, context=ctx, state={}) # type: ignore[arg-type]
|
||||
assert len(provider.stored) == 2
|
||||
assert provider.stored[0].text == "hello"
|
||||
assert provider.stored[1].text == "hi"
|
||||
|
||||
async def test_after_run_skips_inputs_when_disabled(self) -> None:
|
||||
from agent_framework import AgentResponse
|
||||
|
||||
provider = ConcreteHistoryProvider("mem", store_inputs=False)
|
||||
ctx = SessionContext(session_id="s1", input_messages=[ChatMessage(role="user", contents=["hello"])])
|
||||
ctx._response = AgentResponse(messages=[ChatMessage(role="assistant", contents=["hi"])])
|
||||
await provider.after_run(agent=None, session=AgentSession(), context=ctx, state={}) # type: ignore[arg-type]
|
||||
assert len(provider.stored) == 1
|
||||
assert provider.stored[0].text == "hi"
|
||||
|
||||
async def test_after_run_skips_responses_when_disabled(self) -> None:
|
||||
from agent_framework import AgentResponse
|
||||
|
||||
provider = ConcreteHistoryProvider("mem", store_outputs=False)
|
||||
ctx = SessionContext(session_id="s1", input_messages=[ChatMessage(role="user", contents=["hello"])])
|
||||
ctx._response = AgentResponse(messages=[ChatMessage(role="assistant", contents=["hi"])])
|
||||
await provider.after_run(agent=None, session=AgentSession(), context=ctx, state={}) # type: ignore[arg-type]
|
||||
assert len(provider.stored) == 1
|
||||
assert provider.stored[0].text == "hello"
|
||||
|
||||
async def test_after_run_stores_context_messages(self) -> None:
|
||||
from agent_framework import AgentResponse
|
||||
|
||||
provider = ConcreteHistoryProvider("audit", load_messages=False, store_context_messages=True)
|
||||
ctx = SessionContext(session_id="s1", input_messages=[ChatMessage(role="user", contents=["hello"])])
|
||||
ctx.extend_messages("rag", [ChatMessage(role="system", contents=["context"])])
|
||||
ctx._response = AgentResponse(messages=[ChatMessage(role="assistant", contents=["hi"])])
|
||||
await provider.after_run(agent=None, session=AgentSession(), context=ctx, state={}) # type: ignore[arg-type]
|
||||
# Should store: context from rag + input + response
|
||||
texts = [m.text for m in provider.stored]
|
||||
assert "context" in texts
|
||||
assert "hello" in texts
|
||||
assert "hi" in texts
|
||||
|
||||
async def test_after_run_stores_context_from_specific_sources(self) -> None:
|
||||
from agent_framework import AgentResponse
|
||||
|
||||
provider = ConcreteHistoryProvider(
|
||||
"audit", load_messages=False, store_context_messages=True, store_context_from={"rag"}
|
||||
)
|
||||
ctx = SessionContext(session_id="s1", input_messages=[])
|
||||
ctx.extend_messages("rag", [ChatMessage(role="system", contents=["rag-context"])])
|
||||
ctx.extend_messages("other", [ChatMessage(role="system", contents=["other-context"])])
|
||||
ctx._response = AgentResponse(messages=[])
|
||||
await provider.after_run(agent=None, session=AgentSession(), context=ctx, state={}) # type: ignore[arg-type]
|
||||
texts = [m.text for m in provider.stored]
|
||||
assert "rag-context" in texts
|
||||
assert "other-context" not in texts
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AgentSession tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAgentSession:
|
||||
def test_auto_generates_session_id(self) -> None:
|
||||
session = AgentSession()
|
||||
assert session.session_id is not None
|
||||
assert len(session.session_id) > 0
|
||||
|
||||
def test_custom_session_id(self) -> None:
|
||||
session = AgentSession(session_id="custom-123")
|
||||
assert session.session_id == "custom-123"
|
||||
|
||||
def test_state_starts_empty(self) -> None:
|
||||
session = AgentSession()
|
||||
assert session.state == {}
|
||||
|
||||
def test_service_session_id(self) -> None:
|
||||
session = AgentSession(service_session_id="svc-456")
|
||||
assert session.service_session_id == "svc-456"
|
||||
|
||||
def test_to_dict(self) -> None:
|
||||
session = AgentSession(session_id="s1", service_session_id="svc1")
|
||||
session.state = {"key": "value"}
|
||||
d = session.to_dict()
|
||||
assert d["type"] == "session"
|
||||
assert d["session_id"] == "s1"
|
||||
assert d["service_session_id"] == "svc1"
|
||||
assert d["state"] == {"key": "value"}
|
||||
|
||||
def test_from_dict(self) -> None:
|
||||
data = {
|
||||
"type": "session",
|
||||
"session_id": "s1",
|
||||
"service_session_id": "svc1",
|
||||
"state": {"key": "value"},
|
||||
}
|
||||
session = AgentSession.from_dict(data)
|
||||
assert session.session_id == "s1"
|
||||
assert session.service_session_id == "svc1"
|
||||
assert session.state == {"key": "value"}
|
||||
|
||||
def test_roundtrip(self) -> None:
|
||||
session = AgentSession(session_id="rt-1")
|
||||
session.state = {"messages": ["a", "b"], "count": 42}
|
||||
json_str = json.dumps(session.to_dict())
|
||||
restored = AgentSession.from_dict(json.loads(json_str))
|
||||
assert restored.session_id == "rt-1"
|
||||
assert restored.state == {"messages": ["a", "b"], "count": 42}
|
||||
|
||||
def test_from_dict_missing_state(self) -> None:
|
||||
data = {"session_id": "s1"}
|
||||
session = AgentSession.from_dict(data)
|
||||
assert session.state == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# InMemoryHistoryProvider tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestInMemoryHistoryProvider:
|
||||
async def test_empty_state_returns_no_messages(self) -> None:
|
||||
provider = InMemoryHistoryProvider("memory")
|
||||
session = AgentSession()
|
||||
ctx = SessionContext(session_id="s1", input_messages=[])
|
||||
await provider.before_run(agent=None, session=session, context=ctx, state=session.state) # type: ignore[arg-type]
|
||||
assert ctx.context_messages.get("memory", []) == []
|
||||
|
||||
async def test_stores_and_loads_messages(self) -> None:
|
||||
from agent_framework import AgentResponse
|
||||
|
||||
provider = InMemoryHistoryProvider("memory")
|
||||
session = AgentSession()
|
||||
|
||||
# First run: send input, get response
|
||||
input_msg = ChatMessage(role="user", contents=["hello"])
|
||||
resp_msg = ChatMessage(role="assistant", contents=["hi there"])
|
||||
ctx1 = SessionContext(session_id="s1", input_messages=[input_msg])
|
||||
await provider.before_run(agent=None, session=session, context=ctx1, state=session.state) # type: ignore[arg-type]
|
||||
ctx1._response = AgentResponse(messages=[resp_msg])
|
||||
await provider.after_run(agent=None, session=session, context=ctx1, state=session.state) # type: ignore[arg-type]
|
||||
|
||||
# Second run: should load previous messages
|
||||
ctx2 = SessionContext(session_id="s1", input_messages=[ChatMessage(role="user", contents=["again"])])
|
||||
await provider.before_run(agent=None, session=session, context=ctx2, state=session.state) # type: ignore[arg-type]
|
||||
loaded = ctx2.context_messages.get("memory", [])
|
||||
assert len(loaded) == 2
|
||||
assert loaded[0].text == "hello"
|
||||
assert loaded[1].text == "hi there"
|
||||
|
||||
async def test_state_is_serializable(self) -> None:
|
||||
from agent_framework import AgentResponse
|
||||
|
||||
provider = InMemoryHistoryProvider("memory")
|
||||
session = AgentSession()
|
||||
|
||||
input_msg = ChatMessage(role="user", contents=["test"])
|
||||
ctx = SessionContext(session_id="s1", input_messages=[input_msg])
|
||||
await provider.before_run(agent=None, session=session, context=ctx, state=session.state) # type: ignore[arg-type]
|
||||
ctx._response = AgentResponse(messages=[ChatMessage(role="assistant", contents=["reply"])])
|
||||
await provider.after_run(agent=None, session=session, context=ctx, state=session.state) # type: ignore[arg-type]
|
||||
|
||||
# State contains ChatMessage objects (not dicts)
|
||||
assert isinstance(session.state["memory"]["messages"][0], ChatMessage)
|
||||
|
||||
# to_dict() serializes them via SerializationProtocol
|
||||
session_dict = session.to_dict()
|
||||
json_str = json.dumps(session_dict)
|
||||
assert json_str # no error
|
||||
|
||||
# Round-trip through session serialization restores ChatMessage objects
|
||||
restored = AgentSession.from_dict(json.loads(json_str))
|
||||
assert isinstance(restored.state["memory"]["messages"][0], ChatMessage)
|
||||
assert restored.state["memory"]["messages"][0].text == "test"
|
||||
assert restored.state["memory"]["messages"][1].text == "reply"
|
||||
|
||||
async def test_source_id_attribution(self) -> None:
|
||||
provider = InMemoryHistoryProvider("custom-source")
|
||||
assert provider.source_id == "custom-source"
|
||||
ctx = SessionContext(session_id="s1", input_messages=[])
|
||||
ctx.extend_messages("custom-source", [ChatMessage(role="user", contents=["test"])])
|
||||
assert "custom-source" in ctx.context_messages
|
||||
Reference in New Issue
Block a user