Atomic Agents Quickstart

quickstart.py

import time
t0 = time.perf_counter()

from pydantic import Field
from openai import OpenAI
import instructor
from atomic_agents import AtomicAgent, AgentConfig, BasicChatInputSchema, BaseIOSchema
from atomic_agents.context import SystemPromptGenerator, ChatHistory
from dotenv import load_dotenv

def elapsed(label: str, since: float) -> float:
    now = time.perf_counter()
    print(f"[{now - since:.3f}s] {label}")
    return now

t = elapsed("imports loaded", t0)

load_dotenv()
t = elapsed("dotenv loaded", t0)

class CustomOutputSchema(BaseIOSchema):
    """
    docstring for the custom output schema
    """
    chat_message: str = Field(..., description="The chat message from the agent.")
    suggested_questions: list[str] = Field(..., description="Suggested follow-up questions.")

t = elapsed("CustomOutputSchema defined", t0)

system_prompt_generator = SystemPromptGenerator(
    background=["This assistant is knowledgeable, helpful, and suggests follow-up questions."],
    steps=[
        "Analyze the user's input to understand the context and intent.",
        "Formulate a relevant and informative response.",
        "Generate 3 suggested follow-up questions for the user."
    ],
    output_instructions=[
        "Provide clear and concise information in response to user queries.",
        "Conclude each response with 3 relevant suggested questions for the user."
    ]
)

t = elapsed("SystemPromptGenerator created", t0)

client = instructor.from_openai(OpenAI())
t = elapsed("OpenAI + instructor client initialized", t0)

agent = AtomicAgent[BasicChatInputSchema, CustomOutputSchema](
    config=AgentConfig(
        client=client,
        model="gpt-5-mini",
        system_prompt_generator=system_prompt_generator,
        history=ChatHistory(),
    )
)
t = elapsed("AtomicAgent initialized", t0)

if __name__ == "__main__":
    user_input = "Tell me about atomic agents framework"
    print(f"\nSending: '{user_input}'")

    t_request = time.perf_counter()
    response = agent.run(BasicChatInputSchema(chat_message=user_input))
    elapsed(f"agent.run() completed", t_request)
    elapsed(f"total time since script start", t0)

    print(f"\nAgent: {response.chat_message}")
    print("Suggested questions:")
    for question in response.suggested_questions:
        print(f"- {question}")