Integrating MemGPT-like Functionality

agenthub/codeact_agent/action_parser.py

@@ -5,6 +5,7 @@
     Action,
     AgentDelegateAction,
     AgentFinishAction,
+    AgentSummarizeAction,
     CmdRunAction,
     IPythonRunCellAction,
     MessageAction,
@@ -32,6 +33,8 @@ def __init__(self):
         self.default_parser = CodeActActionParserMessage()
 
     def parse(self, response) -> Action:
+        if isinstance(response, AgentSummarizeAction):
+            return response
         action_str = self.parse_response(response)
         return self.parse_action(action_str)
 

agenthub/codeact_agent/codeact_agent.py

@@ -12,6 +12,7 @@
     Action,
     AgentDelegateAction,
     AgentFinishAction,
+    AgentSummarizeAction,
     CmdRunAction,
     IPythonRunCellAction,
     MessageAction,
@@ -24,6 +25,7 @@
 from opendevin.events.observation.observation import Observation
 from opendevin.events.serialization.event import truncate_content
 from opendevin.llm.llm import LLM
+from opendevin.llm.messages import Message
 from opendevin.runtime.plugins import (
     AgentSkillsRequirement,
     JupyterRequirement,
@@ -121,25 +123,39 @@ def action_to_str(self, action: Action) -> str:
             return f'{action.thought}\n<execute_browse>\n{action.inputs["task"]}\n</execute_browse>'
         elif isinstance(action, MessageAction):
             return action.content
+        elif isinstance(action, AgentSummarizeAction):
+            return (
+                'Summary of all Action and Observations till now. \n'
+                + 'Action: '
+                + action.summarized_actions
+                + '\nObservation: '
+                + action.summarized_observations
+            )
         elif isinstance(action, AgentFinishAction) and action.source == 'agent':
             return action.thought
         return ''
 
-    def get_action_message(self, action: Action) -> dict[str, str] | None:
+    def get_action_message(self, action: Action) -> Message | None:
+        message = None
         if (
             isinstance(action, AgentDelegateAction)
             or isinstance(action, CmdRunAction)
             or isinstance(action, IPythonRunCellAction)
             or isinstance(action, MessageAction)
+            or isinstance(action, AgentSummarizeAction)
             or (isinstance(action, AgentFinishAction) and action.source == 'agent')
         ):
-            return {
+            message = {
                 'role': 'user' if action.source == 'user' else 'assistant',
                 'content': self.action_to_str(action),
             }
-        return None
+        if message:
+            return Message(message=message, condensable=True, event_id=action.id)
+        else:
+            return None
 
-    def get_observation_message(self, obs: Observation) -> dict[str, str] | None:
+    def get_observation_message(self, obs: Observation) -> Message | None:
+        message = None
         max_message_chars = self.llm.config.max_message_chars
         if isinstance(obs, CmdOutputObservation):
             content = 'OBSERVATION:\n' + truncate_content(
@@ -148,7 +164,7 @@ def get_observation_message(self, obs: Observation) -> dict[str, str] | None:
             content += (
                 f'\n[Command {obs.command_id} finished with exit code {obs.exit_code}]'
             )
-            return {'role': 'user', 'content': content}
+            message = {'role': 'user', 'content': content}
         elif isinstance(obs, IPythonRunCellObservation):
             content = 'OBSERVATION:\n' + obs.content
             # replace base64 images with a placeholder
@@ -160,13 +176,16 @@ def get_observation_message(self, obs: Observation) -> dict[str, str] | None:
                     )
             content = '\n'.join(splitted)
             content = truncate_content(content, max_message_chars)
-            return {'role': 'user', 'content': content}
+            message = {'role': 'user', 'content': content}
         elif isinstance(obs, AgentDelegateObservation):
             content = 'OBSERVATION:\n' + truncate_content(
                 str(obs.outputs), max_message_chars
             )
-            return {'role': 'user', 'content': content}
-        return None
+            message = {'role': 'user', 'content': content}
+        if message:
+            return Message(message=message, condensable=True, event_id=obs.id)
+        else:
+            return None
 
     def reset(self) -> None:
         """Resets the CodeAct Agent."""
@@ -191,48 +210,69 @@ def step(self, state: State) -> Action:
         if latest_user_message and latest_user_message.strip() == '/exit':
             return AgentFinishAction()
 
-        # prepare what we want to send to the LLM
-        messages: list[dict[str, str]] = self._get_messages(state)
-
-        response = self.llm.completion(
-            messages=messages,
-            stop=[
-                '</execute_ipython>',
-                '</execute_bash>',
-                '</execute_browse>',
-            ],
-            temperature=0.0,
-        )
+        response = None
+        # give it multiple chances to get a response
+        # if it fails, we'll try to condense memory
+        attempt = 0
+        while not response and attempt < self.llm.config.attempts_to_condense:
+            # prepare what we want to send to the LLM
+            messages: list[Message] = self._get_messages(state)
+            print('No of tokens, ' + str(self.llm.get_token_count(messages)) + '\n')
+            response = self.llm.completion(
+                messages=messages,
+                stop=[
+                    '</execute_ipython>',
+                    '</execute_bash>',
+                    '</execute_browse>',
+                ],
+                temperature=0.0,
+                condense=True,
+            )
+            attempt += 1
+
         return self.action_parser.parse(response)
 
-    def _get_messages(self, state: State) -> list[dict[str, str]]:
+    def search_memory(self, query: str) -> list[str]:
+        raise NotImplementedError('Implement this abstract method')
+
+    def _get_messages(self, state: State) -> list[Message]:
         messages = [
-            {'role': 'system', 'content': self.system_message},
-            {'role': 'user', 'content': self.in_context_example},
+            Message(
+                message={'role': 'system', 'content': self.system_message},
+                condensable=False,
+            ),
+            Message(
+                message={'role': 'user', 'content': self.in_context_example},
+                condensable=False,
+            ),
         ]
 
+        if state.history.summary:
+            summary_message = self.get_action_message(state.history.summary)
+            if summary_message:
+                messages.append(summary_message)
         for event in state.history.get_events():
-            # create a regular message from an event
-            if isinstance(event, Action):
-                message = self.get_action_message(event)
-            elif isinstance(event, Observation):
-                message = self.get_observation_message(event)
-            else:
-                raise ValueError(f'Unknown event type: {type(event)}')
-
-            # add regular message
-            if message:
-                messages.append(message)
+            if event.id > state.history.last_summarized_event_id:
+                # create a regular message from an event
+                if isinstance(event, Action):
+                    message = self.get_action_message(event)
+                elif isinstance(event, Observation):
+                    message = self.get_observation_message(event)
+                else:
+                    raise ValueError(f'Unknown event type: {type(event)}')
+                # add regular message
+                if message:
+                    messages.append(message)
 
         # the latest user message is important:
         # we want to remind the agent of the environment constraints
         latest_user_message = next(
-            (m for m in reversed(messages) if m['role'] == 'user'), None
+            (m for m in reversed(messages) if m.message['role'] == 'user'), None
         )
 
         # add a reminder to the prompt
         if latest_user_message:
-            latest_user_message['content'] += (
+            latest_user_message.message['content'] += (
                 f'\n\nENVIRONMENT REMINDER: You have {state.max_iterations - state.iteration} turns left to complete the task. When finished reply with <finish></finish>'
             )
 

opendevin/condenser/condenser.py

@@ -0,0 +1,146 @@
+from opendevin.core.exceptions import (
+    SummarizeError,
+)
+from opendevin.core.logger import opendevin_logger as logger
+from opendevin.events.action import (
+    AgentSummarizeAction,
+)
+from opendevin.llm.messages import Message
+
+from .prompts import (
+    MESSAGE_SUMMARY_WARNING_FRAC,
+    SUMMARY_PROMPT_SYSTEM,
+    parse_summary_response,
+)
+
+
+class CondenserMixin:
+    """Condenses a group of condensable messages as done by MemGPT."""
+
+    def condense(
+        self,
+        messages: list[Message],
+    ):
+        # Start past the system message, and example messages.,
+        # and collect messages for summarization until we reach the desired truncation token fraction (eg 50%)
+        # Do not allow truncation  for in-context examples of function calling
+        token_counts = [
+            self.get_token_count([message])  # type: ignore
+            for message in messages
+            if message.condensable
+        ]
+        message_buffer_token_count = sum(token_counts)  # no system and example message
+
+        desired_token_count_to_summarize = int(
+            message_buffer_token_count * self.config.message_summary_trunc_tokens_frac  # type: ignore
+        )
+
+        candidate_messages_to_summarize = []
+        tokens_so_far = 0
+        for message in messages:
+            if message.condensable:
+                candidate_messages_to_summarize.append(message)
+                tokens_so_far += self.get_token_count([message])  # type: ignore
+            if tokens_so_far > desired_token_count_to_summarize:
+                last_summarized_event_id = message.event_id
+                break
+
+        # TODO: Add functionality for preserving last N messages
+        # MESSAGE_SUMMARY_TRUNC_KEEP_N_LAST = 3
+        # if preserve_last_N_messages:
+        #     candidate_messages_to_summarize = candidate_messages_to_summarize[:-MESSAGE_SUMMARY_TRUNC_KEEP_N_LAST]
+        #     token_counts = token_counts[:-MESSAGE_SUMMARY_TRUNC_KEEP_N_LAST]
+
+        logger.debug(
+            f'message_summary_trunc_tokens_frac={self.config.message_summary_trunc_tokens_frac}'  # type: ignore
+        )
+        # logger.debug(f'MESSAGE_SUMMARY_TRUNC_KEEP_N_LAST={MESSAGE_SUMMARY_TRUNC_KEEP_N_LAST}')
+        logger.debug(f'token_counts={token_counts}')
+        logger.debug(f'message_buffer_token_count={message_buffer_token_count}')
+        logger.debug(
+            f'desired_token_count_to_summarize={desired_token_count_to_summarize}'
+        )
+        logger.debug(
+            f'len(candidate_messages_to_summarize)={len(candidate_messages_to_summarize)}'
+        )
+
+        if len(candidate_messages_to_summarize) == 0:
+            raise SummarizeError(
+                f"Summarize error: tried to run summarize, but couldn't find enough messages to compress [len={len(messages)}]"
+            )
+
+        # TODO: Try to make an assistant message come after the cutoff
+
+        message_sequence_to_summarize = candidate_messages_to_summarize
+
+        if len(message_sequence_to_summarize) <= 1:
+            # This prevents a potential infinite loop of summarizing the same message over and over
+            raise SummarizeError(
+                f"Summarize error: tried to run summarize, but couldn't find enough messages to compress [len={len(message_sequence_to_summarize)} <= 1]"
+            )
+        else:
+            print(
+                f'Attempting to summarize with last summarized event id = {last_summarized_event_id}'
+            )
+
+        action_response = self.summarize_messages(
+            message_sequence_to_summarize=message_sequence_to_summarize
+        )
+        summary_action: AgentSummarizeAction = parse_summary_response(action_response)
+        summary_action.last_summarized_event_id = (
+            last_summarized_event_id if last_summarized_event_id else -1
+        )
+        return summary_action
+
+    def _format_summary_history(self, message_history: list[dict]) -> str:
+        # TODO use existing prompt formatters for this (eg ChatML)
+        return '\n'.join([f'{m["role"]}: {m["content"]}' for m in message_history])
+
+    def summarize_messages(self, message_sequence_to_summarize: list[Message]):
+        """Summarize a message sequence using LLM"""
+        context_window = self.config.max_input_tokens  # type: ignore
+        summary_prompt = SUMMARY_PROMPT_SYSTEM
+        summary_input = self._format_summary_history(
+            self.get_text_messages(message_sequence_to_summarize)  # type: ignore
+        )
+        summary_input_tkns = self.get_token_count(summary_input)  # type: ignore
+        if context_window is None:
+            raise ValueError('context_window should not be None')
+        if summary_input_tkns > MESSAGE_SUMMARY_WARNING_FRAC * context_window:
+            trunc_ratio = (
+                MESSAGE_SUMMARY_WARNING_FRAC * context_window / summary_input_tkns
+            ) * 0.8  # For good measure...
+            cutoff = int(len(message_sequence_to_summarize) * trunc_ratio)
+            curr_summary = self.summarize_messages(
+                message_sequence_to_summarize=message_sequence_to_summarize[:cutoff]
+            )
+            curr_summary_message = (
+                'Summary of all Action and Observations till now. \n'
+                + 'Action: '
+                + curr_summary['args']['summarized_actions']
+                + '\nObservation: '
+                + curr_summary['args']['summarized_observations']
+            )
+            input = [
+                Message({'role': 'assistant', 'content': curr_summary_message})
+            ] + message_sequence_to_summarize[cutoff:]
+            summary_input = self._format_summary_history(self.get_text_messages(input))  # type: ignore
+
+        message_sequence = []
+        message_sequence.append(Message({'role': 'system', 'content': summary_prompt}))
+        message_sequence.append(Message({'role': 'user', 'content': summary_input}))
+
+        response = self.completion(  # type: ignore
+            messages=message_sequence,
+            stop=[
+                '</execute_ipython>',
+                '</execute_bash>',
+                '</execute_browse>',
+            ],
+            temperature=0.0,
+        )
+
+        print(f'summarize_messages gpt reply: {response.choices[0]}')
+
+        action_response = response['choices'][0]['message']['content']
+        return action_response