[unstable] Improve more game master components especially 'next_action_spec'.

PiperOrigin-RevId: 735605779
Change-Id: Idcccdd9f6bb880881c6bfe7cd1c5d4b5d67a8217

Author: jzleibo

concordia/components/game_master/unstable/event_resolution.py

@@ -21,7 +21,7 @@
 from concordia.components.agent.unstable import memory as memory_component
 from concordia.document import interactive_document
 from concordia.language_model import language_model
-from concordia.thought_chains import thought_chains
+from concordia.thought_chains.unstable import thought_chains
 from concordia.typing import logging
 from concordia.typing.unstable import entity as entity_lib
 from concordia.typing.unstable import entity_component
@@ -98,7 +98,6 @@ def pre_act(
       action_spec: entity_lib.ActionSpec,
   ) -> str:
     result = ''
-    prompt_to_log = ''
     self._active_entity_name = None
     self._putative_action = None
     if action_spec.output_type == entity_lib.OutputType.RESOLVE:
@@ -124,10 +123,10 @@ def pre_act(
       result = f'{self._pre_act_key}: {event_statement}'
       prompt_to_log = prompt.view().text()
 
-    self._logging_channel(
-        {'Key': self._pre_act_key,
-         'Value': result,
-         'Prompt': prompt_to_log})
+      self._logging_channel(
+          {'Key': self._pre_act_key,
+           'Value': result,
+           'Prompt': prompt_to_log})
     return result
 
   def get_active_entity_name(self) -> str | None:

concordia/components/game_master/unstable/instructions.py

@@ -39,9 +39,11 @@
 )
 NEXT_ACTING_RESPONSE_EXAMPLE = 'Rowan'
 NEXT_ACTION_SPEC_RESPONSE_EXAMPLE_1 = (
-    'type: choice options: open the door, bar the door, flee'
+    'prompt: What would Kerensa do?;;type: choice;;options: '
+    'open the door, bar the door, flee'
 )
-NEXT_ACTION_SPEC_RESPONSE_EXAMPLE_2 = 'type: free'
+NEXT_ACTION_SPEC_RESPONSE_EXAMPLE_2 = (
+    'prompt: What would Morwenna say?;;type: free')
 RESOLVE_RESPONSE_PUTATIVE_ACTION_EXAMPLE = (
     'What is Yorik attempting to do?\n'
     'Yorik investigates the discarded machinery.')

concordia/components/game_master/unstable/next_acting.py

@@ -21,17 +21,25 @@
 from concordia.components.agent.unstable import memory as memory_component
 from concordia.document import interactive_document
 from concordia.environment.scenes.unstable import runner as scene_runner
+from concordia.environment.unstable import engine as engine_lib
 from concordia.language_model import language_model
 from concordia.typing import logging
 from concordia.typing.unstable import entity as entity_lib
 from concordia.typing.unstable import entity_component
+from concordia.typing.unstable import scene as scene_lib
 
 
 DEFAULT_NEXT_ACTING_COMPONENT_NAME = '__next_acting__'
 # Initiative is the Dungeons & Dragons term for the rule system that controls
 # turn taking.
 DEFAULT_NEXT_ACTING_PRE_ACT_KEY = '\nInitiative'
 
