feat: enhance message context handling by adding message_id parameter
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
Reference in New Issue
Block a user