Skip to main content

API Reference

This reference documents the main classes and functions in the Airtrain Python API.

Core Components​

Agent​

The Agent class is the central component for building AI agents in Airtrain.

from airtrain.core import Agent
from airtrain.models import OpenAIModel

# Create a basic agent
agent = Agent(
    name="my_agent",
    model=OpenAIModel("gpt-4-turbo"),
    skills=["text_generation"],
    tools=[]
)

# Process a user query
response = agent.process("Tell me about AI agents")

Constructor Parameters​

ParameterTypeDescription
namestrUnique identifier for the agent
modelBaseModelThe AI model to use (OpenAI, Anthropic, etc.)
skillslistList of skills the agent can use
toolslistOptional list of tools the agent can use
memoryBaseMemoryOptional memory system for the agent
permissionslistOptional list of permissions the agent has

Methods​

MethodDescription
process(input)Process an input and return a response
add_skill(skill)Add a new skill to the agent
add_tool(tool)Add a new tool to the agent
set_memory(memory)Set the agent's memory system

Models​

Airtrain supports multiple AI models through a unified interface.

# OpenAI Models
from airtrain.models import OpenAIModel
model = OpenAIModel(
    model_name="gpt-4-turbo",
    temperature=0.7,
    max_tokens=2000
)

# Anthropic Models
from airtrain.models import AnthropicModel
model = AnthropicModel(
    model_name="claude-3-opus-20240229",
    temperature=0.5,
    max_tokens=4000
)

# Together AI Models
from airtrain.models import TogetherModel
model = TogetherModel(
    model_name="mistralai/Mixtral-8x7B-Instruct-v0.1",
    temperature=0.8,
    max_tokens=1500
)

# Fireworks Models
from airtrain.models import FireworksModel
model = FireworksModel(
    model_name="accounts/fireworks/models/mixtral-8x7b-instruct",
    temperature=0.5,
    max_tokens=2000
)

Skills​

Skills define the capabilities of an agent:

from airtrain.core.skills import TextGenerationSkill, MultiModalSkill

# Create skills
text_skill = TextGenerationSkill()
multimodal_skill = MultiModalSkill()

# Add to agent
agent.add_skill(text_skill)
agent.add_skill(multimodal_skill)

Tools​

Tools allow agents to interact with external systems:

from airtrain.core import Tool
from airtrain.core.schemas import InputSchema, OutputSchema
from pydantic import BaseModel, Field

# Define schemas
class CalculatorInput(InputSchema):
    expression: str = Field(description="Math expression to evaluate")

class CalculatorOutput(OutputSchema):
    result: float
    error: str = None

# Define function
def calculate(expression):
    try:
        result = eval(expression)
        return CalculatorOutput(result=result)
    except Exception as e:
        return CalculatorOutput(result=0, error=str(e))

# Create tool
calculator_tool = Tool(
    name="calculator",
    function=calculate,
    input_schema=CalculatorInput,
    output_schema=CalculatorOutput,
    description="Calculate mathematical expressions"
)

# Add to agent
agent.add_tool(calculator_tool)

Schemas​

Schemas define the structure of inputs and outputs:

from airtrain.core.schemas import InputSchema, OutputSchema
from pydantic import Field, BaseModel
from typing import List, Optional

# Input schema
class SearchQuery(InputSchema):
    query: str
    num_results: int = Field(default=5, ge=1, le=20)
    
# Output schema with nested model
class SearchResult(BaseModel):
    title: str
    url: str
    snippet: str

class SearchResponse(OutputSchema):
    results: List[SearchResult]
    total_found: int
    suggestion: Optional[str] = None

Memory Systems​

Memory systems allow agents to maintain state:

from airtrain.core.memory import SimpleMemory, VectorMemory
from airtrain.core import Agent

# Simple key-value memory
memory = SimpleMemory()
memory.store("user_name", "Alice")
memory.store("user_preferences", {"theme": "dark", "language": "en"})

# Vector memory for semantic search
vector_memory = VectorMemory(
    embedding_model="text-embedding-3-large"
)
vector_memory.store_text("project_info", "This project aims to build an AI assistant...")
vector_memory.store_text("user_background", "The user is a software developer...")

# Search memory
similar_items = vector_memory.search("What is the project about?", limit=3)

# Add to agent
agent = Agent(name="assistant", model=model, memory=memory)

Credentials​

Credentials manage API keys for different providers:

from airtrain.core.credentials import (
    OpenAICredentials,
    AnthropicCredentials,
    TogetherCredentials,
    FireworksCredentials
)

# Create and save credentials
openai_creds = OpenAICredentials(api_key="your-api-key")
openai_creds.save_to_file(".env")

# Load credentials from file
loaded_creds = OpenAICredentials.from_file(".env")
loaded_creds.load_to_env()

# Validate credentials
try:
    loaded_creds.validate_credentials()
    print("Credentials are valid")
except Exception as e:
    print(f"Invalid credentials: {e}")

Execution Components​

Execution Graph​

Execution graphs define the flow of agent operations:

from airtrain.core.execution import (
    ExecutionGraph,
    ExecutionNode,
    ConditionNode,
    ToolNode
)

# Create nodes
start_node = ExecutionNode(id="start", function=process_input)
decision_node = ConditionNode(
    id="check_intent",
    condition=lambda state: "search" in state.get("intent", "")
)
search_node = ToolNode(id="search", tool=search_tool)
respond_node = ExecutionNode(id="respond", function=generate_response)

# Connect nodes
start_node.connect(decision_node)
decision_node.connect_true(search_node)
decision_node.connect_false(respond_node)
search_node.connect(respond_node)

# Create graph
graph = ExecutionGraph(
    start_node=start_node,
    end_node=respond_node
)

# Execute graph
result = graph.execute(initial_state={"query": "What is Airtrain?"})

Permission System​

Permissions control what agents can access:

from airtrain.core.permissions import (
    Permission,
    PermissionSet,
    ResourcePermission
)

# Define permissions
file_read = Permission("file:read", "Can read files")
file_write = Permission("file:write", "Can write files")
api_call = Permission("api:call", "Can make API calls")

# Create permission set
permission_set = PermissionSet()
permission_set.add(file_read)
permission_set.add(api_call)

# Check permissions
if permission_set.has_permission("file:read"):
    print("Agent can read files")

if not permission_set.has_permission("file:write"):
    print("Agent cannot write files")

# Resource-specific permissions
docs_permission = ResourcePermission(
    "file:read",
    resource="documents/*",
    description="Can read document files"
)

Utilities​

Rate Limiters​

Rate limiters help manage API usage:

from airtrain.utils.rate_limiters import (
    SimpleRateLimiter,
    ExponentialBackoffLimiter
)

# Simple rate limiter
limiter = SimpleRateLimiter(
    requests_per_minute=60,
    max_retries=3
)

# Exponential backoff limiter
backoff_limiter = ExponentialBackoffLimiter(
    initial_delay=1,
    max_delay=60,
    max_retries=5
)

# Use with a model
model = OpenAIModel(
    model_name="gpt-4-turbo",
    rate_limiter=backoff_limiter
)

Logging​

Configure logging for Airtrain:

from airtrain.utils.logging import configure_logging

# Set up logging
configure_logging(
    level="INFO",
    file_path="airtrain.log",
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)

More Resources​

For more detailed information, check out: