Split messages into different types (#66)

- Originally messages were just one type - "Messages". But "Messages"
can be used for different purposes (or not used for other purposes).
- We don't want to take internal messages into consideration when
extracting memories. We only want those to be public messages.
- For that, we have split messages into User, Assistant and Internal.
The Message type is now a union of these types.
- The input for creating these was also created to simplify the public
facing APIs (not in this PR).

---------

Co-authored-by: Copilot <[email protected]>

Author: heyitsaamir

packages/memory_module/__init__.py

@@ -1,16 +1,30 @@
 from memory_module.config import LLMConfig, MemoryModuleConfig
 from memory_module.core.memory_module import MemoryModule
 from memory_module.interfaces.types import (
+    AssistantMessage,
+    AssistantMessageInput,
+    InternalMessage,
+    InternalMessageInput,
     Memory,
     Message,
+    MessageInput,
     ShortTermMemoryRetrievalConfig,
+    UserMessage,
+    UserMessageInput,
 )
 
 __all__ = [
     "MemoryModule",
     "MemoryModuleConfig",
     "LLMConfig",
     "Memory",
+    "InternalMessage",
+    "InternalMessageInput",
+    "UserMessageInput",
+    "UserMessage",
     "Message",
+    "MessageInput",
+    "AssistantMessage",
+    "AssistantMessageInput",
     "ShortTermMemoryRetrievalConfig",
 ]

packages/memory_module/core/memory_core.py

@@ -14,6 +14,7 @@
     Memory,
     MemoryType,
     Message,
+    MessageInput,
     ShortTermMemoryRetrievalConfig,
 )
 from memory_module.services.llm_service import LLMService
@@ -94,7 +95,7 @@ def __init__(
             storage: Optional storage implementation for memory persistence
         """
         self.lm = llm_service
-        self.storage = storage or (
+        self.storage: BaseMemoryStorage = storage or (
             SQLiteMemoryStorage(db_path=config.db_path) if config.db_path is not None else InMemoryStorage()
         )
 
@@ -200,8 +201,11 @@ async def _extract_semantic_fact_from_messages(
         for idx, message in enumerate(messages):
             if message.type == "user":
                 messages_str += f"{idx}. User: {message.content}\n"
-            else:
+            elif message.type == "assistant":
                 messages_str += f"{idx}. Assistant: {message.content}\n"
+            else:
+                # we explicitly ignore internal messages
+                continue
 
         system_message = f"""You are a semantic memory management agent. Your goal is to extract meaningful, facts and preferences from user messages. Focus on recognizing general patterns and interests
 that will remain relevant over time, even if the user is mentioning short-term plans or events.
@@ -262,8 +266,8 @@ async def _extract_episodic_memory_from_messages(self, messages: List[Message])
 
         return await self.lm.completion(messages=messages, response_model=EpisodicMemoryExtraction)
 
-    async def add_short_term_memory(self, message: Message) -> None:
-        await self.storage.store_short_term_memory(message)
+    async def add_short_term_memory(self, message: MessageInput) -> Message:
+        return await self.storage.store_short_term_memory(message)
 
     async def retrieve_chat_history(
         self, conversation_ref: str, config: ShortTermMemoryRetrievalConfig

packages/memory_module/core/memory_module.py

@@ -6,7 +6,7 @@
 from memory_module.interfaces.base_memory_core import BaseMemoryCore
 from memory_module.interfaces.base_memory_module import BaseMemoryModule
 from memory_module.interfaces.base_message_queue import BaseMessageQueue
-from memory_module.interfaces.types import Memory, Message, ShortTermMemoryRetrievalConfig
+from memory_module.interfaces.types import Memory, Message, MessageInput, ShortTermMemoryRetrievalConfig
 from memory_module.services.llm_service import LLMService
 
 
@@ -31,13 +31,17 @@ def __init__(
         self.config = config
 
         self.llm_service = llm_service or LLMService(config=config.llm)
-        self.memory_core = memory_core or MemoryCore(config=config, llm_service=self.llm_service)
-        self.message_queue = message_queue or MessageQueue(config=config, memory_core=self.memory_core)
+        self.memory_core: BaseMemoryCore = memory_core or MemoryCore(config=config, llm_service=self.llm_service)
+        self.message_queue: BaseMessageQueue = message_queue or MessageQueue(
+            config=config, memory_core=self.memory_core
+        )
 
-    async def add_message(self, message: Message) -> None:
+    async def add_message(self, message: MessageInput) -> Message:
         """Add a message to be processed into memory."""
-        await self.memory_core.add_short_term_memory(message)
-        await self.message_queue.enqueue(message)
+        message_res = await self.memory_core.add_short_term_memory(message)
+        await self.message_queue.enqueue(message_res)
+
+        return message_res
 
     async def retrieve_memories(self, query: str, user_id: Optional[str], limit: Optional[int]) -> List[Memory]:
         """Retrieve relevant memories based on a query."""

packages/memory_module/interfaces/base_memory_core.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Dict, List, Optional
 
-from memory_module.interfaces.types import Memory, Message, ShortTermMemoryRetrievalConfig
+from memory_module.interfaces.types import Memory, Message, MessageInput, ShortTermMemoryRetrievalConfig
 
 
 class BaseMemoryCore(ABC):
@@ -43,7 +43,7 @@ async def get_messages(self, memory_ids: List[str]) -> Dict[str, List[Message]]:
         pass
 
     @abstractmethod
-    async def add_short_term_memory(self, message: Message) -> None:
+    async def add_short_term_memory(self, message: MessageInput) -> Message:
         """Add a short-term memory entry."""
         pass
 

packages/memory_module/interfaces/base_memory_module.py

@@ -1,14 +1,14 @@
 from abc import ABC, abstractmethod
 from typing import Dict, List, Optional
 
-from memory_module.interfaces.types import Memory, Message, ShortTermMemoryRetrievalConfig
+from memory_module.interfaces.types import Memory, Message, MessageInput, ShortTermMemoryRetrievalConfig
 
 
 class BaseMemoryModule(ABC):
     """Base class for the memory module interface."""
 
     @abstractmethod
-    async def add_message(self, message: Message) -> None:
+    async def add_message(self, message: MessageInput) -> Message:
         """Add a message to be processed into memory."""
         pass
 

packages/memory_module/interfaces/base_memory_storage.py

@@ -1,7 +1,14 @@
 from abc import ABC, abstractmethod
 from typing import Dict, List, Optional
 
-from memory_module.interfaces.types import BaseMemoryInput, EmbedText, Memory, Message, ShortTermMemoryRetrievalConfig
+from memory_module.interfaces.types import (
+    BaseMemoryInput,
+    EmbedText,
+    Memory,
+    Message,
+    MessageInput,
+    ShortTermMemoryRetrievalConfig,
+)
 
 
 class BaseMemoryStorage(ABC):
@@ -30,7 +37,7 @@ async def update_memory(self, memory_id: str, updated_memory: str, *, embedding_
         pass
 
     @abstractmethod
-    async def store_short_term_memory(self, message: Message) -> None:
+    async def store_short_term_memory(self, message: MessageInput) -> Message:
         """Store a short-term memory entry.
 
         Args:

packages/memory_module/interfaces/types.py

@@ -1,7 +1,8 @@
+from abc import ABC
 from datetime import datetime
 from decimal import Decimal
 from enum import Enum
-from typing import List, Literal, Optional
+from typing import ClassVar, List, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
@@ -12,18 +13,78 @@ class User(BaseModel):
     id: str
 
 
-class Message(BaseModel):
-    """Represents a message in a conversation."""
+class BaseMessageInput(ABC, BaseModel):
+    content: str
+    author_id: str
+    conversation_ref: str
+
+
+class InternalMessageInput(BaseMessageInput):
+    """
+    Input parameter for an internal message. Used when creating a new message.
+    """
+
+    model_config = ConfigDict(from_attributes=True)
+    type: ClassVar = "internal"
+    created_at: Optional[datetime] = None
+
+
+class InternalMessage(InternalMessageInput):
+    """
+    Represents a message that is not meant to be shown to the user.
+    Useful for keeping agentic transcript state.
+    These are not used as part of memory extraction
+    """
 
     model_config = ConfigDict(from_attributes=True)
 
     id: str
-    content: str
-    author_id: Optional[str]
-    conversation_ref: str
+    created_at: datetime  # type: ignore Ignoring because this will exist in the concrete class
+
+
+class UserMessageInput(BaseMessageInput):
+    """
+    Input parameter for a user message. Used when creating a new message.
+    """
+
+    model_config = ConfigDict(from_attributes=True)
+    id: str
+    type: ClassVar = "user"
+    deep_link: Optional[str] = None
     created_at: datetime
-    type: Literal["user", "assistant", "internal"] | None = None
+
+
+class UserMessage(UserMessageInput):
+    """
+    Represents a message that was sent by the user.
+    """
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class AssistantMessageInput(BaseMessageInput):
+    """
+    Input parameter for an assistant message. Used when creating a new message.
+    """
+
+    model_config = ConfigDict(from_attributes=True)
+    id: str
+    type: ClassVar = "assistant"
     deep_link: Optional[str] = None
+    created_at: Optional[datetime] = None
+
+
+class AssistantMessage(AssistantMessageInput):
+    """
+    Represents a message that was sent by the assistant.
+    """
+
+    model_config = ConfigDict(from_attributes=True)
+    created_at: datetime  # type: ignore Ignoring because this will exist in the concrete class
+
+
+type MessageInput = InternalMessageInput | UserMessageInput | AssistantMessageInput
+type Message = InternalMessage | UserMessage | AssistantMessage
 
 
 class MemoryAttribution(BaseModel):

packages/memory_module/storage/in_memory_storage.py

@@ -1,6 +1,7 @@
 import datetime
+import uuid
 from collections import defaultdict
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, TypedDict
 
 import numpy as np
 from memory_module.interfaces.base_memory_storage import BaseMemoryStorage
@@ -9,12 +10,32 @@
 )
 from memory_module.interfaces.base_scheduled_events_service import Event
 from memory_module.interfaces.base_scheduled_events_storage import BaseScheduledEventsStorage
-from memory_module.interfaces.types import BaseMemoryInput, EmbedText, Memory, Message, ShortTermMemoryRetrievalConfig
+from memory_module.interfaces.types import (
+    AssistantMessage,
+    AssistantMessageInput,
+    BaseMemoryInput,
+    EmbedText,
+    InternalMessage,
+    InternalMessageInput,
+    Memory,
+    Message,
+    MessageInput,
+    ShortTermMemoryRetrievalConfig,
+    UserMessage,
+    UserMessageInput,
+)
+
+
+class InMemoryInternalStore(TypedDict):
+    memories: Dict[str, Memory]
+    embeddings: Dict[str, List[List[float]]]
+    buffered_messages: Dict[str, List[Message]]
+    scheduled_events: Dict[str, Event]
 
 
 class InMemoryStorage(BaseMemoryStorage, BaseMessageBufferStorage, BaseScheduledEventsStorage):
     def __init__(self):
-        self.storage: Dict = {
+        self.storage: InMemoryInternalStore = {
             "embeddings": {},
             "buffered_messages": defaultdict(list),
             "scheduled_events": {},
@@ -28,7 +49,8 @@ async def store_memory(
         embedding_vectors: List[List[float]],
     ) -> str | None:
         memory_id = str(len(self.storage["memories"]) + 1)
-        self.storage["memories"][memory_id] = memory
+        memory_obj = Memory(**memory.model_dump(), id=memory_id)
+        self.storage["memories"][memory_id] = memory_obj
         self.storage["embeddings"][memory_id] = embedding_vectors
         return memory_id
 
@@ -37,8 +59,48 @@ async def update_memory(self, memory_id: str, updated_memory: str, *, embedding_
             self.storage["memories"][memory_id].content = updated_memory
             self.storage["embeddings"][memory_id] = embedding_vectors
 
-    async def store_short_term_memory(self, message: Message) -> None:
-        await self.store_buffered_message(message)
+    async def store_short_term_memory(self, message: MessageInput) -> Message:
+        if isinstance(message, InternalMessageInput):
+            id = str(uuid.uuid4())
+        else:
+            id = message.id
+
+        created_at = message.created_at or datetime.datetime.now()
+
+        if isinstance(message, InternalMessageInput):
+            deep_link = None
+        else:
+            deep_link = message.deep_link
+
+        if isinstance(message, UserMessageInput):
+            message_obj = UserMessage(
+                id=id,
+                content=message.content,
+                created_at=created_at,
+                conversation_ref=message.conversation_ref,
+                deep_link=deep_link,
+                author_id=message.author_id,
+            )
+        elif isinstance(message, AssistantMessageInput):
+            message_obj = AssistantMessage(
+                id=id,
+                content=message.content,
+                created_at=created_at,
+                conversation_ref=message.conversation_ref,
+                deep_link=deep_link,
+                author_id=message.author_id,
+            )
+        else:
+            message_obj = InternalMessage(
+                id=id,
+                content=message.content,
+                created_at=created_at,
+                conversation_ref=message.conversation_ref,
+                author_id=message.author_id,
+            )
+
+        await self.store_buffered_message(message_obj)
+        return message_obj
 
     async def retrieve_memories(
         self, embedText: EmbedText, user_id: Optional[str], limit: Optional[int] = None
@@ -66,7 +128,7 @@ async def retrieve_memories(
             )
 
         sorted_memories.sort(key=lambda x: x["distance"], reverse=True)
-        return [Memory(id=item["id"], **item["memory"].__dict__) for item in sorted_memories[:limit]]
+        return [Memory(**item["memory"].__dict__) for item in sorted_memories[:limit]]
 
     async def get_memories(self, memory_ids: List[str]) -> List[Memory]:
         return [
@@ -81,7 +143,7 @@ async def get_messages(self, memory_ids: List[str]) -> Dict[str, List[Message]]:
             str_id = memory_id
             if str_id in self.storage["memories"]:
                 memory = self.storage["memories"][str_id]
-                if hasattr(memory, "message_attributions"):
+                if memory.message_attributions:
                     messages = []
                     for msg_id in memory.message_attributions:
                         # Search through buffered messages to find matching message
@@ -104,7 +166,7 @@ async def clear_memories(self, user_id: str) -> None:
             self.storage["embeddings"].pop(memory_id, None)
             self.storage["memories"].pop(memory_id, None)
 
-    async def get_memory(self, memory_id: int) -> Optional[Memory]:
+    async def get_memory(self, memory_id: str) -> Optional[Memory]:
         return self.storage["memories"].get(memory_id)
 
     async def get_all_memories(self, limit: Optional[int] = None) -> List[Memory]: