Skip to content

Commit 3d9c347

Browse files
authored
Add memory module middleware (#69)
1 parent 8ac8981 commit 3d9c347

File tree

7 files changed

+127
-82
lines changed

7 files changed

+127
-82
lines changed

packages/memory_module/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
UserMessage,
1313
UserMessageInput,
1414
)
15+
from memory_module.utils.teams_bot_middlware import MemoryMiddleware
1516

1617
__all__ = [
1718
"MemoryModule",
@@ -27,4 +28,5 @@
2728
"AssistantMessage",
2829
"AssistantMessageInput",
2930
"ShortTermMemoryRetrievalConfig",
31+
"MemoryMiddleware",
3032
]

packages/memory_module/pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ dependencies = [
1212
"numpy",
1313
"sqlite-vec>=0.1.6",
1414
"litellm==1.54.1",
15+
"botbuilder>=0.0.1",
1516
]
1617

1718
[tool.uv]

packages/memory_module/storage/sqlite_memory_storage.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ async def store_memory(self, memory: BaseMemoryInput, *, embedding_vectors: List
4747
(
4848
memory_id,
4949
memory.content,
50-
memory.created_at,
50+
memory.created_at.isoformat(),
5151
memory.user_id,
5252
memory.memory_type.value,
5353
),
@@ -300,7 +300,7 @@ async def store_short_term_memory(self, message: MessageInput) -> Message:
300300
message.content,
301301
message.author_id,
302302
message.conversation_ref,
303-
created_at,
303+
created_at.isoformat(),
304304
message.type,
305305
deep_link,
306306
),

