feat: agents

[ghost]

Summary

This PR introduces comprehensive Agent functionality to the AI21 Python SDK, providing full CRUD operations and execution capabilities for AI agents through direct Maestro integration.

Key Features Added:

Agent Management:

• Create, retrieve, update, delete agents through Studio API
• Support for different agent types (default, chat, maestro)
• Agent configuration with budget levels, visibility, and requirements
• Both sync and async client support

Agent Execution via Maestro Integration:

• Direct delegation to Maestro client for agent run execution
• Automatic parameter conversion and agent configuration merging
• Support for verbose output, structured RAG, and dynamic planning
• Full run status tracking and retrieval through maestro

Clean Architecture:

• Agents resource exposes maestro runs directly through composition
• Agent configuration is automatically fetched and merged with run parameters
• Type-safe parameter conversion with automatic validation
• No duplication of execution logic - pure delegation to maestro

Usage Examples:

Agent CRUD Operations:

from ai21 import AI21Client
from ai21.models.agents import BudgetLevel, AgentType

client = AI21Client()

# Create
agent = client.beta.agents.create(
    name="Research Assistant",
    description="Specialized in research tasks",
    budget=BudgetLevel.HIGH,
)

# Read
retrieved_agent = client.beta.agents.get(agent.id)
agents_list = client.beta.agents.list()

# Update
modified_agent = client.beta.agents.modify(
    agent.id,
    name="Enhanced Research Assistant",
    description="Updated with enhanced capabilities",
)

# Delete
delete_response = client.beta.agents.delete(agent.id)

Agent Execution:

# Basic agent run
run_response = client.beta.agents.runs.create_and_poll(
    agent_id="your-agent-id",
    input=[{"role": "user", "content": "What is 2+2?"}],
    poll_timeout_sec=120,
)

# Advanced agent run with options
run_response = client.beta.agents.runs.create_and_poll(
    agent_id="your-agent-id",
    input=[
        {"role": "user", "content": "Analyze this data and provide insights"}
    ],
    verbose=True,
    include=["data_sources", "requirements_result"],
    structured_rag_enabled=True,
    dynamic_planning_enabled=True,
    response_language="english",
    poll_timeout_sec=300,
)

print(f"Result: {run_response.result}")

Async Usage:

import asyncio
from ai21 import AsyncAI21Client
from ai21.models.agents import BudgetLevel

async def main():
    client = AsyncAI21Client()
    
    # Create and run agent
    agent = await client.beta.agents.create(
        name="Async Agent",
        description="An agent for async operations",
        budget=BudgetLevel.HIGH,
    )
    
    # Execute agent run
    run_response = await client.beta.agents.runs.create_and_poll(
        agent.id,
        input=[{"role": "user", "content": "Process this request asynchronously"}],
        poll_timeout_sec=180,
    )
    
    # Clean up
    await client.beta.agents.delete(agent.id)

asyncio.run(main())

Technical Implementation:

Core Components:

• Agent CRUD: Full lifecycle management through Studio API (agents.py, async_agents.py)
• Direct Maestro Integration: Agents resource exposes maestro runs API directly
• Parameter Conversion: AgentToMaestroRunConverter handles automatic parameter transformation
• Utility Functions: Helper functions for agent-to-maestro parameter conversion

API Endpoints:

• client.beta.agents.create() - Create new agents
• client.beta.agents.get() - Retrieve agent by ID
• client.beta.agents.list() - List all agents
• client.beta.agents.modify() - Update existing agents
• client.beta.agents.delete() - Delete agents
• client.beta.agents.runs.* - Direct access to maestro runs API with automatic agent integration

Parameter Conversion Flow:

1. Agent configuration fetched from Studio API
2. AgentToMaestroRunConverter converts agent params + run params → maestro params
3. Execution delegated directly to maestro runs API
4. Results returned using maestro's RunResponse model

Benefits:

Maintainability: Single source of truth for run execution logic in maestro
Type Safety: Pydantic handles all conversions with validation
Consistency: Same run behavior whether using agents or maestro directly
Simplicity: Clean delegation pattern with automatic parameter mapping
Extensibility: Easy to add new agent parameters through the conversion schema

Documentation Updates:

This PR includes comprehensive documentation in the README.md file under the "Agents (Beta)" section, providing:

• Complete usage examples for both sync and async operations
• CRUD operation examples
• Agent execution examples with various options
• Links to detailed example files

Examples Added:

• examples/studio/agents/agent_crud.py - Complete CRUD operations
• examples/studio/agents/agent_run.py - Basic agent execution
• examples/studio/agents/async_agent_run.py - Async agent operations

This implementation provides a unified agent experience while leveraging the robust maestro execution engine underneath.

[Josephasafg]

Missing type annotations

[Josephasafg]

Is this comment necessary?

[Josephasafg]

Can you please use this syntax str | NotGiven = NOT_GIVEN?

[Josephasafg]

You might need to add from __future__ import annotations at the top of the file to make it work for older versions

[Josephasafg]

@benshuk PTAL as well