Refactor ContextProcessorComponent to have one hook per phase.

PiperOrigin-RevId: 651805326
Change-Id: I0ddf3e5a7545554d03a6bc9c5626a29d376ddb91

concordia/agents/entity_agent.py

@@ -125,13 +125,14 @@ def act(
   ) -> str:
     self._phase = component_v2.Phase.PRE_ACT
     contexts = self._parallel_call_('pre_act', action_spec)
+    self._context_processor.pre_act(types.MappingProxyType(contexts))
     action_attempt = self._act_component.get_action_attempt(
         contexts, action_spec
     )
 
     self._phase = component_v2.Phase.POST_ACT
     contexts = self._parallel_call_('post_act', action_attempt)
-    self._context_processor.process(contexts)
+    self._context_processor.post_act(contexts)
 
     self._phase = component_v2.Phase.UPDATE
     self._parallel_call_('update')
@@ -142,11 +143,11 @@ def act(
   def observe(self, observation: str) -> None:
     self._phase = component_v2.Phase.PRE_OBSERVE
     contexts = self._parallel_call_('pre_observe', observation)
-    self._context_processor.process(contexts)
+    self._context_processor.pre_observe(contexts)
 
     self._phase = component_v2.Phase.POST_OBSERVE
     contexts = self._parallel_call_('post_observe')
-    self._context_processor.process(contexts)
+    self._context_processor.post_observe(contexts)
 
     self._phase = component_v2.Phase.UPDATE
     self._parallel_call_('update')

concordia/components/agent/v2/no_op_context_processor.py

@@ -16,15 +16,6 @@
 
 from concordia.typing import component_v2
 
-from typing_extensions import override
-
 
 class NoOpContextProcessor(component_v2.ContextProcessorComponent):
   """A context processor component that does nothing."""
-
-  @override
-  def process(
-      self,
-      contexts: component_v2.ComponentContextMapping,
-  ) -> None:
-    del contexts

concordia/typing/component_v2.py

@@ -219,20 +219,36 @@ def get_action_attempt(
 
 
 class ContextProcessorComponent(BaseComponent, metaclass=abc.ABCMeta):
-  """A component that processes context from components."""
+  """A component that processes context from ContextComponents."""
 
-  @abc.abstractmethod
-  def process(
-      self,
-      contexts: ComponentContextMapping,
-  ) -> None:
-    """Processes the context from ContextComponents.
+  def pre_act(self, contexts: ComponentContextMapping) -> None:
+    """Processes the pre_act contexts returned by the ContextComponents.
+
+    Args:
+      contexts: A mapping from ComponentName to ComponentContext.
+    """
+    del contexts
 
-    This function will be called by the entity with the context from other
-    components. The component should process the context and possibly update its
-    internal state or access other components.
+  def post_act(self, contexts: ComponentContextMapping) -> None:
+    """Processes the post_act contexts returned by the ContextComponents.
 
     Args:
-      contexts: The context from ContextComponents.
+      contexts: A mapping from ComponentName to ComponentContext.
     """
-    raise NotImplementedError()
+    del contexts
+
+  def pre_observe(self, contexts: ComponentContextMapping) -> None:
+    """Processes the pre_observe contexts returned by the ContextComponents.
+
+    Args:
+      contexts: A mapping from ComponentName to ComponentContext.
+    """
+    del contexts
+
+  def post_observe(self, contexts: ComponentContextMapping) -> None:
+    """Processes the post_observe contexts returned by the ContextComponents.
+
+    Args:
+      contexts: The context from other components.
+    """
+    del contexts