Skip to main content

OpenAI Integration

Airtrain provides seamless integration with OpenAI's models, allowing you to use them in your AI agents. This guide will show you how to set up and use OpenAI models with Airtrain.

Setting Up OpenAI Credentials​

First, you'll need to set up your OpenAI credentials:

from airtrain.integrations.openai.credentials import OpenAICredentials
import os

# Create credentials
openai_creds = OpenAICredentials(
    openai_api_key=os.getenv("OPENAI_API_KEY"),
    organization_id=os.getenv("OPENAI_ORG_ID"),  # Optional
)

# Save to file
openai_creds.save_to_file("credentials.env")

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

Creating an OpenAI-based Agent​

Once you have your credentials set up, you can create an agent that uses OpenAI models:

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

# Create a model
openai_model = OpenAIModel(
    model_name="gpt-4-turbo",
    temperature=0.7,
    max_tokens=2000
)

# Create an agent with the model
agent = Agent(
    name="content_creator",
    model=openai_model,
    skills=[TextGenerationSkill()]
)

# Use the agent
response = agent.process("Write a short blog post about AI agents")
print(response.content)

Advanced Usage: Function Calling​

OpenAI models support function calling, which you can use with Airtrain to create more structured outputs:

from airtrain.core import Agent, Tool
from airtrain.models import OpenAIModel
from airtrain.core.schemas import InputSchema, OutputSchema
from pydantic import BaseModel, Field
from typing import List, Optional

# Define schemas for the function
class WeatherInputSchema(InputSchema):
    location: str
    unit: str = Field(default="celsius", description="Temperature unit")
    days: int = Field(default=1, ge=1, le=7, description="Number of days to forecast")

class WeatherDataSchema(BaseModel):
    day: str
    temperature: float
    conditions: str
    precipitation_chance: float

class WeatherOutputSchema(OutputSchema):
    location: str
    forecast: List[WeatherDataSchema]
    alerts: Optional[List[str]] = None

# Create a weather tool
def get_weather(location, unit="celsius", days=1):
    """Mock function to get weather data"""
    import random
    conditions = ["Sunny", "Cloudy", "Rainy", "Snowy"]
    forecast = []
    
    for i in range(days):
        day = f"Day {i+1}"
        temp = random.uniform(-5, 30) if unit == "celsius" else random.uniform(20, 90)
        condition = random.choice(conditions)
        precip = random.uniform(0, 1)
        
        forecast.append(WeatherDataSchema(
            day=day,
            temperature=round(temp, 1),
            conditions=condition,
            precipitation_chance=round(precip, 2)
        ))
    
    return WeatherOutputSchema(
        location=location,
        forecast=forecast,
        alerts=["Heavy rain expected"] if any(f.conditions == "Rainy" for f in forecast) else None
    )

# Create the tool
weather_tool = Tool(
    name="get_weather",
    function=get_weather,
    input_schema=WeatherInputSchema,
    output_schema=WeatherOutputSchema,
    description="Get weather forecast for a location"
)

# Create the agent with the tool
agent = Agent(
    name="weather_assistant",
    model=OpenAIModel("gpt-4-turbo"),
    tools=[weather_tool]
)

# Use the agent
user_query = "What's the weather like in New York for the next 3 days?"
response = agent.process(user_query)

# The agent will automatically call the weather tool with appropriate parameters
print(response.content)

Using Different OpenAI Models​

Airtrain supports all OpenAI models. Here's how to use different models for different purposes:

from airtrain.models import OpenAIModel

# For text generation
gpt4_model = OpenAIModel("gpt-4-turbo")

# For embedding generation
embedding_model = OpenAIModel("text-embedding-3-large")

# For image generation
dalle_model = OpenAIModel("dall-e-3")

# With specific parameters
custom_model = OpenAIModel(
    model_name="gpt-4-turbo",
    temperature=0.2,  # More deterministic
    top_p=0.9,
    frequency_penalty=0.5,
    presence_penalty=0.5,
    max_tokens=4000
)

Handling Rate Limits​

When working with OpenAI APIs, you might encounter rate limits. Airtrain provides built-in rate limit handling:

from airtrain.core import Agent
from airtrain.models import OpenAIModel
from airtrain.utils.rate_limiters import ExponentialBackoffLimiter

# Create a model with rate limiting
openai_model = OpenAIModel(
    model_name="gpt-4-turbo",
    rate_limiter=ExponentialBackoffLimiter(
        initial_delay=1,
        max_delay=60,
        max_retries=5
    )
)

# Use the model in an agent
agent = Agent(name="rate_limited_agent", model=openai_model)

Next Steps​

Now that you've set up OpenAI integration, you can: