remove changes in agent folder

src/agent/custom_agent.py

@@ -69,11 +69,6 @@ class CustomAgent(Agent):
             max_actions_per_step: int = 10,
             tool_call_in_content: bool = True,
     ):
-        # Store tool_call_in_content before calling parent's __init__
-        self.tool_call_in_content = tool_call_in_content
-        self.add_infos = add_infos
-        
-        # Call parent's __init__ without tool_call_in_content
         super().__init__(
             task=task,
             llm=llm,
@@ -90,9 +85,9 @@ class CustomAgent(Agent):
             include_attributes=include_attributes,
             max_error_length=max_error_length,
             max_actions_per_step=max_actions_per_step,
+            tool_call_in_content=tool_call_in_content,
         )
-        
-        # Initialize message manager with tool_call_in_content
+        self.add_infos = add_infos
         self.message_manager = CustomMassageManager(
             llm=self.llm,
             task=self.task,
@@ -102,7 +97,7 @@ class CustomAgent(Agent):
             include_attributes=self.include_attributes,
             max_error_length=self.max_error_length,
             max_actions_per_step=self.max_actions_per_step,
-            tool_call_in_content=self.tool_call_in_content,
+            tool_call_in_content=tool_call_in_content,
         )
 
     def _setup_action_models(self) -> None:
@@ -287,4 +282,4 @@ class CustomAgent(Agent):
                 await self.browser_context.close()
 
             if not self.injected_browser and self.browser:
-                await self.browser.close()
+                await self.browser.close()

src/agent/custom_massage_manager.py

@@ -12,8 +12,7 @@ from typing import List, Optional, Type
 from browser_use.agent.message_manager.service import MessageManager
 from browser_use.agent.message_manager.views import MessageHistory
 from browser_use.agent.prompts import SystemPrompt
-from browser_use.agent.views import ActionResult
-from .custom_views import CustomAgentStepInfo
+from browser_use.agent.views import ActionResult, AgentStepInfo
 from browser_use.browser.views import BrowserState
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import (
@@ -41,7 +40,6 @@ class CustomMassageManager(MessageManager):
             max_actions_per_step: int = 10,
             tool_call_in_content: bool = False,
     ):
-        self.tool_call_in_content = tool_call_in_content
         super().__init__(
             llm=llm,
             task=task,
@@ -53,17 +51,13 @@ class CustomMassageManager(MessageManager):
             include_attributes=include_attributes,
             max_error_length=max_error_length,
             max_actions_per_step=max_actions_per_step,
+            tool_call_in_content=tool_call_in_content,
         )
 
         # Custom: Move Task info to state_message
         self.history = MessageHistory()
         self._add_message_with_tokens(self.system_prompt)
-        tool_calls = self._create_tool_calls()
-        example_tool_call = self._create_example_tool_call(tool_calls)
-        self._add_message_with_tokens(example_tool_call)
-
-    def _create_tool_calls(self):
-        return [
+        tool_calls = [
             {
                 'name': 'CustomAgentOutput',
                 'args': {
@@ -80,25 +74,25 @@ class CustomMassageManager(MessageManager):
                 'type': 'tool_call',
             }
         ]
-
-    def _create_example_tool_call(self, tool_calls):
         if self.tool_call_in_content:
             # openai throws error if tool_calls are not responded -> move to content
-            return AIMessage(
+            example_tool_call = AIMessage(
                 content=f'{tool_calls}',
                 tool_calls=[],
             )
         else:
-            return AIMessage(
+            example_tool_call = AIMessage(
                 content=f'',
                 tool_calls=tool_calls,
             )
 
+        self._add_message_with_tokens(example_tool_call)
+
     def add_state_message(
-        self,
-        state: BrowserState,
-        result: Optional[List[ActionResult]] = None,
-        step_info: Optional[CustomAgentStepInfo] = None,
+            self,
+            state: BrowserState,
+            result: Optional[List[ActionResult]] = None,
+            step_info: Optional[AgentStepInfo] = None,
     ) -> None:
         """Add browser state as human message"""
 
@@ -111,7 +105,7 @@ class CustomMassageManager(MessageManager):
                         self._add_message_with_tokens(msg)
                     if r.error:
                         msg = HumanMessage(
-                            content=str(r.error)[-self.max_error_length :]
+                            content=str(r.error)[-self.max_error_length:]
                         )
                         self._add_message_with_tokens(msg)
                     result = None  # if result in history, we dont want to add it again
@@ -124,4 +118,4 @@ class CustomMassageManager(MessageManager):
             max_error_length=self.max_error_length,
             step_info=step_info,
         ).get_user_message()
-        self._add_message_with_tokens(state_message)
+        self._add_message_with_tokens(state_message)