Entities can now be passed a custom Measurements object.

Preserves behavior in default case where no measurements object is passed in (creates a standard Measurements object).

PiperOrigin-RevId: 879629151
Change-Id: I385a1734986c89f6c3ef8c2391c9bd28e6cec648

Author: jzleibo

concordia/agents/entity_agent_with_logging.py

@@ -38,6 +38,7 @@ def __init__(
       context_components: Mapping[str, entity_component.ContextComponent] = (
           types.MappingProxyType({})
       ),
+      measurements: measurements_lib.Measurements | None = None,
   ):
     """Initializes the agent.
 
@@ -54,28 +55,42 @@ def __init__(
       context_processor: The component that will be used to process contexts. If
         None, a NoOpContextProcessor will be used.
       context_components: The ContextComponents that will be used by the agent.
+      measurements: Optional measurements instance to use for logging. Defaults
+        to a standard Measurements().
     """
     super().__init__(agent_name=agent_name,
                      act_component=act_component,
                      context_processor=context_processor,
                      context_components=context_components)
-    self._component_logging = measurements_lib.Measurements()
+    self._component_logging = (
+        measurements
+        if measurements is not None
+        else measurements_lib.Measurements()
+    )
 
     for component_name, component in self._context_components.items():
       if isinstance(component, entity_component.ComponentWithLogging):
         channel_name = component_name
         component.set_logging_channel(
-            self._component_logging.get_channel(channel_name).append
+            lambda datum, ch=channel_name: self._component_logging.publish_datum(
+                ch, datum
+            )
         )
     if isinstance(act_component, entity_component.ComponentWithLogging):
       act_component.set_logging_channel(
-          self._component_logging.get_channel('__act__').append
+          lambda datum: self._component_logging.publish_datum('__act__', datum)
       )
     if isinstance(context_processor, entity_component.ComponentWithLogging):
       context_processor.set_logging_channel(
-          self._component_logging.get_channel('__context_processor__').append
+          lambda datum: self._component_logging.publish_datum(
+              '__context_processor__', datum
+          )
       )
 
+  @property
+  def measurements(self) -> measurements_lib.Measurements:
+    return self._component_logging
+
   def get_all_logs(self):
     return self._component_logging.get_all_channels()
 

concordia/prefabs/entity/basic.py

@@ -189,6 +189,7 @@ def build(
         agent_name=entity_name,
         act_component=act_component,
         context_components=components_of_agent,
+        measurements=self.params.get('measurements'),
     )
 
     return agent

concordia/prefabs/entity/basic_scripted.py

@@ -158,6 +158,7 @@ def build(
         agent_name=entity_name,
         act_component=act_component,
         context_components=components_of_agent,
+        measurements=self.params.get('measurements'),
     )
 
     return agent

concordia/prefabs/entity/basic_with_plan.py

@@ -174,6 +174,7 @@ def build(
         agent_name=entity_name,
         act_component=act_component,
         context_components=components_of_agent,
+        measurements=self.params.get('measurements'),
     )
 
     return agent

concordia/prefabs/entity/conversational.py

@@ -206,6 +206,7 @@ def build(
         agent_name=entity_name,
         act_component=act_component,
         context_components=components_of_agent,
+        measurements=self.params.get('measurements'),
     )
 
     return agent

concordia/prefabs/entity/fake_assistant_with_configurable_system_prompt.py

@@ -90,6 +90,7 @@ def build(
         agent_name=entity_name,
         act_component=act_component,
         context_components=components_of_agent,
+        measurements=self.params.get('measurements'),
     )
 
     return agent

concordia/prefabs/entity/minimal.py

@@ -146,7 +146,7 @@ def build(
         agent_name=agent_name,
         act_component=act_component,
         context_components=components_of_agent,
+        measurements=self.params.get('measurements'),
     )
 
     return agent
-

concordia/prefabs/entity/puppet.py

@@ -81,4 +81,5 @@ def build(
         agent_name=name,
         act_component=act_comp,
         context_components=components,
+        measurements=self.params.get("measurements"),
     )

concordia/prefabs/entity/rational.py

@@ -185,4 +185,5 @@ def build(
         agent_name=entity_name,
         act_component=act_component,
         context_components=components_of_agent,
+        measurements=self.params.get('measurements'),
     )

concordia/prefabs/game_master/async_social_media.py

@@ -242,6 +242,7 @@ def build(
         agent_name=name,
         act_component=act_component,
         context_components=components_of_game_master,
+        measurements=self.params.get('measurements'),
     )
 
     return game_master