From fef06fdc8a5057139c65809af62bc2af4aef2fe8 Mon Sep 17 00:00:00 2001
From: afourney <adamfo@microsoft.com>
Date: Wed, 4 Dec 2024 15:13:48 -0800
Subject: [PATCH] Fixed messages missing from websurfer history. (#4551)

* Fixed messages missing from websurfer history.
* Fix filesurfer too!
---
 .../agents/file_surfer/_file_surfer.py          | 15 ++++++++-------
 .../agents/web_surfer/_multimodal_web_surfer.py | 17 +++++++++++------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/file_surfer/_file_surfer.py b/python/packages/autogen-ext/src/autogen_ext/agents/file_surfer/_file_surfer.py
index fcbc4f58e..bd0988c5d 100644
--- a/python/packages/autogen-ext/src/autogen_ext/agents/file_surfer/_file_surfer.py
+++ b/python/packages/autogen-ext/src/autogen_ext/agents/file_surfer/_file_surfer.py
@@ -11,6 +11,7 @@ from autogen_agentchat.messages import (
 )
 from autogen_core import CancellationToken, FunctionCall
 from autogen_core.components.models import (
+    AssistantMessage,
     ChatCompletionClient,
     LLMMessage,
     SystemMessage,
@@ -69,16 +70,16 @@ class FileSurfer(BaseChatAgent):
                 self._chat_history.append(UserMessage(content=chat_message.content, source=chat_message.source))
             else:
                 raise ValueError(f"Unexpected message in FileSurfer: {chat_message}")
+
         try:
             _, content = await self._generate_reply(cancellation_token=cancellation_token)
-            if isinstance(content, str):
-                return Response(chat_message=TextMessage(content=content, source=self.name))
-            else:
-                return Response(chat_message=MultiModalMessage(content=content, source=self.name))
+            self._chat_history.append(AssistantMessage(content=content, source=self.name))
+            return Response(chat_message=TextMessage(content=content, source=self.name))
+
         except BaseException:
-            return Response(
-                chat_message=TextMessage(content=f"File surfing error:\n\n{traceback.format_exc()}", source=self.name)
-            )
+            content = f"File surfing error:\n\n{traceback.format_exc()}"
+            self._chat_history.append(AssistantMessage(content, source=self.name))
+            return Response(chat_message=TextMessage(content=content, source=self.name))
 
     async def on_reset(self, cancellation_token: CancellationToken) -> None:
         self._chat_history.clear()
diff --git a/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/_multimodal_web_surfer.py b/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/_multimodal_web_surfer.py
index 69cd067c7..3f2e8aeae 100644
--- a/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/_multimodal_web_surfer.py
+++ b/python/packages/autogen-ext/src/autogen_ext/agents/web_surfer/_multimodal_web_surfer.py
@@ -158,6 +158,7 @@ class MultimodalWebSurfer(BaseChatAgent):
         )
         self.default_tools = [
             TOOL_VISIT_URL,
+            TOOL_WEB_SEARCH,
             TOOL_HISTORY_BACK,
             TOOL_CLICK,
             TOOL_TYPE,
@@ -181,14 +182,16 @@ class MultimodalWebSurfer(BaseChatAgent):
 
         try:
             _, content = await self.__generate_reply(cancellation_token=cancellation_token)
+            self._chat_history.append(AssistantMessage(content=message_content_to_str(content), source=self.name))
             if isinstance(content, str):
                 return Response(chat_message=TextMessage(content=content, source=self.name))
             else:
                 return Response(chat_message=MultiModalMessage(content=content, source=self.name))
+
         except BaseException:
-            return Response(
-                chat_message=TextMessage(content=f"Web surfing error:\n\n{traceback.format_exc()}", source=self.name)
-            )
+            content = f"Web surfing error:\n\n{traceback.format_exc()}"
+            self._chat_history.append(AssistantMessage(content, source=self.name))
+            return Response(chat_message=TextMessage(content=content, source=self.name))
 
     async def on_reset(self, cancellation_token: CancellationToken) -> None:
         if not self.did_lazy_init:
@@ -532,9 +535,12 @@ class MultimodalWebSurfer(BaseChatAgent):
         # Clone the messages to give context, removing old screenshots
         history: List[LLMMessage] = []
         for m in self._chat_history:
+            assert isinstance(m, UserMessage | AssistantMessage | SystemMessage)
+            assert isinstance(m.content, str | list)
+
             if isinstance(m.content, str):
                 history.append(m)
-            elif isinstance(m.content, list):
+            else:
                 content = message_content_to_str(m.content)
                 if isinstance(m, UserMessage):
                     history.append(UserMessage(content=content, source=m.source))
@@ -563,8 +569,6 @@ class MultimodalWebSurfer(BaseChatAgent):
         # What tools are available?
         tools = self.default_tools.copy()
 
-        tools.append(TOOL_WEB_SEARCH)
-
         # We can scroll up
         if viewport["pageTop"] > 5:
             tools.append(TOOL_PAGE_UP)
@@ -628,6 +632,7 @@ class MultimodalWebSurfer(BaseChatAgent):
 
         # Add the multimodal message and make the request
         history.append(UserMessage(content=[text_prompt, AGImage.from_pil(scaled_screenshot)], source=self.name))
+
         response = await self._model_client.create(
             history, tools=tools, extra_create_args={"tool_choice": "auto"}, cancellation_token=cancellation_token
         )  # , "parallel_tool_calls": False})