airtrain 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

airtrain/integrations/openai/credentials.py

@@ -0,0 +1,39 @@
+from datetime import datetime, timedelta
+from typing import Optional
+from pydantic import Field, SecretStr, validator
+from openai import OpenAI
+
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+
+
+class OpenAICredentials(BaseCredentials):
+    """OpenAI API credentials with enhanced validation"""
+
+    openai_api_key: SecretStr = Field(..., description="OpenAI API key")
+    openai_organization_id: Optional[str] = Field(
+        None, description="OpenAI organization ID", pattern="^org-[A-Za-z0-9]{24}$"
+    )
+
+    _required_credentials = {"openai_api_key"}
+
+    @validator("openai_api_key")
+    def validate_api_key_format(cls, v: SecretStr) -> SecretStr:
+        key = v.get_secret_value()
+        if not key.startswith("sk-"):
+            raise ValueError("OpenAI API key must start with 'sk-'")
+        if len(key) < 40:
+            raise ValueError("OpenAI API key appears to be too short")
+        return v
+
+    async def validate_credentials(self) -> bool:
+        """Validate credentials by making a test API call"""
+        try:
+            client = OpenAI(
+                api_key=self.openai_api_key.get_secret_value(),
+                organization=self.openai_organization_id,
+            )
+            # Make minimal API call to validate
+            await client.models.list(limit=1)
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid OpenAI credentials: {str(e)}")

airtrain/integrations/openai/list_models.py

@@ -0,0 +1,112 @@
+from typing import Optional, List, Dict, Any
+from pydantic import Field
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import OpenAICredentials
+from .models_config import OPENAI_MODELS, OpenAIModelConfig
+
+
+class OpenAIModel:
+    """Class to represent an OpenAI model."""
+    
+    def __init__(self, model_id: str, config: OpenAIModelConfig):
+        """Initialize the OpenAI model."""
+        self.id = model_id
+        self.display_name = config.display_name
+        self.base_model = config.base_model
+        self.input_price = config.input_price
+        self.cached_input_price = config.cached_input_price
+        self.output_price = config.output_price
+    
+    def dict(self, exclude_none=False):
+        """Convert the model to a dictionary."""
+        result = {
+            "id": self.id,
+            "display_name": self.display_name,
+            "base_model": self.base_model,
+            "input_price": float(self.input_price),
+            "output_price": float(self.output_price),
+        }
+        if self.cached_input_price is not None:
+            result["cached_input_price"] = float(self.cached_input_price)
+        elif not exclude_none:
+            result["cached_input_price"] = None
+        return result
+
+
+class OpenAIListModelsInput(InputSchema):
+    """Schema for OpenAI list models input"""
+    
+    api_models_only: bool = Field(
+        default=False,
+        description=(
+            "If True, fetch models from the API only. If False, use local config."
+        )
+    )
+
+
+class OpenAIListModelsOutput(OutputSchema):
+    """Schema for OpenAI list models output"""
+    
+    models: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="List of OpenAI models"
+    )
+
+
+class OpenAIListModelsSkill(Skill[OpenAIListModelsInput, OpenAIListModelsOutput]):
+    """Skill for listing OpenAI models"""
+
+    input_schema = OpenAIListModelsInput
+    output_schema = OpenAIListModelsOutput
+    
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials
+    
+    def process(
+        self, input_data: OpenAIListModelsInput
+    ) -> OpenAIListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            models = []
+            
+            if input_data.api_models_only:
+                # Fetch models from OpenAI API - requires credentials
+                if not self.credentials:
+                    raise ProcessingError(
+                        "OpenAI credentials required for API models"
+                    )
+                
+                from openai import OpenAI
+                client = OpenAI(
+                    api_key=self.credentials.openai_api_key.get_secret_value(),
+                    organization=self.credentials.openai_organization_id,
+                )
+                
+                # Make API call to get models
+                response = client.models.list()
+                
+                # Convert response to our format
+                for model in response.data:
+                    models.append({
+                        "id": model.id,
+                        "display_name": model.id,  # API doesn't provide display_name
+                        "base_model": model.id,    # API doesn't provide base_model
+                        "created": model.created,
+                        "owned_by": model.owned_by,
+                        # Pricing info not available from API
+                    })
+            else:
+                # Use local model config - no credentials needed
+                for model_id, config in OPENAI_MODELS.items():
+                    model = OpenAIModel(model_id, config)
+                    models.append(model.dict())
+            
+            # Return the output
+            return OpenAIListModelsOutput(models=models)
+            
+        except Exception as e:
+            raise ProcessingError(f"Failed to list OpenAI models: {str(e)}") 

airtrain/integrations/openai/models_config.py

@@ -0,0 +1,224 @@
+from typing import Dict, NamedTuple, Optional
+from decimal import Decimal
+
+
+class OpenAIModelConfig(NamedTuple):
+    display_name: str
+    base_model: str
+    input_price: Decimal
+    cached_input_price: Optional[Decimal]
+    output_price: Decimal
+
+
+OPENAI_MODELS: Dict[str, OpenAIModelConfig] = {
+    "gpt-4.5-preview": OpenAIModelConfig(
+        display_name="GPT-4.5 Preview",
+        base_model="gpt-4.5-preview",
+        input_price=Decimal("75.00"),
+        cached_input_price=Decimal("37.50"),
+        output_price=Decimal("150.00"),
+    ),
+    "gpt-4.5-preview-2025-02-27": OpenAIModelConfig(
+        display_name="GPT-4.5 Preview (2025-02-27)",
+        base_model="gpt-4.5-preview",
+        input_price=Decimal("75.00"),
+        cached_input_price=Decimal("37.50"),
+        output_price=Decimal("150.00"),
+    ),
+    "gpt-4o": OpenAIModelConfig(
+        display_name="GPT-4 Optimized",
+        base_model="gpt-4o",
+        input_price=Decimal("2.50"),
+        cached_input_price=Decimal("1.25"),
+        output_price=Decimal("10.00"),
+    ),
+    "gpt-4o-2024-08-06": OpenAIModelConfig(
+        display_name="GPT-4 Optimized (2024-08-06)",
+        base_model="gpt-4o",
+        input_price=Decimal("2.50"),
+        cached_input_price=Decimal("1.25"),
+        output_price=Decimal("10.00"),
+    ),
+    "gpt-4o-audio-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Audio Preview",
+        base_model="gpt-4o-audio-preview",
+        input_price=Decimal("2.50"),
+        cached_input_price=None,
+        output_price=Decimal("10.00"),
+    ),
+    "gpt-4o-audio-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Audio Preview (2024-12-17)",
+        base_model="gpt-4o-audio-preview",
+        input_price=Decimal("2.50"),
+        cached_input_price=None,
+        output_price=Decimal("10.00"),
+    ),
+    "gpt-4o-realtime-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Realtime Preview",
+        base_model="gpt-4o-realtime-preview",
+        input_price=Decimal("5.00"),
+        cached_input_price=Decimal("2.50"),
+        output_price=Decimal("20.00"),
+    ),
+    "gpt-4o-realtime-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Realtime Preview (2024-12-17)",
+        base_model="gpt-4o-realtime-preview",
+        input_price=Decimal("5.00"),
+        cached_input_price=Decimal("2.50"),
+        output_price=Decimal("20.00"),
+    ),
+    "gpt-4o-mini": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini",
+        base_model="gpt-4o-mini",
+        input_price=Decimal("0.15"),
+        cached_input_price=Decimal("0.075"),
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-mini-2024-07-18": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini (2024-07-18)",
+        base_model="gpt-4o-mini",
+        input_price=Decimal("0.15"),
+        cached_input_price=Decimal("0.075"),
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-mini-audio-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Audio Preview",
+        base_model="gpt-4o-mini-audio-preview",
+        input_price=Decimal("0.15"),
+        cached_input_price=None,
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-mini-audio-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Audio Preview (2024-12-17)",
+        base_model="gpt-4o-mini-audio-preview",
+        input_price=Decimal("0.15"),
+        cached_input_price=None,
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-mini-realtime-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Realtime Preview",
+        base_model="gpt-4o-mini-realtime-preview",
+        input_price=Decimal("0.60"),
+        cached_input_price=Decimal("0.30"),
+        output_price=Decimal("2.40"),
+    ),
+    "gpt-4o-mini-realtime-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Realtime Preview (2024-12-17)",
+        base_model="gpt-4o-mini-realtime-preview",
+        input_price=Decimal("0.60"),
+        cached_input_price=Decimal("0.30"),
+        output_price=Decimal("2.40"),
+    ),
+    "o1": OpenAIModelConfig(
+        display_name="O1",
+        base_model="o1",
+        input_price=Decimal("15.00"),
+        cached_input_price=Decimal("7.50"),
+        output_price=Decimal("60.00"),
+    ),
+    "o1-2024-12-17": OpenAIModelConfig(
+        display_name="O1 (2024-12-17)",
+        base_model="o1",
+        input_price=Decimal("15.00"),
+        cached_input_price=Decimal("7.50"),
+        output_price=Decimal("60.00"),
+    ),
+    "o3-mini": OpenAIModelConfig(
+        display_name="O3 Mini",
+        base_model="o3-mini",
+        input_price=Decimal("1.10"),
+        cached_input_price=Decimal("0.55"),
+        output_price=Decimal("4.40"),
+    ),
+    "o3-mini-2025-01-31": OpenAIModelConfig(
+        display_name="O3 Mini (2025-01-31)",
+        base_model="o3-mini",
+        input_price=Decimal("1.10"),
+        cached_input_price=Decimal("0.55"),
+        output_price=Decimal("4.40"),
+    ),
+    "o1-mini": OpenAIModelConfig(
+        display_name="O1 Mini",
+        base_model="o1-mini",
+        input_price=Decimal("1.10"),
+        cached_input_price=Decimal("0.55"),
+        output_price=Decimal("4.40"),
+    ),
+    "o1-mini-2024-09-12": OpenAIModelConfig(
+        display_name="O1 Mini (2024-09-12)",
+        base_model="o1-mini",
+        input_price=Decimal("1.10"),
+        cached_input_price=Decimal("0.55"),
+        output_price=Decimal("4.40"),
+    ),
+    "gpt-4o-mini-search-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Search Preview",
+        base_model="gpt-4o-mini-search-preview",
+        input_price=Decimal("0.15"),
+        cached_input_price=None,
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-mini-search-preview-2025-03-11": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Search Preview (2025-03-11)",
+        base_model="gpt-4o-mini-search-preview",
+        input_price=Decimal("0.15"),
+        cached_input_price=None,
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-search-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Search Preview",
+        base_model="gpt-4o-search-preview",
+        input_price=Decimal("2.50"),
+        cached_input_price=None,
+        output_price=Decimal("10.00"),
+    ),
+    "gpt-4o-search-preview-2025-03-11": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Search Preview (2025-03-11)",
+        base_model="gpt-4o-search-preview",
+        input_price=Decimal("2.50"),
+        cached_input_price=None,
+        output_price=Decimal("10.00"),
+    ),
+    "computer-use-preview": OpenAIModelConfig(
+        display_name="Computer Use Preview",
+        base_model="computer-use-preview",
+        input_price=Decimal("3.00"),
+        cached_input_price=None,
+        output_price=Decimal("12.00"),
+    ),
+    "computer-use-preview-2025-03-11": OpenAIModelConfig(
+        display_name="Computer Use Preview (2025-03-11)",
+        base_model="computer-use-preview",
+        input_price=Decimal("3.00"),
+        cached_input_price=None,
+        output_price=Decimal("12.00"),
+    ),
+}
+
+
+def get_model_config(model_id: str) -> OpenAIModelConfig:
+    """Get model configuration by model ID"""
+    if model_id not in OPENAI_MODELS:
+        raise ValueError(f"Model {model_id} not found in OpenAI models")
+    return OPENAI_MODELS[model_id]
+
+
+def get_default_model() -> str:
+    """Get the default model ID"""
+    return "gpt-4o"
+
+
+def calculate_cost(
+    model_id: str, input_tokens: int, output_tokens: int, use_cached: bool = False
+) -> Decimal:
+    """Calculate cost for token usage"""
+    config = get_model_config(model_id)
+    input_price = (
+        config.cached_input_price
+        if (use_cached and config.cached_input_price is not None)
+        else config.input_price
+    )
+    return (
+        input_price * Decimal(str(input_tokens))
+        + config.output_price * Decimal(str(output_tokens))
+    ) / Decimal("1000")