Add memory module middleware (#69)

Author: heyitsaamir

packages/memory_module/__init__.py

@@ -12,6 +12,7 @@
     UserMessage,
     UserMessageInput,
 )
+from memory_module.utils.teams_bot_middlware import MemoryMiddleware
 
 __all__ = [
     "MemoryModule",
@@ -27,4 +28,5 @@
     "AssistantMessage",
     "AssistantMessageInput",
     "ShortTermMemoryRetrievalConfig",
+    "MemoryMiddleware",
 ]

packages/memory_module/pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "numpy",
     "sqlite-vec>=0.1.6",
     "litellm==1.54.1",
+    "botbuilder>=0.0.1",
 ]
 
 [tool.uv]

packages/memory_module/storage/sqlite_memory_storage.py

@@ -47,7 +47,7 @@ async def store_memory(self, memory: BaseMemoryInput, *, embedding_vectors: List
                 (
                     memory_id,
                     memory.content,
-                    memory.created_at,
+                    memory.created_at.isoformat(),
                     memory.user_id,
                     memory.memory_type.value,
                 ),
@@ -300,7 +300,7 @@ async def store_short_term_memory(self, message: MessageInput) -> Message:
                 message.content,
                 message.author_id,
                 message.conversation_ref,
-                created_at,
+                created_at.isoformat(),
                 message.type,
                 deep_link,
             ),

packages/memory_module/storage/sqlite_message_buffer_storage.py

@@ -38,7 +38,7 @@ async def store_buffered_message(self, message: Message) -> None:
             (
                 message.id,
                 message.conversation_ref,
-                message.created_at,
+                message.created_at.isoformat(),
             ),
         )
 

packages/memory_module/utils/teams_bot_middlware.py

@@ -0,0 +1,114 @@
+import datetime
+from asyncio import gather
+from typing import Awaitable, Callable, List
+
+from botbuilder.core import TurnContext
+from botbuilder.core.middleware_set import Middleware
+from botbuilder.schema import Activity, ResourceResponse
+from memory_module.interfaces.base_memory_module import BaseMemoryModule
+from memory_module.interfaces.types import (
+    AssistantMessageInput,
+    UserMessageInput,
+)
+
+
+def build_deep_link(context: TurnContext, message_id: str):
+    conversation_ref = TurnContext.get_conversation_reference(context.activity)
+    if conversation_ref.conversation and conversation_ref.conversation.is_group:
+        deeplink_conversation_id = conversation_ref.conversation.id
+    elif conversation_ref.user and conversation_ref.bot:
+        user_aad_object_id = conversation_ref.user.aad_object_id
+        bot_id = conversation_ref.bot.id.replace("28:", "")
+        deeplink_conversation_id = f"19:{user_aad_object_id}_{bot_id}@unq.gbl.spaces"
+    else:
+        return None
+    return f"https://teams.microsoft.com/l/message/{deeplink_conversation_id}/{message_id}?context=%7B%22contextType%22%3A%22chat%22%7D"
+
+
+class MemoryMiddleware(Middleware):
+    def __init__(self, memory_module: BaseMemoryModule):
+        self.memory_module = memory_module
+
+    async def add_user_message(self, context: TurnContext):
+        conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
+        content = context.activity.text
+        if not content:
+            print("content is not text, so ignoring...")
+            return False
+        if conversation_ref_dict is None:
+            print("conversation_ref_dict is None")
+            return False
+        if conversation_ref_dict.user is None:
+            print("conversation_ref_dict.user is None")
+            return False
+        if conversation_ref_dict.conversation is None:
+            print("conversation_ref_dict.conversation is None")
+            return False
+        user_aad_object_id = conversation_ref_dict.user.aad_object_id
+        message_id = context.activity.id
+        await self.memory_module.add_message(
+            UserMessageInput(
+                id=message_id,
+                content=context.activity.text,
+                author_id=user_aad_object_id,
+                conversation_ref=conversation_ref_dict.conversation.id,
+                created_at=context.activity.timestamp if context.activity.timestamp else datetime.datetime.now(),
+                deep_link=build_deep_link(context, context.activity.id),
+            )
+        )
+        return True
+
+    async def add_agent_message(
+        self, context: TurnContext, activities: List[Activity], responses: List[ResourceResponse]
+    ):
+        conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
+        if conversation_ref_dict is None:
+            print("conversation_ref_dict is None")
+            return False
+        if conversation_ref_dict.bot is None:
+            print("conversation_ref_dict.bot is None")
+            return False
+        if conversation_ref_dict.conversation is None:
+            print("conversation_ref_dict.conversation is None")
+            return False
+
+        tasks = []
+        for activity, response in zip(activities, responses, strict=False):
+            if activity.text:
+                tasks.append(
+                    self.memory_module.add_message(
+                        AssistantMessageInput(
+                            id=response.id,
+                            content=activity.text,
+                            author_id=conversation_ref_dict.bot.id,
+                            conversation_ref=conversation_ref_dict.conversation.id,
+                            deep_link=build_deep_link(context, response.id),
+                        )
+                    )
+                )
+
+        if tasks:
+            await gather(*tasks)
+        return True
+
+    async def on_turn(self, context: TurnContext, logic: Callable[[], Awaitable]):  # type: ignore Bug in botbuilder-python https://github.com/microsoft/botbuilder-python/issues/2198
+        # Handle incoming message
+        await self.add_user_message(context)
+
+        # Store the original send_activities method
+        original_send_activities = context.send_activities
+
+        # Create a wrapped version that captures the activities
+        # We need to do this because bot-framework has a bug with how
+        # _on_send_activities middleware is implemented
+        # https://github.com/microsoft/botbuilder-python/issues/2197
+        async def wrapped_send_activities(activities: List[Activity]):
+            responses = await original_send_activities(activities)
+            await self.add_agent_message(context, activities, responses)
+            return responses
+
+        # Replace the send_activities method
+        context.send_activities = wrapped_send_activities
+
+        # Run the bot's logic
+        await logic()

src/bot.py

@@ -1,4 +1,3 @@
-import datetime
 import json
 import os
 import sys
@@ -11,13 +10,12 @@
 from botbuilder.schema import Activity
 from litellm import acompletion
 from memory_module import (
-    AssistantMessageInput,
     InternalMessageInput,
     LLMConfig,
     Memory,
+    MemoryMiddleware,
     MemoryModule,
     MemoryModuleConfig,
-    UserMessageInput,
 )
 from pydantic import BaseModel, Field
 from teams import Application, ApplicationOptions, TeamsAdapter
@@ -63,6 +61,8 @@
     )
 )
 
+bot_app.adapter.use(MemoryMiddleware(memory_module))
+
 
 class TaskConfig(BaseModel):
     task_name: str
@@ -281,75 +281,6 @@ def get_available_functions():
     ]
 
 
-def build_deep_link(context: TurnContext, message_id: str):
-    conversation_ref = TurnContext.get_conversation_reference(context.activity)
-    if conversation_ref.conversation and conversation_ref.conversation.is_group:
-        deeplink_conversation_id = conversation_ref.conversation.id
-    elif conversation_ref.user and conversation_ref.bot:
-        user_aad_object_id = conversation_ref.user.aad_object_id
-        bot_id = conversation_ref.bot.id.replace("28:", "")
-        deeplink_conversation_id = f"19:{user_aad_object_id}_{bot_id}@unq.gbl.spaces"
-    else:
-        return None
-    return f"https://teams.microsoft.com/l/message/{deeplink_conversation_id}/{message_id}?context=%7B%22contextType%22%3A%22chat%22%7D"
-
-
-async def add_user_message(context: TurnContext):
-    conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
-    content = context.activity.text
-    if not content:
-        print("content is not text, so ignoring...")
-        return False
-    if conversation_ref_dict is None:
-        print("conversation_ref_dict is None")
-        return False
-    if conversation_ref_dict.user is None:
-        print("conversation_ref_dict.user is None")
-        return False
-    if conversation_ref_dict.conversation is None:
-        print("conversation_ref_dict.conversation is None")
-        return False
-    user_aad_object_id = conversation_ref_dict.user.aad_object_id
-    message_id = context.activity.id
-    await memory_module.add_message(
-        UserMessageInput(
-            id=message_id,
-            content=context.activity.text,
-            author_id=user_aad_object_id,
-            conversation_ref=conversation_ref_dict.conversation.id,
-            created_at=context.activity.timestamp if context.activity.timestamp else datetime.datetime.now(),
-            deep_link=build_deep_link(context, context.activity.id),
-        )
-    )
-    return True
-
-
-async def add_agent_message(context: TurnContext, message_id: str, content: str):
-    conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
-    if not content:
-        print("content is not text, so ignoring...")
-        return False
-    if conversation_ref_dict is None:
-        print("conversation_ref_dict is None")
-        return False
-    if conversation_ref_dict.bot is None:
-        print("conversation_ref_dict.bot is None")
-        return False
-    if conversation_ref_dict.conversation is None:
-        print("conversation_ref_dict.conversation is None")
-        return False
-    await memory_module.add_message(
-        AssistantMessageInput(
-            id=message_id,
-            content=content,
-            author_id=conversation_ref_dict.bot.id,
-            conversation_ref=conversation_ref_dict.conversation.id,
-            deep_link=build_deep_link(context, message_id),
-        )
-    )
-    return True
-
-
 async def add_internal_message(context: TurnContext, content: str):
     conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
     if not content:
@@ -376,15 +307,12 @@ async def add_internal_message(context: TurnContext, content: str):
 
 @bot_app.conversation_update("membersAdded")
 async def on_members_added(context: TurnContext, state: TurnState):
-    result = await send_string_message(context, "Hello! I'm your IT Support Assistant. How can I assist you today?")
-    if result:
-        await add_agent_message(context, result, "Hello! I'm your IT Support Assistant. How can I assist you today?")
+    await send_string_message(context, "Hello! I'm your IT Support Assistant. How can I assist you today?")
 
 
 @bot_app.activity("message")
 async def on_message(context: TurnContext, state: TurnState):
     conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
-    await add_user_message(context)
     system_prompt = """
 You are an IT Chat Bot that helps users troubleshoot tasks
 
@@ -449,9 +377,7 @@ async def on_message(context: TurnContext, state: TurnState):
         message = response.choices[0].message
 
         if message.tool_calls is None and message.content is not None:
-            agent_message_id = await send_string_message(context, message.content)
-            if agent_message_id:
-                await add_agent_message(context, agent_message_id, message.content)
+            await send_string_message(context, message.content)
             break
         elif message.tool_calls is None and message.content is None:
             print("No tool calls and no content")