Add ComponentName typing annotation, instead of str, and improve …

…docsctrings. PiperOrigin-RevId: 653987315 Change-Id: Ifed1472778b2a739628f250c5f5f58a99261e43e
google-deepmind · Jul 19, 2024 · f699c2d · f699c2d
1 parent 66b4daa
commit f699c2d
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 10 deletions.
diff --git a/concordia/components/agent/v2/all_similar_memories.py b/concordia/components/agent/v2/all_similar_memories.py
@@ -23,6 +23,7 @@
 from concordia.document import interactive_document
 from concordia.language_model import language_model
 from concordia.memory_bank import legacy_associative_memory
+from concordia.typing import component_v2
 from concordia.typing import logging
 
 
@@ -37,7 +38,8 @@ def __init__(
       model: language_model.LanguageModel,
       memory_component_name: str = (
           memory_component.DEFAULT_MEMORY_COMPONENT_NAME),
-      components: Mapping[str, str] = types.MappingProxyType({}),
+      components: Mapping[
+          component_v2.ComponentName, str] = types.MappingProxyType({}),
       num_memories_to_retrieve: int = 25,
       pre_act_key: str = 'Relevant memories',
       logging_channel: logging.LoggingChannel = logging.NoOpLoggingChannel,
@@ -49,7 +51,7 @@ def __init__(
       memory_component_name: The name of the memory component from which to
         retrieve related memories.
       components: The components to condition the answer on. This is a mapping
-        of the component name to a prefix to use in the prompt.
+        of the component name to a label to use in the prompt.
       num_memories_to_retrieve: The number of memories to retrieve.
       pre_act_key: Prefix to add to the output of the component when called
         in `pre_act`.

diff --git a/concordia/components/agent/v2/person_by_situation.py b/concordia/components/agent/v2/person_by_situation.py
@@ -23,6 +23,7 @@
 from concordia.document import interactive_document
 from concordia.language_model import language_model
 from concordia.memory_bank import legacy_associative_memory
+from concordia.typing import component_v2
 import termcolor
 
 DEFAULT_PRE_ACT_KEY = (
@@ -37,7 +38,8 @@ def __init__(
       model: language_model.LanguageModel,
       memory_component_name: str = (
           memory_component.DEFAULT_MEMORY_COMPONENT_NAME),
-      components: Mapping[str, str] = types.MappingProxyType({}),
+      components: Mapping[
+          component_v2.ComponentName, str] = types.MappingProxyType({}),
       clock_now: Callable[[], datetime.datetime] | None = None,
       num_memories_to_retrieve: int = 25,
       pre_act_key: str = DEFAULT_PRE_ACT_KEY,
@@ -51,7 +53,7 @@ def __init__(
       memory_component_name: The name of the memory component from which to
         retrieve recent memories.
       components: The components to condition the answer on. This is a mapping
-        of the component name to a prefix to use in the prompt.
+        of the component name to a label to use in the prompt.
       clock_now: time callback to use.
       num_memories_to_retrieve: The number of recent memories to retrieve.
       pre_act_key: Prefix to add to the output of the component when called

diff --git a/concordia/components/agent/v2/plan.py b/concordia/components/agent/v2/plan.py
@@ -24,6 +24,7 @@
 from concordia.document import interactive_document
 from concordia.language_model import language_model
 from concordia.memory_bank import legacy_associative_memory
+from concordia.typing import component_v2
 import termcolor
 
 DEFAULT_PRE_ACT_KEY = 'Plan'
@@ -39,7 +40,8 @@ def __init__(
       observation_component_name: str,
       memory_component_name: str = (
           memory_component.DEFAULT_MEMORY_COMPONENT_NAME),
-      components: Mapping[str, str] = types.MappingProxyType({}),
+      components: Mapping[
+          component_v2.ComponentName, str] = types.MappingProxyType({}),
       clock_now: Callable[[], datetime.datetime] | None = None,
       goal_component_name: str | None = None,
       num_memories_to_retrieve: int = 10,
@@ -57,7 +59,7 @@ def __init__(
       memory_component_name: The name of the memory component from which to
         retrieve memories
       components: components to build the context of planning. This is a mapping
-        of the component name to a prefix to use in the prompt.
+        of the component name to a label to use in the prompt.
       clock_now: time callback to use for the state.
       goal_component_name: index into `components` to use to represent the goal
         of planning

diff --git a/concordia/components/agent/v2/self_perception.py b/concordia/components/agent/v2/self_perception.py
@@ -22,6 +22,7 @@
 from concordia.document import interactive_document
 from concordia.language_model import language_model
 from concordia.memory_bank import legacy_associative_memory
+from concordia.typing import component_v2
 from concordia.typing import logging
 
 DEFAULT_PRE_ACT_KEY = 'Who they are'
@@ -35,7 +36,8 @@ def __init__(
       model: language_model.LanguageModel,
       memory_component_name: str = (
           memory_component.DEFAULT_MEMORY_COMPONENT_NAME),
-      components: Mapping[str, str] = types.MappingProxyType({}),
+      components: Mapping[
+          component_v2.ComponentName, str] = types.MappingProxyType({}),
       num_memories_to_retrieve: int = 100,
       pre_act_key: str = DEFAULT_PRE_ACT_KEY,
       logging_channel: logging.LoggingChannel = logging.NoOpLoggingChannel,
@@ -47,7 +49,7 @@ def __init__(
       memory_component_name: The name of the memory component from which to
         retrieve recent memories.
       components: The components to condition the answer on. This is a mapping
-        of the component name to a prefix to use in the prompt.
+        of the component name to a label to use in the prompt.
       num_memories_to_retrieve: Number of memories to retrieve.
       pre_act_key: Prefix to add to the output of the component when called
         in `pre_act`.

diff --git a/concordia/components/agent/v2/situation_perception.py b/concordia/components/agent/v2/situation_perception.py
@@ -22,6 +22,7 @@
 from concordia.document import interactive_document
 from concordia.language_model import language_model
 from concordia.memory_bank import legacy_associative_memory
+from concordia.typing import component_v2
 from concordia.typing import logging
 
 DEFAULT_PRE_ACT_KEY = 'What situation they are in'
@@ -35,7 +36,8 @@ def __init__(
       model: language_model.LanguageModel,
       memory_component_name: str = (
           memory_component.DEFAULT_MEMORY_COMPONENT_NAME),
-      components: Mapping[str, str] = types.MappingProxyType({}),
+      components: Mapping[
+          component_v2.ComponentName, str] = types.MappingProxyType({}),
       clock_now: Callable[[], datetime.datetime] | None = None,
       num_memories_to_retrieve: int = 25,
       pre_act_key: str = DEFAULT_PRE_ACT_KEY,
@@ -48,7 +50,7 @@ def __init__(
       memory_component_name: The name of the memory component from which to
         retrieve recent memories.
       components: The components to condition the answer on. This is a mapping
-        of the component name to a prefix to use in the prompt.
+        of the component name to a label to use in the prompt.
       clock_now: time callback to use.
       num_memories_to_retrieve: The number of recent memories to retrieve.
       pre_act_key: Prefix to add to the output of the component when called