refactor(conversations): Update conversation store message type handling

- Add support for GeneratedAssistantMessage in conversation store implementations
- Update type hints to include generic type parameter for GeneratedAssistantMessage
- Modify add_message_async method signatures across conversation store classes
- Ensure consistent message type handling across different conversation store implementations
- Improve type flexibility for message storage and retrieval

Author: arthurbrenno

agentle/agents/conversations/callback_conversation_store.py

@@ -46,11 +46,14 @@ async def clear_messages(chat_id: str):
 from __future__ import annotations
 
 from collections.abc import Awaitable, Callable, Sequence
-from typing import override
+from typing import Any, override
 
 from agentle.agents.conversations.conversation_store import ConversationStore
 from agentle.generations.models.messages.assistant_message import AssistantMessage
 from agentle.generations.models.messages.developer_message import DeveloperMessage
+from agentle.generations.models.messages.generated_assistant_message import (
+    GeneratedAssistantMessage,
+)
 from agentle.generations.models.messages.user_message import UserMessage
 
 
@@ -78,10 +81,23 @@ def __init__(
         self,
         get_callback: Callable[
             [str],
-            Awaitable[Sequence[DeveloperMessage | UserMessage | AssistantMessage]],
+            Awaitable[
+                Sequence[
+                    DeveloperMessage
+                    | UserMessage
+                    | AssistantMessage
+                ]
+            ],
         ],
         add_callback: Callable[
-            [str, DeveloperMessage | UserMessage | AssistantMessage], Awaitable[None]
+            [
+                str,
+                DeveloperMessage
+                | UserMessage
+                | AssistantMessage
+                | GeneratedAssistantMessage[Any],
+            ],
+            Awaitable[None],
         ],
         clear_callback: Callable[[str], Awaitable[None]],
         message_limit: int | None = None,
@@ -146,8 +162,13 @@ async def get_conversation_history_async(
         return await self._get_callback(chat_id)
 
     @override
-    async def add_message_async(
-        self, chat_id: str, message: DeveloperMessage | UserMessage | AssistantMessage
+    async def add_message_async[T = Any](
+        self,
+        chat_id: str,
+        message: DeveloperMessage
+        | UserMessage
+        | AssistantMessage
+        | GeneratedAssistantMessage[T],
     ) -> None:
         """
         Add a message to the conversation using the user-provided add callback.

agentle/agents/conversations/conversation_store.py

@@ -1,8 +1,10 @@
 import abc
 from collections.abc import Sequence
+from typing import Any
 
 from agentle.generations.models.messages.assistant_message import AssistantMessage
 from agentle.generations.models.messages.developer_message import DeveloperMessage
+from agentle.generations.models.messages.generated_assistant_message import GeneratedAssistantMessage
 from agentle.generations.models.messages.user_message import UserMessage
 
 
@@ -42,8 +44,8 @@ async def get_conversation_history_async(
     ) -> Sequence[DeveloperMessage | UserMessage | AssistantMessage]: ...
 
     @abc.abstractmethod
-    async def add_message_async(
-        self, chat_id: str, message: DeveloperMessage | UserMessage | AssistantMessage
+    async def add_message_async[T = Any](
+        self, chat_id: str, message: DeveloperMessage | UserMessage | AssistantMessage | GeneratedAssistantMessage[T]
     ) -> None: ...
 
     @abc.abstractmethod

agentle/agents/conversations/firebase_conversation_store.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, override
+from typing import TYPE_CHECKING, Any, override
 
 from agentle.agents.conversations.conversation_store import ConversationStore
 from agentle.generations.models.messages.assistant_message import AssistantMessage
 from agentle.generations.models.messages.developer_message import DeveloperMessage
+from agentle.generations.models.messages.generated_assistant_message import (
+    GeneratedAssistantMessage,
+)
 from agentle.generations.models.messages.user_message import UserMessage
 
 if TYPE_CHECKING:
@@ -28,8 +31,8 @@ def __init__(
         self._collection_name = collection_name
 
     @override
-    async def add_message_async(
-        self, chat_id: str, message: DeveloperMessage | UserMessage | AssistantMessage
+    async def add_message_async[T = Any](
+        self, chat_id: str, message: DeveloperMessage | UserMessage | AssistantMessage | GeneratedAssistantMessage[T]
     ) -> None:
         from google.cloud import firestore
 

agentle/agents/conversations/json_file_conversation_store.py

@@ -7,6 +7,9 @@
 from agentle.agents.conversations.conversation_store import ConversationStore
 from agentle.generations.models.messages.assistant_message import AssistantMessage
 from agentle.generations.models.messages.developer_message import DeveloperMessage
+from agentle.generations.models.messages.generated_assistant_message import (
+    GeneratedAssistantMessage,
+)
 from agentle.generations.models.messages.user_message import UserMessage
 
 
@@ -87,8 +90,12 @@ def _json_serializer(self, obj: Any) -> Any:
         # For other non-serializable objects, convert to string
         return str(obj)
 
-    def _message_to_dict(
-        self, message: DeveloperMessage | UserMessage | AssistantMessage
+    def _message_to_dict[T](
+        self,
+        message: DeveloperMessage
+        | UserMessage
+        | AssistantMessage
+        | GeneratedAssistantMessage[T],
     ) -> dict[str, Any]:
         """Convert a message object to a dictionary for JSON serialization."""
         # Get the basic message dictionary
@@ -177,8 +184,13 @@ def _dict_to_message(
             return UserMessage.model_validate(message_data)
 
     @override
-    async def add_message_async(
-        self, chat_id: str, message: DeveloperMessage | UserMessage | AssistantMessage
+    async def add_message_async[T = Any](
+        self,
+        chat_id: str,
+        message: DeveloperMessage
+        | UserMessage
+        | AssistantMessage
+        | GeneratedAssistantMessage[T],
     ) -> None:
         """Add a message to the conversation."""
         messages_data = self._load_messages(chat_id)

agentle/agents/conversations/local_conversation_store.py

@@ -1,14 +1,19 @@
 from collections.abc import MutableMapping, MutableSequence, Sequence
-from typing import override
+from typing import Any, override
+
 from agentle.agents.conversations.conversation_store import ConversationStore
 from agentle.generations.models.messages.assistant_message import AssistantMessage
 from agentle.generations.models.messages.developer_message import DeveloperMessage
+from agentle.generations.models.messages.generated_assistant_message import (
+    GeneratedAssistantMessage,
+)
 from agentle.generations.models.messages.user_message import UserMessage
 
 
 class LocalConversationStore(ConversationStore):
     __messages: MutableMapping[
-        str, MutableSequence[DeveloperMessage | UserMessage | AssistantMessage]
+        str,
+        MutableSequence[DeveloperMessage | UserMessage | AssistantMessage],
     ]
 
     def __init__(
@@ -20,8 +25,13 @@ def __init__(
         self.__messages = {}
 
     @override
-    async def add_message_async(
-        self, chat_id: str, message: DeveloperMessage | UserMessage | AssistantMessage
+    async def add_message_async[T = Any](
+        self,
+        chat_id: str,
+        message: DeveloperMessage
+        | UserMessage
+        | AssistantMessage
+        | GeneratedAssistantMessage[T],
     ) -> None:
         if chat_id not in self.__messages:
             self.__messages[chat_id] = []
@@ -41,6 +51,9 @@ async def add_message_async(
                     # Don't add message if limit reached and not overriding
                     return
 
+        if isinstance(message, GeneratedAssistantMessage):
+            message = message.to_assistant_message()
+
         self.__messages[chat_id].append(message)
 
     @override

agentle/agents/conversations/mysql_conversation_store.py

@@ -8,6 +8,9 @@
 from agentle.agents.conversations.conversation_store import ConversationStore
 from agentle.generations.models.messages.assistant_message import AssistantMessage
 from agentle.generations.models.messages.developer_message import DeveloperMessage
+from agentle.generations.models.messages.generated_assistant_message import (
+    GeneratedAssistantMessage,
+)
 from agentle.generations.models.messages.user_message import UserMessage
 
 if TYPE_CHECKING:
@@ -74,7 +77,11 @@ async def _ensure_table_exists(self) -> None:
                 await conn.commit()
 
     def _message_to_dict(
-        self, message: DeveloperMessage | UserMessage | AssistantMessage
+        self,
+        message: DeveloperMessage
+        | UserMessage
+        | AssistantMessage
+        | GeneratedAssistantMessage[Any],
     ) -> dict[str, Any]:
         """
         Convert a Message object to a dictionary for JSON serialization.
@@ -125,8 +132,13 @@ def _dict_to_message(
             return UserMessage.model_validate(message_dict)
 
     @override
-    async def add_message_async(
-        self, chat_id: str, message: DeveloperMessage | UserMessage | AssistantMessage
+    async def add_message_async[T = Any](
+        self,
+        chat_id: str,
+        message: DeveloperMessage
+        | UserMessage
+        | AssistantMessage
+        | GeneratedAssistantMessage[T],
     ) -> None:
         """
         Add a message to the conversation store.

agentle/agents/conversations/postgres_conversation_store.py

@@ -7,6 +7,9 @@
 from agentle.agents.conversations.conversation_store import ConversationStore
 from agentle.generations.models.messages.assistant_message import AssistantMessage
 from agentle.generations.models.messages.developer_message import DeveloperMessage
+from agentle.generations.models.messages.generated_assistant_message import (
+    GeneratedAssistantMessage,
+)
 from agentle.generations.models.messages.user_message import UserMessage
 
 if TYPE_CHECKING:
@@ -84,7 +87,11 @@ async def ensure_table_exists(self) -> None:
             """)
 
     def _message_to_dict(
-        self, message: DeveloperMessage | UserMessage | AssistantMessage
+        self,
+        message: DeveloperMessage
+        | UserMessage
+        | AssistantMessage
+        | GeneratedAssistantMessage[Any],
     ) -> dict[str, Any]:
         """Convert a message object to a dictionary for JSON serialization."""
         message_dict = message.model_dump()
@@ -119,8 +126,13 @@ def _dict_to_message(
             return UserMessage.model_validate(message_dict)
 
     @override
-    async def add_message_async(
-        self, chat_id: str, message: DeveloperMessage | UserMessage | AssistantMessage
+    async def add_message_async[T = Any](
+        self,
+        chat_id: str,
+        message: DeveloperMessage
+        | UserMessage
+        | AssistantMessage
+        | GeneratedAssistantMessage[T],
     ) -> None:
         """Add a message to the conversation."""
         async with self._pool.acquire() as conn: