feat: enhance message context handling by adding message_id parameter

This commit is contained in:
Darye
2026-02-18 20:18:27 +01:00
parent 8de36d398f
commit c865b293a9
2 changed files with 15 additions and 7 deletions

View File

@@ -133,11 +133,11 @@ class AgentLoop:
await self._mcp_stack.__aenter__() await self._mcp_stack.__aenter__()
await connect_mcp_servers(self._mcp_servers, self.tools, self._mcp_stack) await connect_mcp_servers(self._mcp_servers, self.tools, self._mcp_stack)
def _set_tool_context(self, channel: str, chat_id: str) -> None: def _set_tool_context(self, channel: str, chat_id: str, message_id: str | None = None) -> None:
"""Update context for all tools that need routing info.""" """Update context for all tools that need routing info."""
if message_tool := self.tools.get("message"): if message_tool := self.tools.get("message"):
if isinstance(message_tool, MessageTool): if isinstance(message_tool, MessageTool):
message_tool.set_context(channel, chat_id) message_tool.set_context(channel, chat_id, message_id)
if spawn_tool := self.tools.get("spawn"): if spawn_tool := self.tools.get("spawn"):
if isinstance(spawn_tool, SpawnTool): if isinstance(spawn_tool, SpawnTool):
@@ -321,7 +321,7 @@ class AgentLoop:
if len(session.messages) > self.memory_window: if len(session.messages) > self.memory_window:
asyncio.create_task(self._consolidate_memory(session)) asyncio.create_task(self._consolidate_memory(session))
self._set_tool_context(msg.channel, msg.chat_id) self._set_tool_context(msg.channel, msg.chat_id, msg.metadata.get("message_id"))
initial_messages = self.context.build_messages( initial_messages = self.context.build_messages(
history=session.get_history(max_messages=self.memory_window), history=session.get_history(max_messages=self.memory_window),
current_message=msg.content, current_message=msg.content,
@@ -379,7 +379,7 @@ class AgentLoop:
session_key = f"{origin_channel}:{origin_chat_id}" session_key = f"{origin_channel}:{origin_chat_id}"
session = self.sessions.get_or_create(session_key) session = self.sessions.get_or_create(session_key)
self._set_tool_context(origin_channel, origin_chat_id) self._set_tool_context(origin_channel, origin_chat_id, msg.metadata.get("message_id"))
initial_messages = self.context.build_messages( initial_messages = self.context.build_messages(
history=session.get_history(max_messages=self.memory_window), history=session.get_history(max_messages=self.memory_window),
current_message=msg.content, current_message=msg.content,

View File

@@ -13,16 +13,19 @@ class MessageTool(Tool):
self, self,
send_callback: Callable[[OutboundMessage], Awaitable[None]] | None = None, send_callback: Callable[[OutboundMessage], Awaitable[None]] | None = None,
default_channel: str = "", default_channel: str = "",
default_chat_id: str = "" default_chat_id: str = "",
default_message_id: str | None = None
): ):
self._send_callback = send_callback self._send_callback = send_callback
self._default_channel = default_channel self._default_channel = default_channel
self._default_chat_id = default_chat_id self._default_chat_id = default_chat_id
self._default_message_id = default_message_id
def set_context(self, channel: str, chat_id: str) -> None: def set_context(self, channel: str, chat_id: str, message_id: str | None = None) -> None:
"""Set the current message context.""" """Set the current message context."""
self._default_channel = channel self._default_channel = channel
self._default_chat_id = chat_id self._default_chat_id = chat_id
self._default_message_id = message_id
def set_send_callback(self, callback: Callable[[OutboundMessage], Awaitable[None]]) -> None: def set_send_callback(self, callback: Callable[[OutboundMessage], Awaitable[None]]) -> None:
"""Set the callback for sending messages.""" """Set the callback for sending messages."""
@@ -67,11 +70,13 @@ class MessageTool(Tool):
content: str, content: str,
channel: str | None = None, channel: str | None = None,
chat_id: str | None = None, chat_id: str | None = None,
message_id: str | None = None,
media: list[str] | None = None, media: list[str] | None = None,
**kwargs: Any **kwargs: Any
) -> str: ) -> str:
channel = channel or self._default_channel channel = channel or self._default_channel
chat_id = chat_id or self._default_chat_id chat_id = chat_id or self._default_chat_id
message_id = message_id or self._default_message_id
if not channel or not chat_id: if not channel or not chat_id:
return "Error: No target channel/chat specified" return "Error: No target channel/chat specified"
@@ -83,7 +88,10 @@ class MessageTool(Tool):
channel=channel, channel=channel,
chat_id=chat_id, chat_id=chat_id,
content=content, content=content,
media=media or [] media=media or [],
metadata={
"message_id": message_id,
}
) )
try: try: