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

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")

airtrain/integrations/openai/skills.py

@@ -0,0 +1,342 @@
+from typing import AsyncGenerator, List, Optional, Dict, TypeVar, Type, Generator, Union
+from pydantic import Field, BaseModel
+from openai import OpenAI, AsyncOpenAI
+from openai.types.chat import ChatCompletionChunk
+import numpy as np
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import OpenAICredentials
+
+
+class OpenAIInput(InputSchema):
+    """Schema for OpenAI chat input"""
+
+    user_input: str = Field(..., description="User's input text")
+    system_prompt: str = Field(
+        default="You are a helpful assistant.",
+        description="System prompt to guide the model's behavior",
+    )
+    conversation_history: List[Dict[str, str]] = Field(
+        default_factory=list,
+        description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
+    )
+    model: str = Field(
+        default="gpt-4o",
+        description="OpenAI model to use",
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+    max_tokens: Optional[int] = Field(
+        default=131072, description="Maximum tokens in response"
+    )
+    stream: bool = Field(
+        default=False,
+        description="Whether to stream the response token by token",
+    )
+
+
+class OpenAIOutput(OutputSchema):
+    """Schema for OpenAI chat output"""
+
+    response: str
+    used_model: str
+    usage: Dict[str, int]
+
+
+class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
+    """Skill for interacting with OpenAI models with async support"""
+
+    input_schema = OpenAIInput
+    output_schema = OpenAIOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+        self.async_client = AsyncOpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def _build_messages(self, input_data: OpenAIInput) -> List[Dict[str, str]]:
+        """Build messages list from input data including conversation history."""
+        messages = [{"role": "system", "content": input_data.system_prompt}]
+
+        if input_data.conversation_history:
+            messages.extend(input_data.conversation_history)
+
+        messages.append({"role": "user", "content": input_data.user_input})
+        return messages
+
+    def process_stream(self, input_data: OpenAIInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
+        try:
+            messages = self._build_messages(input_data)
+
+            stream = self.client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+                stream=True,
+            )
+
+            for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
+
+        except Exception as e:
+            raise ProcessingError(f"OpenAI streaming failed: {str(e)}")
+
+    def process(self, input_data: OpenAIInput) -> OpenAIOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                # For streaming, collect the entire response
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+            else:
+                # For non-streaming, use regular completion
+                messages = self._build_messages(input_data)
+                completion = self.client.chat.completions.create(
+                    model=input_data.model,
+                    messages=messages,
+                    temperature=input_data.temperature,
+                    max_tokens=input_data.max_tokens,
+                    stream=False,
+                )
+                response = completion.choices[0].message.content
+
+            return OpenAIOutput(
+                response=response,
+                used_model=input_data.model,
+                usage={
+                    "total_tokens": completion.usage.total_tokens,
+                    "prompt_tokens": completion.usage.prompt_tokens,
+                    "completion_tokens": completion.usage.completion_tokens,
+                },
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"OpenAI chat failed: {str(e)}")
+
+    async def process_async(self, input_data: OpenAIInput) -> OpenAIOutput:
+        """Async version of process method"""
+        try:
+            messages = self._build_messages(input_data)
+            completion = await self.async_client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+            return OpenAIOutput(
+                response=completion.choices[0].message.content,
+                used_model=completion.model,
+                usage={
+                    "total_tokens": completion.usage.total_tokens,
+                    "prompt_tokens": completion.usage.prompt_tokens,
+                    "completion_tokens": completion.usage.completion_tokens,
+                },
+            )
+        except Exception as e:
+            raise ProcessingError(f"OpenAI async chat failed: {str(e)}")
+
+    async def process_stream_async(
+        self, input_data: OpenAIInput
+    ) -> AsyncGenerator[str, None]:
+        """Async version of stream processor"""
+        try:
+            messages = self._build_messages(input_data)
+            stream = await self.async_client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+                stream=True,
+            )
+            async for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
+        except Exception as e:
+            raise ProcessingError(f"OpenAI async streaming failed: {str(e)}")
+
+
+ResponseT = TypeVar("ResponseT", bound=BaseModel)
+
+
+class OpenAIParserInput(InputSchema):
+    """Schema for OpenAI structured output input"""
+
+    user_input: str
+    system_prompt: str = "You are a helpful assistant that provides structured data."
+    model: str = "gpt-4o"
+    temperature: float = 0.7
+    max_tokens: Optional[int] = None
+    response_model: Type[ResponseT]
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class OpenAIParserOutput(OutputSchema):
+    """Schema for OpenAI structured output"""
+
+    parsed_response: BaseModel
+    used_model: str
+    tokens_used: int
+
+
+class OpenAIParserSkill(Skill[OpenAIParserInput, OpenAIParserOutput]):
+    """Skill for getting structured responses from OpenAI"""
+
+    input_schema = OpenAIParserInput
+    output_schema = OpenAIParserOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def process(self, input_data: OpenAIParserInput) -> OpenAIParserOutput:
+        try:
+            # Use parse method instead of create
+            completion = self.client.beta.chat.completions.parse(
+                model=input_data.model,
+                messages=[
+                    {"role": "system", "content": input_data.system_prompt},
+                    {"role": "user", "content": input_data.user_input},
+                ],
+                response_format=input_data.response_model,
+            )
+
+            if completion.choices[0].message.parsed is None:
+                raise ProcessingError("Failed to parse response")
+
+            return OpenAIParserOutput(
+                parsed_response=completion.choices[0].message.parsed,
+                used_model=completion.model,
+                tokens_used=completion.usage.total_tokens,
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"OpenAI parsing failed: {str(e)}")
+
+
+class OpenAIEmbeddingsInput(InputSchema):
+    """Schema for OpenAI embeddings input"""
+
+    texts: Union[str, List[str]] = Field(
+        ..., description="Text or list of texts to generate embeddings for"
+    )
+    model: str = Field(
+        default="text-embedding-3-large", description="OpenAI embeddings model to use"
+    )
+    encoding_format: str = Field(
+        default="float", description="The format of the embeddings: 'float' or 'base64'"
+    )
+    dimensions: Optional[int] = Field(
+        default=None, description="Optional number of dimensions for the embeddings"
+    )
+
+
+class OpenAIEmbeddingsOutput(OutputSchema):
+    """Schema for OpenAI embeddings output"""
+
+    embeddings: List[List[float]] = Field(..., description="List of embeddings vectors")
+    used_model: str = Field(..., description="Model used for generating embeddings")
+    tokens_used: int = Field(..., description="Number of tokens used")
+
+
+class OpenAIEmbeddingsSkill(Skill[OpenAIEmbeddingsInput, OpenAIEmbeddingsOutput]):
+    """Skill for generating embeddings using OpenAI models"""
+
+    input_schema = OpenAIEmbeddingsInput
+    output_schema = OpenAIEmbeddingsOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+        self.async_client = AsyncOpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def process(self, input_data: OpenAIEmbeddingsInput) -> OpenAIEmbeddingsOutput:
+        """Generate embeddings for the input text(s)"""
+        try:
+            # Handle single text input
+            texts = (
+                [input_data.texts]
+                if isinstance(input_data.texts, str)
+                else input_data.texts
+            )
+
+            # Create embeddings
+            response = self.client.embeddings.create(
+                model=input_data.model,
+                input=texts,
+                encoding_format=input_data.encoding_format,
+                dimensions=input_data.dimensions,
+            )
+
+            # Extract embeddings
+            embeddings = [data.embedding for data in response.data]
+
+            return OpenAIEmbeddingsOutput(
+                embeddings=embeddings,
+                used_model=response.model,
+                tokens_used=response.usage.total_tokens,
+            )
+        except Exception as e:
+            raise ProcessingError(f"OpenAI embeddings generation failed: {str(e)}")
+
+    async def process_async(
+        self, input_data: OpenAIEmbeddingsInput
+    ) -> OpenAIEmbeddingsOutput:
+        """Async version of the embeddings generation"""
+        try:
+            # Handle single text input
+            texts = (
+                [input_data.texts]
+                if isinstance(input_data.texts, str)
+                else input_data.texts
+            )
+
+            # Create embeddings
+            response = await self.async_client.embeddings.create(
+                model=input_data.model,
+                input=texts,
+                encoding_format=input_data.encoding_format,
+                dimensions=input_data.dimensions,
+            )
+
+            # Extract embeddings
+            embeddings = [data.embedding for data in response.data]
+
+            return OpenAIEmbeddingsOutput(
+                embeddings=embeddings,
+                used_model=response.model,
+                tokens_used=response.usage.total_tokens,
+            )
+        except Exception as e:
+            raise ProcessingError(
+                f"OpenAI async embeddings generation failed: {str(e)}"
+            )

