You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
swarms/playground/tools/simple_structured_output.py

68 lines
1.9 KiB

"""
* WORKING
What this script does:
Structured output example
Requirements:
Add the folowing API key(s) in your .env file:
- OPENAI_API_KEY (this example works best with Openai bc it uses openai function calling structure)
"""
################ Adding project root to PYTHONPATH ################################
# If you are running playground examples in the project files directly, use this:
import sys
import os
sys.path.insert(0, os.getcwd())
################ Adding project root to PYTHONPATH ################################
from pydantic import BaseModel, Field
from swarms import Agent, OpenAIChat
# Initialize the schema for the person's information
class PersonInfo(BaseModel):
"""
To create a PersonInfo, you need to return a JSON object with the following format:
{
"function_call": "PersonInfo",
"parameters": {
...
}
}
"""
name: str = Field(..., title="Name of the person")
age: int = Field(..., title="Age of the person")
is_student: bool = Field(..., title="Whether the person is a student")
courses: list[str] = Field(
..., title="List of courses the person is taking"
)
# Initialize the agent
agent = Agent(
agent_name="Person Information Generator",
system_prompt=(
"Generate a person's information"
),
llm=OpenAIChat(),
max_loops=1,
verbose=True,
# List of pydantic models that the agent can use
list_base_models=[PersonInfo],
output_validation=True
)
# Define the task to generate a person's information
task = "Generate a person's information for Paul Graham 56 years old and is a student at Harvard University and is taking 3 courses: Math, Science, and History."
# Run the agent to generate the person's information
generated_data = agent.run(task)
# Print the generated data
print(type(generated_data))
print(f"Generated data: {generated_data}")