add pydantic and documentation to generate_object_agent.py and update generate_object_example.py

Author: createreadupdate

examples/generate_object_example.py

@@ -1,13 +1,17 @@
 import asyncio
-from core.generate_object_agent import GenerateObjectAgent
+from core.generate_object_agent import GenerateObjectAgent, ObjectGenerationInput
 
 async def run_generate_object_example():
-    description = "A machine that can sort fruits."
-    goal = "Generate a high-level design of the machine."
-    agent = GenerateObjectAgent(object_description=description, goal=goal)
+    input_data = ObjectGenerationInput(
+        object_description="A machine that can sort fruits.",
+        goal="Generate a high-level design of the machine.",
+        max_tokens=1000
+    )
+    
+    agent = GenerateObjectAgent(input_data)
     generated_object = await agent.generate_object()
 
-    print("Object description:", description)
+    print("Object description:", input_data.object_description)
     print("Generated object:", generated_object)
 
 if __name__ == "__main__":

src/core/generate_object_agent.py

@@ -1,15 +1,46 @@
-import asyncio
+from pydantic import BaseModel, Field
 from typing import Dict
 from .openai_api import OpenAIClient
 
+class ObjectGenerationInput(BaseModel):
+    object_description: str = Field(..., description="A description of the object to generate")
+    goal: str = Field(..., description="The goal of the generation process")
+    max_tokens: int = Field(1000, description="The maximum number of tokens to generate")
+
 class GenerateObjectAgent:
-    def __init__(self, object_description: str, goal: str, max_tokens: int = 1000):
-        self.object_description = object_description
-        self.goal = goal
-        self.max_tokens = max_tokens
+    """
+    A class to generate objects based on a given description and goal using the OpenAI API.
+
+    Attributes:
+        object_description (str): A description of the object to generate.
+        goal (str): The goal of the generation process.
+        max_tokens (int): The maximum number of tokens to generate.
+        openai_client (OpenAIClient): An instance of OpenAIClient to interact with the API.
+
+    Methods:
+        generate_object(): Generates an object based on the description and goal.
+    """
+
+    def __init__(self, data: ObjectGenerationInput):
+        """
+        Constructs all the necessary attributes for the GenerateObjectAgent object.
+
+        Args:
+            data (ObjectGenerationInput): An instance of ObjectGenerationInput containing 
+            the object description, goal, and max_tokens.
+        """
+        self.object_description = data.object_description
+        self.goal = data.goal
+        self.max_tokens = data.max_tokens
         self.openai_client = OpenAIClient()
 
     async def generate_object(self) -> Dict:
+        """
+        Generates an object based on the given description and goal.
+
+        Returns:
+            dict: A dictionary containing the original object description and the generated object.
+        """
         system_prompt = f"You are an assistant tasked with generating objects based on a given description. The goal is: {self.goal}."
         user_prompt = f"Generate an object based on the following description: {self.object_description}."