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

* cleanup of threads and serialization * fix for sliding window * fix redis test * updated from comments * updated context provider and threads * updated lock * add asyncio default * fix redis tests * fix tests * fix tests * renamed to invoking * fixed tests * fix for instructions
2026-06-16 21:04:09 +08:00 · 2025-09-29 18:22:34 +02:00
parent bf5931932e
commit 10d10364a9
52 changed files with 1642 additions and 1411 deletions
@@ -22,7 +22,7 @@ class RedisStoreState(AFBaseModel):


 class RedisChatMessageStore:
-    """Redis-backed implementation of ChatMessageStore using Redis Lists.
+    """Redis-backed implementation of ChatMessageStoreProtocol using Redis Lists.

    This implementation provides persistent, thread-safe chat message storage using Redis Lists.
    Messages are stored as JSON-serialized strings in chronological order, with each conversation
@@ -153,9 +153,9 @@ class RedisChatMessageStore:
            await pipe.execute()

    async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
-        """Add messages to the Redis store (ChatMessageStore protocol method).
+        """Add messages to the Redis store (ChatMessageStoreProtocol protocol method).

-        This method implements the required ChatMessageStore protocol for adding messages.
+        This method implements the required ChatMessageStoreProtocol protocol for adding messages.
        Messages are appended to the Redis list in chronological order, with automatic
        trimming if message limits are configured.

@@ -190,9 +190,9 @@ class RedisChatMessageStore:
                await self._redis_client.ltrim(self.redis_key, -self.max_messages, -1)  # type: ignore[misc]

    async def list_messages(self) -> list[ChatMessage]:
-        """Get all messages from the store in chronological order (ChatMessageStore protocol method).
+        """Get all messages from the store in chronological order (ChatMessageStoreProtocol protocol method).

-        This method implements the required ChatMessageStore protocol for retrieving messages.
+        This method implements the required ChatMessageStoreProtocol protocol for retrieving messages.
        Returns all messages stored in Redis, ordered from oldest (index 0) to newest (index -1).

        Returns:
@@ -220,10 +220,10 @@ class RedisChatMessageStore:

        return messages

-    async def serialize_state(self, **kwargs: Any) -> Any:
-        """Serialize the current store state for persistence (ChatMessageStore protocol method).
+    async def serialize(self, **kwargs: Any) -> Any:
+        """Serialize the current store state for persistence (ChatMessageStoreProtocol protocol method).

-        This method implements the required ChatMessageStore protocol for state serialization.
+        This method implements the required ChatMessageStoreProtocol protocol for state serialization.
        Captures the Redis connection configuration and thread information needed to
        reconstruct the store and reconnect to the same conversation data.

@@ -243,10 +243,43 @@ class RedisChatMessageStore:
        )
        return state.model_dump(**kwargs)

-    async def deserialize_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
-        """Deserialize state data into this store instance (ChatMessageStore protocol method).
+    @classmethod
+    async def deserialize(cls, serialized_store_state: Any, **kwargs: Any) -> RedisChatMessageStore:
+        """Deserialize state data into a new store instance (ChatMessageStoreProtocol protocol method).

-        This method implements the required ChatMessageStore protocol for state deserialization.
+        This method implements the required ChatMessageStoreProtocol protocol for state deserialization.
+        Creates a new RedisChatMessageStore instance from previously serialized data,
+        allowing the store to reconnect to the same conversation data in Redis.
+
+        Args:
+            serialized_store_state: Previously serialized state data from serialize_state().
+                                   Should be a dictionary with thread_id, redis_url, etc.
+            **kwargs: Additional arguments passed to Pydantic model validation.
+
+        Returns:
+            A new RedisChatMessageStore instance configured from the serialized state.
+
+        Raises:
+            ValueError: If required fields are missing or invalid in the serialized state.
+        """
+        if not serialized_store_state:
+            raise ValueError("serialized_store_state is required for deserialization")
+
+        # Validate and parse the serialized state using Pydantic
+        state = RedisStoreState.model_validate(serialized_store_state, **kwargs)
+
+        # Create and return a new store instance with the deserialized configuration
+        return cls(
+            redis_url=state.redis_url,
+            thread_id=state.thread_id,
+            key_prefix=state.key_prefix,
+            max_messages=state.max_messages,
+        )
+
+    async def update_from_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
+        """Deserialize state data into this store instance (ChatMessageStoreProtocol protocol method).
+
+        This method implements the required ChatMessageStoreProtocol protocol for state deserialization.
        Restores the store configuration from previously serialized data, allowing the store
        to reconnect to the same conversation data in Redis.

@@ -1,32 +1,34 @@
 # Copyright (c) Microsoft. All rights reserved.

-from __future__ import annotations
-
+import json
 import sys
 from collections.abc import MutableSequence, Sequence
 from functools import reduce
 from operator import and_
 from typing import Any, Literal, cast

-from agent_framework import ChatMessage, Context, ContextProvider, Role, TextContent
+import numpy as np
+from agent_framework import ChatMessage, Context, ContextProvider, Role
 from agent_framework.exceptions import (
+    AgentException,
    ServiceInitializationError,
    ServiceInvalidRequestError,
 )
+from redisvl.index import AsyncSearchIndex
+from redisvl.query import FilterQuery, HybridQuery, TextQuery
+from redisvl.query.filter import FilterExpression, Tag
+from redisvl.utils.token_escaper import TokenEscaper
+from redisvl.utils.vectorize import BaseVectorizer

 if sys.version_info >= (3, 11):
    from typing import Self  # pragma: no cover
 else:
    from typing_extensions import Self  # pragma: no cover

-import json
-
-import numpy as np
-from redisvl.index import AsyncSearchIndex
-from redisvl.query import FilterQuery, HybridQuery, TextQuery
-from redisvl.query.filter import FilterExpression, Tag
-from redisvl.utils.token_escaper import TokenEscaper
-from redisvl.utils.vectorize import BaseVectorizer
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover


 class RedisProvider(ContextProvider):
@@ -36,41 +38,73 @@ class RedisProvider(ContextProvider):
    Uses full-text or optional hybrid vector search to ground model responses.
    """

-    # Connection and indexing
-    redis_url: str = "redis://localhost:6379"
-    index_name: str = "context"
-    prefix: str = "context"
+    def __init__(
+        self,
+        redis_url: str = "redis://localhost:6379",
+        index_name: str = "context",
+        prefix: str = "context",
+        # Redis vectorizer configuration (optional, injected by client)
+        redis_vectorizer: BaseVectorizer | None = None,
+        vector_field_name: str | None = None,
+        vector_algorithm: Literal["flat", "hnsw"] | None = None,
+        vector_distance_metric: Literal["cosine", "ip", "l2"] | None = None,
+        # Partition fields (indexed for filtering)
+        application_id: str | None = None,
+        agent_id: str | None = None,
+        user_id: str | None = None,
+        thread_id: str | None = None,
+        scope_to_per_operation_thread_id: bool = False,
+        # Prompt and runtime
+        context_prompt: str = ContextProvider.DEFAULT_CONTEXT_PROMPT,
+        redis_index: Any = None,
+        overwrite_index: bool = False,
+    ):
+        """Create a Redis Context Provider.

-    # Redis vectorizer configuration (optional, injected by client)
-    redis_vectorizer: BaseVectorizer | None = None
-    vector_field_name: str | None = None
-    vector_algorithm: Literal["flat", "hnsw"] | None = None
-    vector_distance_metric: Literal["cosine", "ip", "l2"] | None = None
+        Args:
+            redis_url: The Redis server URL.
+            index_name: The name of the Redis index.
+            prefix: The prefix for all keys in the Redis database.
+            redis_vectorizer: The vectorizer to use for Redis.
+            vector_field_name: The name of the vector field in Redis.
+            vector_algorithm: The algorithm to use for vector search.
+            vector_distance_metric: The distance metric to use for vector search.
+            application_id: The application ID to scope the context.
+            agent_id: The agent ID to scope the context.
+            user_id: The user ID to scope the context.
+            thread_id: The thread ID to scope the context.
+            scope_to_per_operation_thread_id: Whether to scope to the per-operation thread ID.
+            context_prompt: The context prompt to use for the provider.
+            redis_index: The Redis index to use for the provider.
+            overwrite_index: Whether to overwrite the existing Redis index.

-    # Partition fields (indexed for filtering)
-    application_id: str | None = None
-    agent_id: str | None = None
-    user_id: str | None = None
-    thread_id: str | None = None
-    scope_to_per_operation_thread_id: bool = False
-
-    # Prompt and runtime
-    context_prompt: str = ContextProvider.DEFAULT_CONTEXT_PROMPT
-    redis_index: Any = None
-    overwrite_index: bool = False
-    _per_operation_thread_id: str | None = None
-    _token_escaper: TokenEscaper = TokenEscaper()
-    _conversation_id: str | None = None
-    _index_initialized: bool = False
-    _schema_dict: dict[str, Any] | None = None
-
-    def model_post_init(self, __context: Any) -> None:
-        """Post-initialization hook to set up computed fields after Pydantic initialization.
-
-        This is called automatically by Pydantic after the model is initialized.
        """
-        # Create Redis index using the cached schema_dict property
-        self.redis_index = AsyncSearchIndex.from_dict(self.schema_dict, redis_url=self.redis_url, validate_on_load=True)
+        self.redis_url = redis_url
+        self.index_name = index_name
+        self.prefix = prefix
+        if redis_vectorizer is not None and not isinstance(redis_vectorizer, BaseVectorizer):
+            raise AgentException(
+                f"The redis vectorizer is not a valid type, got: {type(redis_vectorizer)}, expected: BaseVectorizer."
+            )
+        self.redis_vectorizer = redis_vectorizer
+        self.vector_field_name = vector_field_name
+        self.vector_algorithm: Literal["flat", "hnsw"] | None = vector_algorithm
+        self.vector_distance_metric: Literal["cosine", "ip", "l2"] | None = vector_distance_metric
+        self.application_id = application_id
+        self.agent_id = agent_id
+        self.user_id = user_id
+        self.thread_id = thread_id
+        self.scope_to_per_operation_thread_id = scope_to_per_operation_thread_id
+        self.context_prompt = context_prompt
+        self.overwrite_index = overwrite_index
+        self._per_operation_thread_id: str | None = None
+        self._token_escaper: TokenEscaper = TokenEscaper()
+        self._conversation_id: str | None = None
+        self._index_initialized: bool = False
+        self._schema_dict: dict[str, Any] | None = None
+        self.redis_index = redis_index or AsyncSearchIndex.from_dict(
+            self.schema_dict, redis_url=self.redis_url, validate_on_load=True
+        )

    @property
    def schema_dict(self) -> dict[str, Any]:
@@ -429,6 +463,7 @@ class RedisProvider(ContextProvider):
        """
        return self._per_operation_thread_id if self.scope_to_per_operation_thread_id else self.thread_id

+    @override
    async def thread_created(self, thread_id: str | None) -> None:
        """Called when a new thread is created.

@@ -442,20 +477,27 @@ class RedisProvider(ContextProvider):
        # Track current conversation id (Agent passes conversation_id here)
        self._conversation_id = thread_id or self._conversation_id

-    async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
-        """Called when a new message is being added to the thread.
-
-        Validates scope, normalizes allowed roles, and persists messages to Redis via add().
-
-        Args:
-            thread_id: The ID of the thread or None.
-            new_messages: New messages to add.
-        """
+    @override
+    async def invoked(
+        self,
+        request_messages: ChatMessage | Sequence[ChatMessage],
+        response_messages: ChatMessage | Sequence[ChatMessage] | None = None,
+        invoke_exception: Exception | None = None,
+        **kwargs: Any,
+    ) -> None:
        self._validate_filters()
-        self._validate_per_operation_thread_id(thread_id)
-        self._per_operation_thread_id = self._per_operation_thread_id or thread_id

-        messages_list = [new_messages] if isinstance(new_messages, ChatMessage) else list(new_messages)
+        request_messages_list = (
+            [request_messages] if isinstance(request_messages, ChatMessage) else list(request_messages)
+        )
+        response_messages_list = (
+            [response_messages]
+            if isinstance(response_messages, ChatMessage)
+            else list(response_messages)
+            if response_messages
+            else []
+        )
+        messages_list = [*request_messages_list, *response_messages_list]

        messages: list[dict[str, Any]] = []
        for message in messages_list:
@@ -475,7 +517,8 @@ class RedisProvider(ContextProvider):
        if messages:
            await self._add(data=messages)

-    async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
+    @override
+    async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
        """Called before invoking the model to provide scoped context.

        Concatenates recent messages into a query, fetches matching memories from Redis.
@@ -483,6 +526,7 @@ class RedisProvider(ContextProvider):

        Args:
            messages: List of new messages in the thread.
+            kwargs: not used at present at present.

        Returns:
            Context: Context object containing instructions with memories.
@@ -495,8 +539,12 @@ class RedisProvider(ContextProvider):
        line_separated_memories = "\n".join(
            str(memory.get("content", "")) for memory in memories if memory.get("content")
        )
-        content = TextContent(f"{self.context_prompt}\n{line_separated_memories}") if line_separated_memories else None
-        return Context(contents=[content] if content else None)
+
+        return Context(
+            messages=[ChatMessage(role="user", text=f"{self.context_prompt}\n{line_separated_memories}")]
+            if line_separated_memories
+            else None
+        )

    async def __aenter__(self) -> Self:
        """Async context manager entry.