Refactoring observation component that uses memory retrieval by time …

…to represent observations.

PiperOrigin-RevId: 592503275
Change-Id: I9a68bb7457ae597d5701349d1b2a012e6d263aa2

concordia/associative_memory/associative_memory.py

@@ -268,6 +268,31 @@ def retrieve_by_regex(
 
     return self._pd_to_text(data, add_time=add_time, sort_by_time=sort_by_time)
 
+  def retrieve_time_interval(
+      self,
+      time_from: datetime.datetime,
+      time_until: datetime.datetime,
+      add_time: bool = False,
+  ):
+    """Retrieve memories within a time interval.
+
+    Args:
+      time_from: the start time of the interval
+      time_until: the end time of the interval
+      add_time: whether to add time stamp to the output
+
+    Returns:
+      List of strings corresponding to memories
+    """
+
+    with self._memory_bank_lock:
+      data = self._memory_bank[
+          (self._memory_bank['time'] >= time_from)
+          & (self._memory_bank['time'] <= time_until)
+      ]
+
+    return self._pd_to_text(data, add_time=add_time, sort_by_time=True)
+
   def retrieve_recent(
       self,
       k: int = 1,

concordia/components/agent/observation.py

@@ -15,6 +15,8 @@
 
 """Agent components for representing observation stream."""
 
+from collections.abc import Callable
+import datetime
 from concordia.associative_memory import associative_memory
 from concordia.document import interactive_document
 from concordia.language_model import language_model
@@ -23,126 +25,157 @@
 
 
 class Observation(component.Component):
-  """Component that stacks current observations together, clears on update."""
+  """Component that displays and adds observations to memory."""
 
   def __init__(
       self,
       agent_name: str,
+      clock_now: Callable[[], datetime.datetime],
+      timeframe: datetime.timedelta,
       memory: associative_memory.AssociativeMemory,
       component_name: str = 'Current observation',
       verbose: bool = False,
       log_colour='green',
   ):
-    """Initialize the observation component.
+    """Initializes the component.
 
     Args:
-      agent_name: the name of the agent
-      memory: memory for writing observations into
-      component_name: the name of this component
-      verbose: whether or not to print intermediate reasoning steps
-      log_colour: colour for logging
+      agent_name: Name of the agent.
+      clock_now: Function that returns the current time.
+      timeframe: Delta from current moment to display observations from, e.g. 1h
+        would display all observations made in the last hour.
+      memory: Associative memory to add and retrieve observations.
+      component_name: Name of this component.
+      verbose: Whether to print the observations.
+      log_colour: Colour to print the log.
     """
     self._agent_name = agent_name
     self._log_colour = log_colour
     self._name = component_name
     self._memory = memory
-
-    self._last_observation = []
+    self._timeframe = timeframe
+    self._clock_now = clock_now
 
     self._verbose = verbose
 
   def name(self) -> str:
     return self._name
 
   def state(self):
+    mems = self._memory.retrieve_time_interval(
+        self._clock_now() - self._timeframe, self._clock_now(), add_time=True
+    )
     if self._verbose:
-      self._log('\n'.join(self._last_observation) + '\n')
-    return '\n'.join(self._last_observation) + '\n'
+      self._log('\n'.join(mems) + '\n')
+    return '\n'.join(mems) + '\n'
 
   def _log(self, entry: str):
     print(termcolor.colored(entry, self._log_colour), end='')
 
   def observe(self, observation: str):
-    self._last_observation.append(observation)
     self._memory.add(
         f'[observation] {observation}',
         tags=['observation'],
     )
 
-  def update(self):
-    self._last_observation = []
-    return ''
-
 
 class ObservationSummary(component.Component):
-  """Component that summarises current observations on update."""
+  """Component that summarises observations from a segment of time."""
 
   def __init__(
       self,
-      model: language_model.LanguageModel,
       agent_name: str,
+      model: language_model.LanguageModel,
+      clock_now: Callable[[], datetime.datetime],
+      timeframe_delta_from: datetime.timedelta,
+      timeframe_delta_until: datetime.timedelta,
+      memory: associative_memory.AssociativeMemory,
       components: list[component.Component],
+      component_name: str = 'Summary of observations',
+      display_timeframe: bool = True,
       verbose: bool = False,
       log_colour='green',
   ):
-    """Summarize the agent's observations.
+    """Initializes the component.
 
     Args:
-      model: a language model
-      agent_name: the name of the agent
-      components: components to condition observation summarisation
-      verbose: whether or not to print intermediate reasoning steps
-      log_colour: colour for logging
+      agent_name: Name of the agent.
+      model: Language model to summarise the observations.
+      clock_now: Function that returns the current time.
+      timeframe_delta_from: delta from the current moment to the begnning of the
+        segment to summarise, e.g. 4h would summarise all observations that
+        happened from 4h ago intil clock_now minus timeframe_delta_until.
+      timeframe_delta_until: delta from the current moment to the end of the
+        segment to summarise.
+      memory: Associative memory retrieve observations from.
+      components: List of components to summarise.
+      component_name: Name of the component.
+      display_timeframe: Whether to display the time interval as text.
+      verbose: Whether to print the observations.
+      log_colour: Colour to print the log.
     """
     self._model = model
-    self._state = ''
     self._agent_name = agent_name
     self._log_colour = log_colour
+    self._name = component_name
+    self._memory = memory
+    self._timeframe_delta_from = timeframe_delta_from
+    self._timeframe_delta_until = timeframe_delta_until
+    self._clock_now = clock_now
     self._components = components
-
-    self._last_observation = []
+    self._state = ''
+    self._display_timeframe = display_timeframe
 
     self._verbose = verbose
 
   def name(self) -> str:
-    return 'Summary of recent observations'
+    return self._name
 
   def state(self):
     return self._state
 
   def _log(self, entry: str):
     print(termcolor.colored(entry, self._log_colour), end='')
 
-  def observe(self, observation: str):
-    self._last_observation.append(observation)
-
   def update(self):
-    context = '\n'.join(
-        [
-            f"{self._agent_name}'s "
-            + (comp.name() + ':\n' + comp.state())
-            for comp in self._components
-        ]
+    context = '\n'.join([
+        f"{self._agent_name}'s " + (comp.name() + ':\n' + comp.state())
+        for comp in self._components
+    ])
+
+    segment_start = self._clock_now() - self._timeframe_delta_from
+    segment_end = self._clock_now() - self._timeframe_delta_until
+
+    mems = self._memory.retrieve_time_interval(
+        segment_start,
+        segment_end,
+        add_time=True,
     )
 
-    numbered_observations = [
-        f'{i}. {observation}'
-        for i, observation in enumerate(self._last_observation)
-    ]
-    current_observations = '\n'.join(numbered_observations)
-
     prompt = interactive_document.InteractiveDocument(self._model)
     prompt.statement(context + '\n')
-    prompt.statement(
-        'Current observations, numbered in chronological order:\n'
-        + f'{current_observations}\n'
-    )
-    self._state = prompt.open_question(
-        'Summarize the observations into one sentence.'
+    prompt.statement(f'Recent memories of {self._agent_name}:\n' + f'{mems}\n')
+    self._state = (
+        self._agent_name
+        + ' '
+        + prompt.open_question(
+            'Summarize the memories above into one sentence about'
+            f' {self._agent_name}.',
+            answer_prefix=f'{self._agent_name} ',
+            max_characters=500,
+        )
     )
 
-    self._last_observation = []
+    if self._display_timeframe:
+      if segment_start.date() == segment_end.date():
+        interval = segment_start.strftime(
+            '%d %b %Y [%H:%M:%S  '
+        ) + segment_end.strftime('- %H:%M:%S]: ')
+      else:
+        interval = segment_start.strftime(
+            '[%d %b %Y %H:%M:%S  '
+        ) + segment_end.strftime('- %d %b %Y  %H:%M:%S]: ')
+      self._state = f'{interval} {self._state}'
 
     if self._verbose:
-      self._log('\nObservation summary:')
-      self._log('\n' + prompt.view().text() + '\n')
+      self._log(self._state)

concordia/components/game_master/conversation.py

@@ -16,6 +16,7 @@
 """Externality for the Game Master, which generates conversations."""
 
 from collections.abc import Sequence
+import datetime
 
 from concordia import components as generic_components
 from concordia.agents import basic_agent
@@ -123,7 +124,12 @@ def _make_npc(
             generic_components.constant.ConstantComponent(
                 name='General knowledge:', state=context
             ),
-            sim_components.observation.Observation(agent_name=name, memory=mem),
+            sim_components.observation.Observation(
+                agent_name=name,
+                memory=mem,
+                clock_now=scene_clock.now,
+                timeframe=datetime.timedelta(days=1),
+            ),
         ],
         verbose=True,
     )