airtrain/integrations/perplexity/__init__.py

@@ -0,0 +1,49 @@
+"""Perplexity AI integration module"""
+
+from .credentials import PerplexityCredentials
+from .skills import (
+    PerplexityInput,
+    PerplexityOutput,
+    PerplexityChatSkill,
+    PerplexityCitation,
+    PerplexityStreamingChatSkill,
+    PerplexityStreamOutput,
+)
+from .list_models import (
+    PerplexityListModelsSkill,
+    StandalonePerplexityListModelsSkill,
+    PerplexityListModelsInput,
+    PerplexityListModelsOutput,
+)
+from .models_config import (
+    get_model_config,
+    get_default_model,
+    supports_citations,
+    supports_search,
+    get_models_by_category,
+    PERPLEXITY_MODELS_CONFIG,
+)
+
+__all__ = [
+    # Credentials
+    "PerplexityCredentials",
+    # Skills
+    "PerplexityInput",
+    "PerplexityOutput",
+    "PerplexityChatSkill",
+    "PerplexityCitation",
+    "PerplexityStreamingChatSkill",
+    "PerplexityStreamOutput",
+    # List Models
+    "PerplexityListModelsSkill",
+    "StandalonePerplexityListModelsSkill",
+    "PerplexityListModelsInput",
+    "PerplexityListModelsOutput",
+    # Model Config
+    "get_model_config",
+    "get_default_model",
+    "supports_citations",
+    "supports_search",
+    "get_models_by_category",
+    "PERPLEXITY_MODELS_CONFIG",
+]

airtrain/integrations/perplexity/credentials.py

@@ -0,0 +1,43 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+import requests
+
+
+class PerplexityCredentials(BaseCredentials):
+    """Perplexity AI API credentials"""
+
+    perplexity_api_key: SecretStr = Field(..., description="Perplexity AI API key")
+
+    _required_credentials = {"perplexity_api_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Perplexity AI credentials by making a test API call"""
+        try:
+            headers = {
+                "Authorization": f"Bearer {self.perplexity_api_key.get_secret_value()}",
+                "Content-Type": "application/json",
+            }
+
+            # Small API call to check if credentials are valid
+            data = {
+                "model": "sonar-pro",
+                "messages": [{"role": "user", "content": "Test"}],
+                "max_tokens": 1,
+            }
+
+            # Make a synchronous request for validation
+            response = requests.post(
+                "https://api.perplexity.ai/chat/completions", headers=headers, json=data
+            )
+
+            if response.status_code == 200:
+                return True
+            else:
+                raise CredentialValidationError(
+                    f"Invalid Perplexity AI credentials: {response.status_code} - {response.text}"
+                )
+
+        except Exception as e:
+            raise CredentialValidationError(
+                f"Invalid Perplexity AI credentials: {str(e)}"
+            )