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,522 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unified context management types for the agent framework.
|
||||
|
||||
This module provides the core types for the context provider pipeline:
|
||||
- SessionContext: Per-invocation state passed through providers
|
||||
- BaseContextProvider: Base class for context providers (renamed to ContextProvider in PR2)
|
||||
- BaseHistoryProvider: Base class for history storage providers (renamed to HistoryProvider in PR2)
|
||||
- AgentSession: Lightweight session state container
|
||||
- InMemoryHistoryProvider: Built-in in-memory history provider
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import uuid
|
||||
from abc import abstractmethod
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ._tools import ToolProtocol
|
||||
from ._types import AgentResponse, ChatMessage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._agents import SupportsAgentRun
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AgentSession",
|
||||
"BaseContextProvider",
|
||||
"BaseHistoryProvider",
|
||||
"InMemoryHistoryProvider",
|
||||
"SessionContext",
|
||||
]
|
||||
|
||||
|
||||
# Registry of known types for state deserialization
|
||||
_STATE_TYPE_REGISTRY: dict[str, type] = {}
|
||||
|
||||
|
||||
def _register_state_type(cls: type) -> None:
|
||||
"""Register a type for automatic deserialization in session state."""
|
||||
type_id: str = getattr(cls, "_get_type_identifier", lambda: cls.__name__.lower())()
|
||||
_STATE_TYPE_REGISTRY[type_id] = cls
|
||||
|
||||
|
||||
def _serialize_value(value: Any) -> Any:
|
||||
"""Serialize a single value, handling objects with to_dict()."""
|
||||
if hasattr(value, "to_dict") and callable(value.to_dict):
|
||||
return value.to_dict() # pyright: ignore[reportUnknownMemberType]
|
||||
if isinstance(value, list):
|
||||
return [_serialize_value(item) for item in value] # pyright: ignore[reportUnknownVariableType]
|
||||
if isinstance(value, dict):
|
||||
return {str(k): _serialize_value(v) for k, v in value.items()} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]
|
||||
return value
|
||||
|
||||
|
||||
def _deserialize_value(value: Any) -> Any:
|
||||
"""Deserialize a single value, restoring registered types."""
|
||||
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 isinstance(value, list):
|
||||
return [_deserialize_value(item) for item in value] # pyright: ignore[reportUnknownVariableType]
|
||||
if isinstance(value, dict):
|
||||
return {str(k): _deserialize_value(v) for k, v in value.items()} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]
|
||||
return value
|
||||
|
||||
|
||||
def _serialize_state(state: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Deep-serialize a state dict, converting SerializationProtocol objects to dicts."""
|
||||
return {k: _serialize_value(v) for k, v in state.items()}
|
||||
|
||||
|
||||
def _deserialize_state(state: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Deep-deserialize a state dict, restoring SerializationProtocol objects."""
|
||||
return {k: _deserialize_value(v) for k, v in state.items()}
|
||||
|
||||
|
||||
# Register known types
|
||||
_register_state_type(ChatMessage)
|
||||
|
||||
|
||||
class SessionContext:
|
||||
"""Per-invocation state passed through the context provider pipeline.
|
||||
|
||||
Created fresh for each agent.run() call. Providers read from and write to
|
||||
the mutable fields to add context before invocation and process responses after.
|
||||
|
||||
Attributes:
|
||||
session_id: The ID of the current session.
|
||||
service_session_id: Service-managed session ID (if present, service handles storage).
|
||||
input_messages: The new messages being sent to the agent (set by caller).
|
||||
context_messages: Dict mapping source_id -> messages added by that provider.
|
||||
Maintains insertion order (provider execution order).
|
||||
instructions: Additional instructions added by providers.
|
||||
tools: Additional tools added by providers.
|
||||
response: After invocation, contains the full AgentResponse, should not be changed.
|
||||
options: Options passed to agent.run() - read-only, for reflection only.
|
||||
metadata: Shared metadata dictionary for cross-provider communication.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session_id: str | None = None,
|
||||
service_session_id: str | None = None,
|
||||
input_messages: list[ChatMessage],
|
||||
context_messages: dict[str, list[ChatMessage]] | None = None,
|
||||
instructions: list[str] | None = None,
|
||||
tools: list[ToolProtocol] | None = None,
|
||||
options: dict[str, Any] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
):
|
||||
"""Initialize the session context.
|
||||
|
||||
Args:
|
||||
session_id: The ID of the current session.
|
||||
service_session_id: Service-managed session ID.
|
||||
input_messages: The new messages being sent to the agent.
|
||||
context_messages: Pre-populated context messages by source.
|
||||
instructions: Pre-populated instructions.
|
||||
tools: Pre-populated tools.
|
||||
options: Options from agent.run() - read-only for providers.
|
||||
metadata: Shared metadata for cross-provider communication.
|
||||
"""
|
||||
self.session_id = session_id
|
||||
self.service_session_id = service_session_id
|
||||
self.input_messages = input_messages
|
||||
self.context_messages: dict[str, list[ChatMessage]] = context_messages or {}
|
||||
self.instructions: list[str] = instructions or []
|
||||
self.tools: list[ToolProtocol] = tools or []
|
||||
self._response: AgentResponse | None = None
|
||||
self.options: dict[str, Any] = options or {}
|
||||
self.metadata: dict[str, Any] = metadata or {}
|
||||
|
||||
@property
|
||||
def response(self) -> AgentResponse | None:
|
||||
"""The agent's response. Set by the framework after invocation, read-only for providers."""
|
||||
return self._response
|
||||
|
||||
def extend_messages(self, source: str | object, messages: Sequence[ChatMessage]) -> None:
|
||||
"""Add context messages from a specific source.
|
||||
|
||||
Messages are copied before attribution is added, so the caller's
|
||||
original message objects are never mutated. The copies are stored
|
||||
keyed by source_id, maintaining insertion order based on provider
|
||||
execution order. Each message gets an ``attribution`` marker in
|
||||
``additional_properties`` for downstream filtering.
|
||||
|
||||
Args:
|
||||
source: Either a plain ``source_id`` string, or an object with a
|
||||
``source_id`` attribute (e.g. a context provider). When an
|
||||
object is passed, its class name is recorded as
|
||||
``source_type`` in the attribution.
|
||||
messages: The messages to add.
|
||||
"""
|
||||
if isinstance(source, str):
|
||||
source_id = source
|
||||
attribution: dict[str, str] = {"source_id": source_id}
|
||||
else:
|
||||
source_id = source.source_id # type: ignore[attr-defined]
|
||||
attribution = {"source_id": source_id, "source_type": type(source).__name__}
|
||||
|
||||
copied: list[ChatMessage] = []
|
||||
for message in messages:
|
||||
msg_copy = copy.copy(message)
|
||||
msg_copy.additional_properties = dict(message.additional_properties)
|
||||
msg_copy.additional_properties.setdefault("_attribution", attribution)
|
||||
copied.append(msg_copy)
|
||||
if source_id not in self.context_messages:
|
||||
self.context_messages[source_id] = []
|
||||
self.context_messages[source_id].extend(copied)
|
||||
|
||||
def extend_instructions(self, source_id: str, instructions: str | Sequence[str]) -> None:
|
||||
"""Add instructions to be prepended to the conversation.
|
||||
|
||||
Args:
|
||||
source_id: The provider source_id adding these instructions.
|
||||
instructions: A single instruction string or sequence of strings.
|
||||
"""
|
||||
if isinstance(instructions, str):
|
||||
instructions = [instructions]
|
||||
self.instructions.extend(instructions)
|
||||
|
||||
def extend_tools(self, source_id: str, tools: Sequence[ToolProtocol]) -> None:
|
||||
"""Add tools to be available for this invocation.
|
||||
|
||||
Tools are added with source attribution in their metadata.
|
||||
|
||||
Args:
|
||||
source_id: The provider source_id adding these tools.
|
||||
tools: The tools to add.
|
||||
"""
|
||||
for tool in tools:
|
||||
if hasattr(tool, "additional_properties") and isinstance(tool.additional_properties, dict):
|
||||
tool.additional_properties["context_source"] = source_id
|
||||
self.tools.extend(tools)
|
||||
|
||||
def get_messages(
|
||||
self,
|
||||
*,
|
||||
sources: set[str] | None = None,
|
||||
exclude_sources: set[str] | None = None,
|
||||
include_input: bool = False,
|
||||
include_response: bool = False,
|
||||
) -> list[ChatMessage]:
|
||||
"""Get context messages, optionally filtered and including input/response.
|
||||
|
||||
Returns messages in provider execution order (dict insertion order),
|
||||
with input and response appended if requested.
|
||||
|
||||
Args:
|
||||
sources: If provided, only include context messages from these sources.
|
||||
exclude_sources: If provided, exclude context messages from these sources.
|
||||
include_input: If True, append input_messages after context.
|
||||
include_response: If True, append response.messages at the end.
|
||||
|
||||
Returns:
|
||||
Flattened list of messages in conversation order.
|
||||
"""
|
||||
result: list[ChatMessage] = []
|
||||
for source_id, messages in self.context_messages.items():
|
||||
if sources is not None and source_id not in sources:
|
||||
continue
|
||||
if exclude_sources is not None and source_id in exclude_sources:
|
||||
continue
|
||||
result.extend(messages)
|
||||
if include_input and self.input_messages:
|
||||
result.extend(self.input_messages)
|
||||
if include_response and self.response and self.response.messages:
|
||||
result.extend(self.response.messages)
|
||||
return result
|
||||
|
||||
|
||||
class BaseContextProvider:
|
||||
"""Base class for context providers (hooks pattern).
|
||||
|
||||
Context providers participate in the context engineering pipeline,
|
||||
adding context before model invocation and processing responses after.
|
||||
|
||||
Note:
|
||||
This class uses a temporary name prefixed with ``_`` to avoid collision
|
||||
with the existing ``ContextProvider`` in ``_memory.py``. It will be
|
||||
renamed to ``ContextProvider`` in PR2 when the old class is removed.
|
||||
|
||||
Attributes:
|
||||
source_id: Unique identifier for this provider instance (required).
|
||||
Used for message/tool attribution so other providers can filter.
|
||||
"""
|
||||
|
||||
def __init__(self, source_id: str):
|
||||
"""Initialize the provider.
|
||||
|
||||
Args:
|
||||
source_id: Unique identifier for this provider instance.
|
||||
"""
|
||||
self.source_id = source_id
|
||||
|
||||
async def before_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called before model invocation.
|
||||
|
||||
Override to add context (messages, instructions, tools) to the
|
||||
SessionContext before the model is invoked.
|
||||
|
||||
Args:
|
||||
agent: The agent running this invocation.
|
||||
session: The current session.
|
||||
context: The invocation context - add messages/instructions/tools here.
|
||||
state: The session's mutable state dict.
|
||||
"""
|
||||
|
||||
async def after_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after model invocation.
|
||||
|
||||
Override to process the response (store messages, extract info, etc.).
|
||||
The context.response will be populated at this point.
|
||||
|
||||
Args:
|
||||
agent: The agent that ran this invocation.
|
||||
session: The current session.
|
||||
context: The invocation context with response populated.
|
||||
state: The session's mutable state dict.
|
||||
"""
|
||||
|
||||
|
||||
class BaseHistoryProvider(BaseContextProvider):
|
||||
"""Base class for conversation history storage providers.
|
||||
|
||||
A single class configurable for different use cases:
|
||||
- Primary memory storage (loads + stores messages)
|
||||
- Audit/logging storage (stores only, doesn't load)
|
||||
- Evaluation storage (stores only for later analysis)
|
||||
|
||||
Note:
|
||||
This class uses a temporary name prefixed with ``_`` to avoid collision
|
||||
with existing types. It will be renamed to ``HistoryProvider`` in PR2.
|
||||
|
||||
Subclasses only need to implement ``get_messages()`` and ``save_messages()``.
|
||||
The default ``before_run``/``after_run`` handle loading and storing based on
|
||||
configuration flags. Override them for custom behavior.
|
||||
|
||||
Attributes:
|
||||
load_messages: Whether to load messages before invocation (default True).
|
||||
When False, the agent skips calling ``before_run`` entirely.
|
||||
store_inputs: Whether to store input messages (default True).
|
||||
store_context_messages: Whether to store context from other providers (default False).
|
||||
store_context_from: If set, only store context from these source_ids.
|
||||
store_outputs: Whether to store response messages (default True).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str,
|
||||
*,
|
||||
load_messages: bool = True,
|
||||
store_inputs: bool = True,
|
||||
store_context_messages: bool = False,
|
||||
store_context_from: set[str] | None = None,
|
||||
store_outputs: bool = True,
|
||||
):
|
||||
"""Initialize the history provider.
|
||||
|
||||
Args:
|
||||
source_id: Unique identifier for this provider instance.
|
||||
load_messages: Whether to load messages before invocation.
|
||||
store_inputs: Whether to store input messages.
|
||||
store_context_messages: Whether to store context from other providers.
|
||||
store_context_from: If set, only store context from these source_ids.
|
||||
store_outputs: Whether to store response messages.
|
||||
"""
|
||||
super().__init__(source_id)
|
||||
self.load_messages = load_messages
|
||||
self.store_inputs = store_inputs
|
||||
self.store_context_messages = store_context_messages
|
||||
self.store_context_from = store_context_from
|
||||
self.store_outputs = store_outputs
|
||||
|
||||
@abstractmethod
|
||||
async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[ChatMessage]:
|
||||
"""Retrieve stored messages for this session.
|
||||
|
||||
Args:
|
||||
session_id: The session ID to retrieve messages for.
|
||||
**kwargs: Additional arguments (e.g., ``state`` for in-memory providers).
|
||||
|
||||
Returns:
|
||||
List of stored messages.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def save_messages(self, session_id: str | None, messages: Sequence[ChatMessage], **kwargs: Any) -> None:
|
||||
"""Persist messages for this session.
|
||||
|
||||
Args:
|
||||
session_id: The session ID to store messages for.
|
||||
messages: The messages to persist.
|
||||
**kwargs: Additional arguments (e.g., ``state`` for in-memory providers).
|
||||
"""
|
||||
...
|
||||
|
||||
def _get_context_messages_to_store(self, context: SessionContext) -> list[ChatMessage]:
|
||||
"""Get context messages that should be stored based on configuration."""
|
||||
if not self.store_context_messages:
|
||||
return []
|
||||
if self.store_context_from is not None:
|
||||
return context.get_messages(sources=self.store_context_from)
|
||||
return context.get_messages(exclude_sources={self.source_id})
|
||||
|
||||
async def before_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Load history into context. Skipped by the agent when load_messages=False."""
|
||||
history = await self.get_messages(context.session_id, state=state)
|
||||
context.extend_messages(self, history)
|
||||
|
||||
async def after_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Store messages based on configuration."""
|
||||
messages_to_store: list[ChatMessage] = []
|
||||
messages_to_store.extend(self._get_context_messages_to_store(context))
|
||||
if self.store_inputs:
|
||||
messages_to_store.extend(context.input_messages)
|
||||
if self.store_outputs and context.response and context.response.messages:
|
||||
messages_to_store.extend(context.response.messages)
|
||||
if messages_to_store:
|
||||
await self.save_messages(context.session_id, messages_to_store, state=state)
|
||||
|
||||
|
||||
class AgentSession:
|
||||
"""A conversation session with an agent.
|
||||
|
||||
Lightweight state container. Provider instances are owned by the agent,
|
||||
not the session. The session only holds session IDs and a mutable state dict.
|
||||
|
||||
Attributes:
|
||||
session_id: Unique identifier for this session.
|
||||
service_session_id: Service-managed session ID (if using service-side storage).
|
||||
state: Mutable state dict shared with all providers.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session_id: str | None = None,
|
||||
service_session_id: str | None = None,
|
||||
):
|
||||
"""Initialize the session.
|
||||
|
||||
Args:
|
||||
session_id: Optional session ID (generated if not provided).
|
||||
service_session_id: Optional service-managed session ID.
|
||||
"""
|
||||
self._session_id = session_id or str(uuid.uuid4())
|
||||
self.service_session_id = service_session_id
|
||||
self.state: dict[str, Any] = {}
|
||||
|
||||
@property
|
||||
def session_id(self) -> str:
|
||||
"""The unique identifier for this session."""
|
||||
return self._session_id
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Serialize session to a plain dict for storage/transfer.
|
||||
|
||||
Values in ``state`` that implement ``SerializationProtocol`` (i.e. have
|
||||
``to_dict``/``from_dict``) are serialized automatically. Built-in types
|
||||
(str, int, float, bool, None, list, dict) are kept as-is.
|
||||
"""
|
||||
return {
|
||||
"type": "session",
|
||||
"session_id": self._session_id,
|
||||
"service_session_id": self.service_session_id,
|
||||
"state": _serialize_state(self.state),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> AgentSession:
|
||||
"""Restore session from a previously serialized dict.
|
||||
|
||||
Values in ``state`` that were serialized via ``SerializationProtocol``
|
||||
(containing a ``type`` key) are restored to their original types.
|
||||
|
||||
Args:
|
||||
data: Dict from a previous ``to_dict()`` call.
|
||||
|
||||
Returns:
|
||||
Restored AgentSession instance.
|
||||
"""
|
||||
session = cls(
|
||||
session_id=data["session_id"],
|
||||
service_session_id=data.get("service_session_id"),
|
||||
)
|
||||
session.state = _deserialize_state(data.get("state", {}))
|
||||
return session
|
||||
|
||||
|
||||
class InMemoryHistoryProvider(BaseHistoryProvider):
|
||||
"""Built-in history provider that stores messages in session.state.
|
||||
|
||||
Messages are stored in ``state[source_id]["messages"]`` as a list of
|
||||
``ChatMessage`` objects. Serialization to/from dicts is handled by
|
||||
``AgentSession.to_dict()``/``from_dict()`` using ``SerializationProtocol``.
|
||||
|
||||
This provider holds no instance state — all data lives in the session's
|
||||
state dict, passed as a named ``state`` parameter to ``get_messages``/``save_messages``.
|
||||
|
||||
This is the default provider auto-added by the agent when no providers
|
||||
are configured and ``conversation_id`` or ``store=True`` is set.
|
||||
"""
|
||||
|
||||
async def get_messages(
|
||||
self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any
|
||||
) -> list[ChatMessage]:
|
||||
"""Retrieve messages from session state."""
|
||||
if state is None:
|
||||
return []
|
||||
my_state = state.get(self.source_id, {})
|
||||
return list(my_state.get("messages", []))
|
||||
|
||||
async def save_messages(
|
||||
self,
|
||||
session_id: str | None,
|
||||
messages: Sequence[ChatMessage],
|
||||
*,
|
||||
state: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Persist messages to session state."""
|
||||
if state is None:
|
||||
return
|
||||
my_state = state.setdefault(self.source_id, {})
|
||||
existing = my_state.get("messages", [])
|
||||
my_state["messages"] = [*existing, *messages]
|
||||
@@ -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