packages/memory_module/storage/sqlite_message_buffer_storage.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ async def store_buffered_message(self, message: Message) -> None:
3838
(
3939
message.id,
4040
message.conversation_ref,
41-
message.created_at,
41+
message.created_at.isoformat(),
4242
),
4343
)
4444

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
import datetime
2+
from asyncio import gather
3+
from typing import Awaitable, Callable, List
4+
5+
from botbuilder.core import TurnContext
6+
from botbuilder.core.middleware_set import Middleware
7+
from botbuilder.schema import Activity, ResourceResponse
8+
from memory_module.interfaces.base_memory_module import BaseMemoryModule
9+
from memory_module.interfaces.types import (
10+
AssistantMessageInput,
11+
UserMessageInput,
12+
)
13+
14+
15+
def build_deep_link(context: TurnContext, message_id: str):
16+
conversation_ref = TurnContext.get_conversation_reference(context.activity)
17+
if conversation_ref.conversation and conversation_ref.conversation.is_group:
18+
deeplink_conversation_id = conversation_ref.conversation.id
19+
elif conversation_ref.user and conversation_ref.bot:
20+
user_aad_object_id = conversation_ref.user.aad_object_id
21+
bot_id = conversation_ref.bot.id.replace("28:", "")
22+
deeplink_conversation_id = f"19:{user_aad_object_id}_{bot_id}@unq.gbl.spaces"
23+
else:
24+
return None
25+
return f"https://teams.microsoft.com/l/message/{deeplink_conversation_id}/{message_id}?context=%7B%22contextType%22%3A%22chat%22%7D"
26+
27+
28+
class MemoryMiddleware(Middleware):
29+
def __init__(self, memory_module: BaseMemoryModule):
30+
self.memory_module = memory_module
31+
32+
async def add_user_message(self, context: TurnContext):
33+
conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
34+
content = context.activity.text
35+
if not content:
36+
print("content is not text, so ignoring...")
37+
return False
38+
if conversation_ref_dict is None:
39+
print("conversation_ref_dict is None")
40+
return False
41+
if conversation_ref_dict.user is None:
42+
print("conversation_ref_dict.user is None")
43+
return False
44+
if conversation_ref_dict.conversation is None:
45+
print("conversation_ref_dict.conversation is None")
46+
return False
47+
user_aad_object_id = conversation_ref_dict.user.aad_object_id
48+
message_id = context.activity.id
49+
await self.memory_module.add_message(
50+
UserMessageInput(
51+
id=message_id,
52+
content=context.activity.text,
53+
author_id=user_aad_object_id,
54+
conversation_ref=conversation_ref_dict.conversation.id,
55+
created_at=context.activity.timestamp if context.activity.timestamp else datetime.datetime.now(),
56+
deep_link=build_deep_link(context, context.activity.id),
57+
)
58+
)
59+
return True
60+
61+
async def add_agent_message(
62+
self, context: TurnContext, activities: List[Activity], responses: List[ResourceResponse]
63+
):
64+
conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
65+
if conversation_ref_dict is None:
66+
print("conversation_ref_dict is None")
67+
return False
68+
if conversation_ref_dict.bot is None:
69+
print("conversation_ref_dict.bot is None")
70+
return False
71+
if conversation_ref_dict.conversation is None:
72+
print("conversation_ref_dict.conversation is None")
73+
return False
74+
75+
tasks = []
76+
for activity, response in zip(activities, responses, strict=False):
77+
if activity.text:
78+
tasks.append(
79+
self.memory_module.add_message(
80+
AssistantMessageInput(
81+
id=response.id,
82+
content=activity.text,
83+
author_id=conversation_ref_dict.bot.id,
84+
conversation_ref=conversation_ref_dict.conversation.id,
85+
deep_link=build_deep_link(context, response.id),
86+
)
87+
)
88+
)
89+
90+
if tasks:
91+
await gather(*tasks)
92+
return True
93+
94+
async def on_turn(self, context: TurnContext, logic: Callable[[], Awaitable]): # type: ignore Bug in botbuilder-python https://github.com/microsoft/botbuilder-python/issues/2198
95+
# Handle incoming message
96+
await self.add_user_message(context)
97+
98+
# Store the original send_activities method
99+
original_send_activities = context.send_activities
100+
101+
# Create a wrapped version that captures the activities
102+
# We need to do this because bot-framework has a bug with how
103+
# _on_send_activities middleware is implemented
104+
# https://github.com/microsoft/botbuilder-python/issues/2197
105+
async def wrapped_send_activities(activities: List[Activity]):
106+
responses = await original_send_activities(activities)
107+
await self.add_agent_message(context, activities, responses)
108+
return responses
109+
110+
# Replace the send_activities method
111+
context.send_activities = wrapped_send_activities
112+
113+
# Run the bot's logic
114+
await logic()

src/bot.py

Lines changed: 5 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import datetime
21
import json
32
import os
43
import sys
@@ -11,13 +10,12 @@
1110
from botbuilder.schema import Activity
1211
from litellm import acompletion
1312
from memory_module import (
14-
AssistantMessageInput,
1513
InternalMessageInput,
1614
LLMConfig,
1715
Memory,
16+
MemoryMiddleware,
1817
MemoryModule,
1918
MemoryModuleConfig,
20-
UserMessageInput,
2119
)
2220
from pydantic import BaseModel, Field
2321
from teams import Application, ApplicationOptions, TeamsAdapter
@@ -63,6 +61,8 @@
6361
)
6462
)
6563

64+
bot_app.adapter.use(MemoryMiddleware(memory_module))
65+
6666

6767
class TaskConfig(BaseModel):
6868
task_name: str
@@ -281,75 +281,6 @@ def get_available_functions():
281281
]
282282

283283

284-
def build_deep_link(context: TurnContext, message_id: str):
285-
conversation_ref = TurnContext.get_conversation_reference(context.activity)
286-
if conversation_ref.conversation and conversation_ref.conversation.is_group:
287-
deeplink_conversation_id = conversation_ref.conversation.id
288-
elif conversation_ref.user and conversation_ref.bot:
289-
user_aad_object_id = conversation_ref.user.aad_object_id
290-
bot_id = conversation_ref.bot.id.replace("28:", "")
291-
deeplink_conversation_id = f"19:{user_aad_object_id}_{bot_id}@unq.gbl.spaces"
292-
else:
293-
return None
294-
return f"https://teams.microsoft.com/l/message/{deeplink_conversation_id}/{message_id}?context=%7B%22contextType%22%3A%22chat%22%7D"
295-
296-
297-
async def add_user_message(context: TurnContext):
298-
conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
299-
content = context.activity.text
300-
if not content:
301-
print("content is not text, so ignoring...")
302-
return False
303-
if conversation_ref_dict is None:
304-
print("conversation_ref_dict is None")
305-
return False
306-
if conversation_ref_dict.user is None:
307-
print("conversation_ref_dict.user is None")
308-
return False
309-
if conversation_ref_dict.conversation is None:
310-
print("conversation_ref_dict.conversation is None")
311-
return False
312-
user_aad_object_id = conversation_ref_dict.user.aad_object_id
313-
message_id = context.activity.id
314-
await memory_module.add_message(
315-
UserMessageInput(
316-
id=message_id,
317-
content=context.activity.text,
318-
author_id=user_aad_object_id,
319-
conversation_ref=conversation_ref_dict.conversation.id,
320-
created_at=context.activity.timestamp if context.activity.timestamp else datetime.datetime.now(),
321-
deep_link=build_deep_link(context, context.activity.id),
322-
)
323-
)
324-
return True
325-
326-
327-
async def add_agent_message(context: TurnContext, message_id: str, content: str):
328-
conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
329-
if not content:
330-
print("content is not text, so ignoring...")
331-
return False
332-
if conversation_ref_dict is None:
333-
print("conversation_ref_dict is None")
334-
return False
335-
if conversation_ref_dict.bot is None:
336-
print("conversation_ref_dict.bot is None")
337-
return False
338-
if conversation_ref_dict.conversation is None:
339-
print("conversation_ref_dict.conversation is None")
340-
return False
341-
await memory_module.add_message(
342-
AssistantMessageInput(
343-
id=message_id,
344-
content=content,
345-
author_id=conversation_ref_dict.bot.id,
346-
conversation_ref=conversation_ref_dict.conversation.id,
347-
deep_link=build_deep_link(context, message_id),
348-
)
349-
)
350-
return True
351-
352-
353284
async def add_internal_message(context: TurnContext, content: str):
354285
conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
355286
if not content:
@@ -376,15 +307,12 @@ async def add_internal_message(context: TurnContext, content: str):
376307

377308
@bot_app.conversation_update("membersAdded")
378309
async def on_members_added(context: TurnContext, state: TurnState):
379-
result = await send_string_message(context, "Hello! I'm your IT Support Assistant. How can I assist you today?")
380-
if result:
381-
await add_agent_message(context, result, "Hello! I'm your IT Support Assistant. How can I assist you today?")
310+
await send_string_message(context, "Hello! I'm your IT Support Assistant. How can I assist you today?")
382311

383312

384313
@bot_app.activity("message")
385314
async def on_message(context: TurnContext, state: TurnState):
386315
conversation_ref_dict = TurnContext.get_conversation_reference(context.activity)
387-
await add_user_message(context)
388316
system_prompt = """
389317
You are an IT Chat Bot that helps users troubleshoot tasks
390318
@@ -449,9 +377,7 @@ async def on_message(context: TurnContext, state: TurnState):
449377
message = response.choices[0].message
450378

451379
if message.tool_calls is None and message.content is not None:
452-
agent_message_id = await send_string_message(context, message.content)
453-
if agent_message_id:
454-
await add_agent_message(context, agent_message_id, message.content)
380+
await send_string_message(context, message.content)
455381
break
456382
elif message.tool_calls is None and message.content is None:
457383
print("No tool calls and no content")

uv.lock

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)