From f699c2d003c84cd6df53e83bc599c237984964d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edgar=20Du=C3=A9=C3=B1ez-Guzm=C3=A1n?= Date: Fri, 19 Jul 2024 06:20:58 -0700 Subject: [PATCH] Add `ComponentName` typing annotation, instead of `str`, and improve docsctrings. PiperOrigin-RevId: 653987315 Change-Id: Ifed1472778b2a739628f250c5f5f58a99261e43e --- concordia/components/agent/v2/all_similar_memories.py | 6 ++++-- concordia/components/agent/v2/person_by_situation.py | 6 ++++-- concordia/components/agent/v2/plan.py | 6 ++++-- concordia/components/agent/v2/self_perception.py | 6 ++++-- concordia/components/agent/v2/situation_perception.py | 6 ++++-- 5 files changed, 20 insertions(+), 10 deletions(-) diff --git a/concordia/components/agent/v2/all_similar_memories.py b/concordia/components/agent/v2/all_similar_memories.py index 5d5bdedf..74362cec 100644 --- 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 index 7c8cebfc..b79a8cca 100644 --- 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 index 05224c5f..efe9813e 100644 --- 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 index 3bdd2f79..a5e83f76 100644 --- 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 index 202e101c..87d64aa0 100644 --- 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