+DEFAULT_NEXT_ACTION_SPEC_COMPONENT_NAME = '__next_action_spec__'
+DEFAULT_NEXT_ACTION_SPEC_PRE_ACT_KEY = '\nType of action'
+
+_DEFAULT_CALL_TO_NEXT_ACTION_SPEC = (
+    'In what action spec format should {name} respond?')
+
 
 class NextActing(entity_component.ContextComponent):
   """A component that decides whose turn is next.
@@ -82,7 +90,6 @@ def pre_act(
       action_spec: entity_lib.ActionSpec,
   ) -> str:
     result = ''
-    prompt_to_log = ''
     if action_spec.output_type == entity_lib.OutputType.NEXT_ACTING:
       entity_name = self.get_entity().name
       prompt = interactive_document.InteractiveDocument(self._model)
@@ -96,12 +103,7 @@ def pre_act(
           question='Whose turn is next?',
           answers=self._player_names)
       result = self._player_names[idx]
-      prompt_to_log = prompt.view().text()
 
-    self._logging_channel(
-        {'Key': self._pre_act_key,
-         'Value': result,
-         'Prompt': prompt_to_log})
     return result
 
 
@@ -162,7 +164,6 @@ def pre_act(
       action_spec: entity_lib.ActionSpec,
   ) -> str:
     result = ''
-    prompt_to_log = ''
     if action_spec.output_type == entity_lib.OutputType.NEXT_ACTING:
       entity_name = self.get_entity().name
       prompt = interactive_document.InteractiveDocument(self._model)
@@ -178,10 +179,152 @@ def pre_act(
           question='Whose turn is next?',
           answers=scene_participants)
       result = scene_participants[idx]
-      prompt_to_log = prompt.view().text()
 
-    self._logging_channel(
-        {'Key': self._pre_act_key,
-         'Value': result,
-         'Prompt': prompt_to_log})
     return result
+
+
+class NextActionSpec(entity_component.ContextComponent):
+  """A component that decides whose turn is next.
+  """
+
+  def __init__(
+      self,
+      model: language_model.LanguageModel,
+      player_names: Sequence[str],
+      components: Mapping[
+          entity_component.ComponentName, str
+      ] = types.MappingProxyType({}),
+      call_to_next_action_spec: str = _DEFAULT_CALL_TO_NEXT_ACTION_SPEC,
+      pre_act_key: str = DEFAULT_NEXT_ACTION_SPEC_PRE_ACT_KEY,
+      logging_channel: logging.LoggingChannel = logging.NoOpLoggingChannel,
+  ):
+    """Initializes the component.
+
+    Args:
+      model: The language model to use for the component.
+      player_names: Names of players to choose from.
+      components: The components to condition the answer on. This is a mapping
+        of the component name to a label to use in the prompt.
+      call_to_next_action_spec: prompt to use for the game master to decide on
+        what action spec to use for the next turn. Will be formatted to
+        substitute {name} for the name of the player whose turn is next.
+      pre_act_key: Prefix to add to the output of the component when called
+        in `pre_act`.
+      logging_channel: The channel to use for debug logging.
+
+    Raises:
+      ValueError: If the component order is not None and contains duplicate
+        components.
+    """
+    super().__init__()
+    self._model = model
+    self._player_names = player_names
+    self._components = dict(components)
+    self._call_to_next_action_spec = call_to_next_action_spec
+    self._pre_act_key = pre_act_key
+    self._logging_channel = logging_channel
+
+  def _get_named_component_pre_act_value(self, component_name: str) -> str:
+    """Returns the pre-act value of a named component of the parent entity."""
+    return (
+        self.get_entity().get_component(
+            component_name, type_=action_spec_ignored.ActionSpecIgnored
+        ).get_pre_act_value()
+    )
+
+  def pre_act(
+      self,
+      action_spec: entity_lib.ActionSpec,
+  ) -> str:
+    result = ''
+    if action_spec.output_type == entity_lib.OutputType.NEXT_ACTION_SPEC:
+      entity_name = self.get_entity().name
+      prompt = interactive_document.InteractiveDocument(self._model)
+      component_states = '\n'.join([
+          f"{entity_name}'s"
+          f' {prefix}:\n{self._get_named_component_pre_act_value(key)}'
+          for key, prefix in self._components.items()
+      ])
+      prompt.statement(f'{component_states}\n')
+      idx = prompt.multiple_choice_question(
+          question='Whose turn is next?',
+          answers=self._player_names)
+      active_player = self._player_names[idx]
+      prompt.statement(
+          'Example formatted action specs:\n1). "prompt: p;;type: free"\n'
+          '2). "prompt: p;;type: choice;;options: x, y, z".\nNote that p is a '
+          'string of any length, typically a question, and x, y, z, etc are '
+          'multiple choice answer responses. For instance, a valid format '
+          'could be indicated as '
+          'prompt: Where will Edgar go?;;type: choice;;'
+          'options: home, London, Narnia, the third moon of Jupiter')
+      result = prompt.open_question(
+          question=self._call_to_next_action_spec.format(name=active_player),
+          max_tokens=512)
+
+    return result
+
+
+class NextActionSpecFromSceneSpec(entity_component.ContextComponent):
+  """A component that decides the next action spec using the current scene spec.
+  """
+
+  def __init__(
+      self,
+      scenes: Sequence[scene_lib.ExperimentalSceneSpec],
+      memory_component_name: str = (
+          memory_component.DEFAULT_MEMORY_COMPONENT_NAME
+      ),
+      pre_act_key: str = DEFAULT_NEXT_ACTION_SPEC_PRE_ACT_KEY,
+      logging_channel: logging.LoggingChannel = logging.NoOpLoggingChannel,
+  ):
+    """Initializes the component.
+
+    Args:
+      scenes: All scenes to be used in the episode.
+      memory_component_name: The name of the memory component.
+      pre_act_key: Prefix to add to the output of the component when called
+        in `pre_act`.
+      logging_channel: The channel to use for debug logging.
+
+    Raises:
+      ValueError: If the component order is not None and contains duplicate
+        components.
+    """
+    super().__init__()
+    self._memory_component_name = memory_component_name
+    self._pre_act_key = pre_act_key
+    self._logging_channel = logging_channel
+
+    # Extract all scene type specs from the provided scenes.
+    self._scene_type_specs = {}
+    for scene in scenes:
+      self._scene_type_specs[scene.scene_type.name] = scene.scene_type
+
+  def _get_named_component_pre_act_value(self, component_name: str) -> str:
+    """Returns the pre-act value of a named component of the parent entity."""
+    return (
+        self.get_entity().get_component(
+            component_name, type_=action_spec_ignored.ActionSpecIgnored
+        ).get_pre_act_value()
+    )
+
+  def _get_current_scene_type(self) -> scene_lib.ExperimentalSceneTypeSpec:
+    memory = self.get_entity().get_component(
+        self._memory_component_name, type_=memory_component.Memory
+    )
+    scene_type_str = scene_runner.get_current_scene_type(memory=memory)
+    scene_type_spec = self._scene_type_specs[scene_type_str]
+    return scene_type_spec
+
+  def pre_act(
+      self,
+      action_spec: entity_lib.ActionSpec,
+  ) -> str:
+    action_spec_string = ''
+    if action_spec.output_type == entity_lib.OutputType.NEXT_ACTION_SPEC:
+      scene_type_spec = self._get_current_scene_type()
+      action_spec = scene_type_spec.action_spec
+      action_spec_string = engine_lib.action_spec_to_string(action_spec)
+
+    return action_spec_string

concordia/components/game_master/unstable/switch_act.py

@@ -173,9 +173,9 @@ def _next_entity_action_spec(
       action_spec: entity_lib.ActionSpec) -> str:
     context = self._context_for_action(contexts)
     if DEFAULT_NEXT_ACTION_SPEC_COMPONENT_NAME in contexts:
-      # action_spec_string = _convert_to_string(
-      #     next_action_spec['scene_type'].action_spec)
-      result = ''
+      result = str(contexts[DEFAULT_NEXT_ACTION_SPEC_COMPONENT_NAME])
+      if not result:
+        result = f'prompt: {entity_lib.DEFAULT_CALL_TO_ACTION};;type: free'
       self._log(result, context)
     else:
       # YOLO case
@@ -185,12 +185,18 @@ def _next_entity_action_spec(
       # Then ask the GM to reformat their answer in whatever string format can
       # be used by the engine and its parser.
       chain_of_thought.statement(
-          'Example formatted action specs:\n"type: free"\n'
-          '"type: choice options: x, y, z"')
+          'Example formatted action specs:\n1). "prompt: p;;type: free"\n'
+          '2). "prompt: p;;type: choice;;options: x, y, z".\nNote that p is a '
+          'string of any length, typically a question, and x, y, z, etc are '
+          'multiple choice answer responses. For instance, a valid format '
+          'could be indicated as '
+          'prompt: Where will Edgar go?;;type: choice;;'
+          'options: home, London, Narnia, the third moon of Jupiter')
       next_action_spec_string = chain_of_thought.open_question(
-          question='Format the decision prompt type as an action spec.')
+          question='In what action spec format should the next player respond?')
       if 'type:' not in next_action_spec_string:
-        next_action_spec_string = 'type: free' + next_action_spec_string
+        next_action_spec_string = (
+            f'prompt: {entity_lib.DEFAULT_CALL_TO_ACTION};;type: free')
 
       result = next_action_spec_string
       self._log(result, chain_of_thought)

concordia/environment/unstable/engine.py

@@ -72,20 +72,45 @@ def run_loop(
 def action_spec_parser(next_action_spec_string: str) -> entity_lib.ActionSpec:
   """Parse the next action spec string into an action spec."""
   if 'type: free' in next_action_spec_string:
+    splits = next_action_spec_string.split(';;')
+    if len(splits) != 2:
+      call_to_action = entity_lib.DEFAULT_CALL_TO_ACTION
+    else:
+      call_to_action = splits[0].split('prompt: ')[1]
     return entity_lib.ActionSpec(
-        call_to_action=next_action_spec_string,
+        call_to_action=call_to_action,
         output_type=entity_lib.OutputType.FREE,
     )
   elif (
       'type: choice' in next_action_spec_string
       and 'options: ' in next_action_spec_string
   ):
+    splits = next_action_spec_string.split(';;')
+    if 'prompt: ' not in next_action_spec_string:
+      call_to_action = entity_lib.DEFAULT_CALL_TO_ACTION
+    else:
+      call_to_action = splits[0].split('prompt: ')[1]
     return entity_lib.ActionSpec(
-        call_to_action=next_action_spec_string,
+        call_to_action=call_to_action,
         output_type=entity_lib.OutputType.CHOICE,
         options=tuple(next_action_spec_string.split('options: ')[1].split(',')),
     )
   else:
     raise RuntimeError(
         'Invalid next action spec string: \"{}\"'.format(
             next_action_spec_string))
+
+
+def action_spec_to_string(action_spec: entity_lib.ActionSpec) -> str:
+  """Convert an action spec to a string."""
+  if action_spec.output_type == entity_lib.OutputType.FREE:
+    return f'prompt: {action_spec.call_to_action};;type: free'
+  elif action_spec.output_type == entity_lib.OutputType.CHOICE:
+    return (
+        f'prompt: {action_spec.call_to_action};;type: choice options: '
+        + ', '.join(action_spec.options)
+    )
+  else:
+    raise RuntimeError(
+        'Invalid action spec output type: \"{}\"'.format(
+            action_spec.output_type))

concordia/environment/unstable/engines/asynchronous.py

@@ -19,7 +19,6 @@
 from typing import Any
 
 from concordia.environment.unstable import engine as engine_lib
-from concordia.typing.unstable import agent as agent_lib
 from concordia.typing.unstable import entity as entity_lib
 
 
@@ -120,8 +119,8 @@ def terminate(self,
 
   def run_loop(
       self,
-      game_master: agent_lib.GenerativeAgent,
-      entities: Sequence[agent_lib.GenerativeAgent],
+      game_master: entity_lib.Entity,
+      entities: Sequence[entity_lib.Entity],
       premise: str = '',
       max_steps: int = 100,
       verbose: bool = False,