enhanced output formatting

src/agent_interfaces.py

@@ -5,17 +5,10 @@ from openai import OpenAI
 from colorama import Fore, Style
 from pydantic import BaseModel, ValidationError
 
+from output_formats import OutputFormat, OutputFormatModel
 from message import Message, AgentMessage
 from data_collection import save
 
-FORMAT_INSTRUCTIONS = """The output should be reformatted as a JSON instance that conforms to the JSON schema below.
-Here is the output schema:
-```
-{schema}
-```
-"""
-
-
 class BaseAgentInterface:
     """
     The interface that agents use to receive info from and interact with the game.
@@ -44,11 +37,11 @@ class BaseAgentInterface:
         self.add_message(response)
         return response
 
-    def respond_to_formatted(self, message: Message, output_format: Type[BaseModel], max_retries = 3) -> Type[BaseModel]:
+    def respond_to_formatted(self, message: Message, output_format: OutputFormat, max_retries = 3) -> OutputFormatModel:
         """Adds a message to the message history, and generates a response matching the provided format."""
         initial_response = self.respond_to(message)
 
-        reformat_message = Message(type="format", content=self._get_format_instructions(output_format))
+        reformat_message = Message(type="format", content=output_format.get_format_instructions())
 
         output = None
         retries = 0
@@ -56,11 +49,17 @@ class BaseAgentInterface:
         while not output and retries < max_retries:
             try:
                 formatted_response = self.respond_to(reformat_message)
-                output = output_format.model_validate_json(formatted_response.content)
+                output = output_format.output_format_model.model_validate_json(formatted_response.content)
             except ValidationError as e:
                 if retries > max_retries:
                     raise e
-                reformat_message = Message(type="retry", content=f"Error formatting response: {e} \n\n Please try again.")
+                retry_message = Message(type="retry", content=f"Error formatting response: {e} \n\n Please try again.")
+                if output_format.few_shot_examples:
+                    self.add_message(retry_message)
+                    reformat_message = Message(type="few_shot", content=output_format.get_few_shot()) # not implemented
+                else:
+                    reformat_message = retry_message
+
                 retries += 1
 
         return output
@@ -75,23 +74,6 @@ class BaseAgentInterface:
     def is_ai(self):
         return not self.is_human
 
-    # This should probably be put on a theoretical output format class...
-    # Or maybe on the Message class with a from format constructor
-    @staticmethod
-    def _get_format_instructions(output_format: Type[BaseModel]):
-        schema = output_format.model_json_schema()
-
-        reduced_schema = schema
-        if "title" in reduced_schema:
-            del reduced_schema["title"]
-        if "type" in reduced_schema:
-            del reduced_schema["type"]
-
-        schema_str = json.dumps(reduced_schema, indent=4)
-
-        return FORMAT_INSTRUCTIONS.format(schema=schema_str)
-
-
 AgentInterface = NewType("AgentInterface", BaseAgentInterface)
 
 

src/game.py

@@ -5,7 +5,7 @@ from typing import Optional, Type
 from colorama import Fore, Style
 
 from game_utils import *
-from models import *
+from output_formats import *
 from player import Player
 from prompts import fetch_prompt, format_prompt
 from message import Message
@@ -175,7 +175,7 @@ class Game:
 
             # Get Player Animal Description
             message = Message(type="prompt", content=prompt)
-            response = current_player.interface.respond_to_formatted(message, AnimalDescriptionModel)
+            response = current_player.interface.respond_to_formatted(message, OutputFormat(AnimalDescriptionFormat))
 
             self.player_responses.append({"sender": current_player.name, "response": response.description})
 
@@ -190,7 +190,7 @@ class Game:
         prompt = fetch_prompt("chameleon_guess_animal")
 
         message = Message(type="prompt", content=prompt)
-        response = chameleon.interface.respond_to_formatted(message, ChameleonGuessAnimalModel)
+        response = chameleon.interface.respond_to_formatted(message, OutputFormat(ChameleonGuessFormat))
 
         chameleon_animal_guess = response.animal
 
@@ -207,9 +207,8 @@ class Game:
 
                 # Get Player Vote
                 message = Message(type="prompt", content=prompt)
-                response = player.interface.respond_to_formatted(message, VoteModel)
-
-                # check if a valid player was voted for...
+                player_names = [p.name for p in self.players]
+                response = player.interface.respond_to_formatted(message, OutputFormat(HerdVoteFormat, player_names))
 
                 # Add Vote to Player Votes
                 player_votes.append({"voter": player, "vote": response.vote})

src/models.py

@@ -1,69 +0,0 @@
-import random
-from typing import Annotated, Type, List
-import json
-
-from pydantic import BaseModel, field_validator
-
-MAX_DESCRIPTION_LEN = 10
-
-
-class AnimalDescriptionModel(BaseModel):
-    # Define fields of our class here
-    description: Annotated[str, "A brief description of the animal"]
-
-    # @field_validator('description')
-    # @classmethod
-    # def check_starting_character(cls, v) -> str:
-    #     if not v[0].upper() == 'I':
-    #         raise ValueError("Description must begin with 'I'")
-    #     return v
-    #
-    # @field_validator('description')
-    # @classmethod
-    # def wordcount(cls, v) -> str:
-    #     count = len(v.split())
-    #     if count > MAX_DESCRIPTION_LEN:
-    #         raise ValueError(f"Animal Description must be {MAX_DESCRIPTION_LEN} words or less")
-    #     return v
-
-
-class ChameleonDecisionModel(BaseModel):
-    will_guess: bool
-
-
-class AnimalGuessModel(BaseModel):
-    animal_name: str
-
-
-class ChameleonGuessDecisionModel(BaseModel):
-    decision: Annotated[str, "Must be one of: ['guess', 'pass']"]
-
-    @field_validator('decision')
-    @classmethod
-    def check_decision(cls, v) -> str:
-        if v.lower() not in ['guess', 'pass']:
-            raise ValueError("Decision must be one of: ['guess', 'pass']")
-        return v
-
-
-class ChameleonGuessAnimalModel(BaseModel):
-    animal: Annotated[str, "The name of the animal the chameleon is guessing"]
-
-    @field_validator('animal')
-    @classmethod
-    def is_one_word(cls, v) -> str:
-        if len(v.split()) > 1:
-            raise ValueError("Animal's name must be one word")
-        return v
-
-
-class VoteModel(BaseModel):
-    vote: Annotated[str, "The name of the player you are voting for"]
-
-    # @field_validator('vote')
-    # @classmethod
-    # def check_player_exists(cls, v) -> str:
-    #     if v.lower() not in [player.lower() for player in players]:
-    #         raise ValueError(f"Player {v} does not exist")
-    #     return v
-

src/output_formats.py

@@ -0,0 +1,92 @@
+import random
+from typing import Annotated, NewType, List, Optional, Type, Literal
+import json
+
+from pydantic import BaseModel, field_validator, Field
+
+MAX_DESCRIPTION_LEN = 10
+FORMAT_INSTRUCTIONS = """Please reformat your previous response as a JSON instance that conforms to the JSON structure below.
+Here is the output format:
+{schema}
+"""
+FEW_SHOT_INSTRUCTIONS = """Here are a few examples of correctly formatted responses: \n
+{examples}
+"""
+
+OutputFormatModel = NewType("OutputFormatModel", BaseModel)
+
+
+class OutputFormat:
+    """The base class for all output formats."""
+
+    format_instructions: str = FORMAT_INSTRUCTIONS
+    """Instructions for formatting the output, it is combined with the JSON schema of the output format."""
+    few_shot_instructions: str = FEW_SHOT_INSTRUCTIONS
+    """Instructions for the few shot examples, it is combined with the few shot examples."""
+    few_shot_examples: Optional[List[dict]] = None
+    """A list of examples to be shown to the agent to help them understand the desired format of the output."""
+
+    def __init__(self, output_format_model: Type[OutputFormatModel], player_names: List[str] = None):
+        self.output_format_model = output_format_model
+        self.output_format_model.player_names = player_names
+
+    def get_format_instructions(self) -> str:
+        json_format = self.output_format_model().model_dump_json()
+
+        return self.format_instructions.format(schema=json_format)
+
+    def get_few_shot(self, max_examples=3):
+        if len(self.few_shot_examples) <= max_examples:
+            examples = self.few_shot_examples
+        else:
+            examples = random.sample(self.few_shot_examples, max_examples)
+
+        few_shot = "\n\n".join([f"Example Response:\n{json.dumps(example)}" for example in examples])
+
+        return self.few_shot_instructions.format(examples=few_shot)
+
+
+class AnimalDescriptionFormat(BaseModel):
+    # Define fields of our class here
+    description: str = Field("A brief description of the animal")
+    """A brief description of the animal"""
+
+    @field_validator('description')
+    @classmethod
+    def check_starting_character(cls, v) -> str:
+        if not v[0].upper() == 'I':
+            raise ValueError("Description must begin with 'I'")
+        return v
+
+    @field_validator('description')
+    @classmethod
+    def wordcount(cls, v) -> str:
+        count = len(v.split())
+        if count > MAX_DESCRIPTION_LEN:
+            raise ValueError(f"Animal Description must be {MAX_DESCRIPTION_LEN} words or less")
+        return v
+
+
+class ChameleonGuessFormat(BaseModel):
+    animal: str = Field("The name of the animal you think the chameleon is")
+
+    @field_validator('animal')
+    @classmethod
+    def is_one_word(cls, v) -> str:
+        if len(v.split()) > 1:
+            raise ValueError("Animal's name must be one word")
+        return v
+
+
+class HerdVoteFormat(BaseModel):
+    vote: str = Field("The name of the player you are voting for")
+    """The name of the player you are voting for"""
+    player_names: List[str] = Field([], exclude=True)
+    """The names of the players in the game"""
+
+    @field_validator('vote')
+    @classmethod
+    def check_player_exists(cls, v) -> str:
+        if v.lower() not in [player.lower() for player in cls.player_names]:
+            raise ValueError(f"Player {v} does not exist")
+        return v

src/prompts.py

@@ -1,66 +1,13 @@
-from models import *
-from langchain.prompts.few_shot import FewShotPromptTemplate
-from langchain.prompts.prompt import PromptTemplate
-
-
 def fetch_prompt(prompt_name):
     """Fetches a static prompt."""
     return prompts[prompt_name]
+
+
 def format_prompt(prompt_name, **kwargs):
     """Fetches a template prompt and populates it."""
     return fetch_prompt(prompt_name).format(**kwargs)
 
 
-class Task:
-    def __init__(self, prompt: str, response_format: Type[BaseModel], few_shot_examples: List[dict] = None):
-        self.prompt = prompt
-        self.response_format = response_format
-        self.few_shot_examples = few_shot_examples
-
-    def full_prompt(self, **kwargs):
-        prompt = self.prompt.format(**kwargs)
-
-        format_instructions = self.get_format_instructions()
-        if self.few_shot_examples:
-            few_shot = self.get_few_shot()
-
-    def get_format_instructions(self):
-        schema = self.get_input_schema()
-        format_instructions = FORMAT_INSTRUCTIONS.format(schema=schema)
-
-        return format_instructions
-
-    def get_input_schema(self):
-        schema = self.response_format.model_json_schema()
-
-        reduced_schema = schema
-        if "title" in reduced_schema:
-            del reduced_schema["title"]
-        if "type" in reduced_schema:
-            del reduced_schema["type"]
-
-        schema_str = json.dumps(reduced_schema, indent=4)
-
-        return schema_str
-
-    def get_few_shot(self, max_examples=3):
-        if len(self.few_shot_examples) <= max_examples:
-            examples = self.few_shot_examples
-        else:
-            examples = random.sample(self.few_shot_examples, max_examples)
-
-        few_shot = "\n\n".join([self.format_example(ex) for ex in examples])
-
-        return few_shot
-
-    def format_example(self, example):
-        ex_prompt = self.prompt.format(**example['inputs'])
-        ex_response = example['response']
-
-        return f"Prompt: {ex_prompt}\nResponse: {ex_response}"
-
-
-
 _game_rules = '''\
 You are playing a social deduction game where every player pretends the be the same animal. 
 During the round players go around the room and make an "I"-statement as if they were the animal.