Merge main into sweep/add-stale-action

gpt_engineer/ai.py

@@ -38,6 +38,16 @@ class TokenUsage:
 
 class AI:
     def __init__(self, model_name="gpt-4", temperature=0.1):
+        """
+        Initialize the AI class.
+
+        Parameters
+        ----------
+        model_name : str, optional
+            The name of the model to use, by default "gpt-4".
+        temperature : float, optional
+            The temperature to use for the model, by default 0.1.
+        """
         self.temperature = temperature
         self.model_name = fallback_model(model_name)
         self.llm = create_chat_model(self.model_name, temperature)
@@ -50,19 +60,75 @@ def __init__(self, model_name="gpt-4", temperature=0.1):
         self.token_usage_log = []
 
     def start(self, system: str, user: str, step_name: str) -> List[Message]:
+        """
+        Start the conversation with a system message and a user message.
+    
+        Parameters
+        ----------
+        system : str
+            The content of the system message.
+        user : str
+            The content of the user message.
+        step_name : str
+            The name of the step.
+    
+        Returns
+        -------
+        List[Message]
+            The list of messages in the conversation.
+        """
         messages: List[Message] = [
             SystemMessage(content=system),
             HumanMessage(content=user),
         ]
         return self.next(messages, step_name=step_name)
 
     def fsystem(self, msg: str) -> SystemMessage:
+        """
+        Create a system message.
+    
+        Parameters
+        ----------
+        msg : str
+            The content of the message.
+    
+        Returns
+        -------
+        SystemMessage
+            The created system message.
+        """
         return SystemMessage(content=msg)
 
     def fuser(self, msg: str) -> HumanMessage:
+        """
+        Create a user message.
+    
+        Parameters
+        ----------
+        msg : str
+            The content of the message.
+    
+        Returns
+        -------
+        HumanMessage
+            The created user message.
+        """
         return HumanMessage(content=msg)
 
     def fassistant(self, msg: str) -> AIMessage:
+        """
+        Create an AI message.
+    
+        Parameters
+        ----------
+        msg : str
+            The content of the message.
+    
+        Returns
+        -------
+        AIMessage
+            The created AI message.
+        """
         return AIMessage(content=msg)
 
     def next(
@@ -75,6 +141,24 @@ def next(
         """
         Advances the conversation by sending message history
         to LLM and updating with the response.
+    
+        Parameters
+        ----------
+        messages : List[Message]
+            The list of messages in the conversation.
+        prompt : Optional[str], optional
+            The prompt to use, by default None.
+        step_name : str
+            The name of the step.
+    
+        Returns
+        -------
+        List[Message]
+            The updated list of messages in the conversation.
+        """
+        """
+        Advances the conversation by sending message history
+        to LLM and updating with the response.
         """
         if prompt:
             messages.append(self.fuser(prompt))
@@ -95,15 +179,53 @@ def next(
 
     @staticmethod
     def serialize_messages(messages: List[Message]) -> str:
+        """
+        Serialize a list of messages to a JSON string.
+    
+        Parameters
+        ----------
+        messages : List[Message]
+            The list of messages to serialize.
+    
+        Returns
+        -------
+        str
+            The serialized messages as a JSON string.
+        """
         return json.dumps(messages_to_dict(messages))
 
     @staticmethod
     def deserialize_messages(jsondictstr: str) -> List[Message]:
+        """
+        Deserialize a JSON string to a list of messages.
+    
+        Parameters
+        ----------
+        jsondictstr : str
+            The JSON string to deserialize.
+    
+        Returns
+        -------
+        List[Message]
+            The deserialized list of messages.
+        """
         return list(messages_from_dict(json.loads(jsondictstr)))  # type: ignore
 
     def update_token_usage_log(
         self, messages: List[Message], answer: str, step_name: str
     ) -> None:
+        """
+        Update the token usage log with the number of tokens used in the current step.
+    
+        Parameters
+        ----------
+        messages : List[Message]
+            The list of messages in the conversation.
+        answer : str
+            The answer from the AI.
+        step_name : str
+            The name of the step.
+        """
         prompt_tokens = self.num_tokens_from_messages(messages)
         completion_tokens = self.num_tokens(answer)
         total_tokens = prompt_tokens + completion_tokens
@@ -125,6 +247,14 @@ def update_token_usage_log(
         )
 
     def format_token_usage_log(self) -> str:
+        """
+        Format the token usage log as a CSV string.
+    
+        Returns
+        -------
+        str
+            The token usage log formatted as a CSV string.
+        """
         result = "step_name,"
         result += "prompt_tokens_in_step,completion_tokens_in_step,total_tokens_in_step"
         result += ",total_prompt_tokens,total_completion_tokens,total_tokens\n"
@@ -139,9 +269,35 @@ def format_token_usage_log(self) -> str:
         return result
 
     def num_tokens(self, txt: str) -> int:
+        """
+        Get the number of tokens in a text.
+    
+        Parameters
+        ----------
+        txt : str
+            The text to count the tokens in.
+    
+        Returns
+        -------
+        int
+            The number of tokens in the text.
+        """
         return len(self.tokenizer.encode(txt))
 
     def num_tokens_from_messages(self, messages: List[Message]) -> int:
+        """
+        Get the total number of tokens used by a list of messages.
+    
+        Parameters
+        ----------
+        messages : List[Message]
+            The list of messages to count the tokens in.
+    
+        Returns
+        -------
+        int
+            The total number of tokens used by the messages.
+        """
         """Returns the number of tokens used by a list of messages."""
         n_tokens = 0
         for message in messages:
@@ -154,6 +310,19 @@ def num_tokens_from_messages(self, messages: List[Message]) -> int:
 
 
 def fallback_model(model: str) -> str:
+    """
+    Retrieve the specified model, or fallback to "gpt-3.5-turbo" if the model is not available.
+
+    Parameters
+    ----------
+    model : str
+        The name of the model to retrieve.
+
+    Returns
+    -------
+    str
+        The name of the retrieved model, or "gpt-3.5-turbo" if the specified model is not available.
+    """
     try:
         openai.Model.retrieve(model)
         return model
@@ -167,6 +336,21 @@ def fallback_model(model: str) -> str:
 
 
 def create_chat_model(model: str, temperature) -> BaseChatModel:
+    """
+    Create a chat model with the specified model name and temperature.
+
+    Parameters
+    ----------
+    model : str
+        The name of the model to create.
+    temperature : float
+        The temperature to use for the model.
+
+    Returns
+    -------
+    BaseChatModel
+        The created chat model.
+    """
     if model == "gpt-4":
         return ChatOpenAI(
             model="gpt-4",
@@ -186,6 +370,19 @@ def create_chat_model(model: str, temperature) -> BaseChatModel:
 
 
 def get_tokenizer(model: str):
+    """
+    Get the tokenizer for the specified model.
+
+    Parameters
+    ----------
+    model : str
+        The name of the model to get the tokenizer for.
+
+    Returns
+    -------
+    Tokenizer
+        The tokenizer for the specified model.
+    """
     if "gpt-4" in model or "gpt-3.5" in model:
         return tiktoken.encoding_for_model(model)