diff --git a/packages/memory_module/__init__.py b/packages/memory_module/__init__.py index f7eea423..a336ceee 100644 --- a/packages/memory_module/__init__.py +++ b/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", ] diff --git a/packages/memory_module/pyproject.toml b/packages/memory_module/pyproject.toml index 67a7d23a..74aa491a 100644 --- a/packages/memory_module/pyproject.toml +++ b/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] diff --git a/packages/memory_module/storage/sqlite_memory_storage.py b/packages/memory_module/storage/sqlite_memory_storage.py index 52a6bf37..fa089a55 100644 --- a/packages/memory_module/storage/sqlite_memory_storage.py +++ b/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, ), diff --git a/packages/memory_module/storage/sqlite_message_buffer_storage.py b/packages/memory_module/storage/sqlite_message_buffer_storage.py index edd5b91e..1d46fc7e 100644 --- a/packages/memory_module/storage/sqlite_message_buffer_storage.py +++ b/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(), ), ) diff --git a/packages/memory_module/utils/teams_bot_middlware.py b/packages/memory_module/utils/teams_bot_middlware.py new file mode 100644 index 00000000..da88ac91 --- /dev/null +++ b/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() diff --git a/src/bot.py b/src/bot.py index 57263cb5..ce7a5b68 100644 --- a/src/bot.py +++ b/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") diff --git a/uv.lock b/uv.lock index c6795b01..4c547a21 100644 --- a/uv.lock +++ b/uv.lock @@ -1298,6 +1298,7 @@ version = "0.0.0" source = { virtual = "packages/memory_module" } dependencies = [ { name = "aiosqlite" }, + { name = "botbuilder" }, { name = "instructor" }, { name = "litellm" }, { name = "numpy" }, @@ -1309,6 +1310,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "aiosqlite", specifier = ">=0.20.0" }, + { name = "botbuilder", specifier = ">=0.0.1" }, { name = "instructor", specifier = ">=1.6.4" }, { name = "litellm", specifier = "==1.54.1" }, { name = "numpy" },