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

airtrain/integrations/search/exa/skills.py

@@ -0,0 +1,115 @@
+"""
+Skills for Exa Search API.
+
+This module provides skills for using the Exa search API.
+"""
+
+import json
+import logging
+import httpx
+from typing import Optional, Dict, Any, List, cast
+
+from pydantic import ValidationError
+
+from airtrain.core.skills import Skill
+from airtrain.core.errors import ProcessingError
+from .credentials import ExaCredentials
+from .schemas import ExaSearchInputSchema, ExaSearchOutputSchema, ExaSearchResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class ExaSearchSkill(Skill[ExaSearchInputSchema, ExaSearchOutputSchema]):
+    """Skill for searching the web using the Exa search API."""
+
+    input_schema = ExaSearchInputSchema
+    output_schema = ExaSearchOutputSchema
+
+    EXA_API_ENDPOINT = "https://api.exa.ai/search"
+
+    def __init__(
+        self,
+        credentials: ExaCredentials,
+        timeout: float = 60.0,
+        max_retries: int = 3,
+        **kwargs,
+    ):
+        """
+        Initialize the Exa search skill.
+
+        Args:
+            credentials: Credentials for accessing the Exa API
+            timeout: Timeout for API requests in seconds
+            max_retries: Maximum number of retries for failed requests
+        """
+        super().__init__(**kwargs)
+        self.credentials = credentials
+        self.timeout = timeout
+        self.max_retries = max_retries
+
+    async def process(self, input_data: ExaSearchInputSchema) -> ExaSearchOutputSchema:
+        """
+        Process a search request using the Exa API.
+
+        Args:
+            input_data: Search input parameters
+
+        Returns:
+            Search results from Exa
+
+        Raises:
+            ProcessingError: If there's an issue with the API request
+        """
+        try:
+            # Prepare request payload
+            payload = input_data.model_dump(exclude_none=True)
+
+            # Build request headers
+            headers = {
+                "content-type": "application/json",
+                "Authorization": f"Bearer {self.credentials.api_key.get_secret_value()}",
+            }
+
+            # Make the API request
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    self.EXA_API_ENDPOINT,
+                    headers=headers,
+                    json=payload,
+                    timeout=self.timeout,
+                )
+
+                # Check for successful response
+                if response.status_code == 200:
+                    result_data = response.json()
+
+                    # Construct the output schema
+                    output = ExaSearchOutputSchema(
+                        results=result_data.get("results", []),
+                        query=input_data.query,
+                        autopromptString=result_data.get("autopromptString"),
+                        costDollars=result_data.get("costDollars"),
+                    )
+
+                    return output
+                else:
+                    # Handle error responses
+                    error_message = f"Exa API returned status code {response.status_code}: {response.text}"
+                    logger.error(error_message)
+                    raise ProcessingError(error_message)
+
+        except httpx.TimeoutException:
+            error_message = f"Timeout while querying Exa API (timeout={self.timeout}s)"
+            logger.error(error_message)
+            raise ProcessingError(error_message)
+
+        except ValidationError as e:
+            error_message = f"Schema validation error: {str(e)}"
+            logger.error(error_message)
+            raise ProcessingError(error_message)
+
+        except Exception as e:
+            error_message = f"Unexpected error while querying Exa API: {str(e)}"
+            logger.error(error_message)
+            raise ProcessingError(error_message)

airtrain/integrations/together/__init__.py

@@ -0,0 +1,33 @@
+"""Together AI integration module"""
+
+from .credentials import TogetherAICredentials
+from .skills import TogetherAIChatSkill, TogetherAIInput, TogetherAIOutput
+from .models_config import (
+    get_model_config_with_capabilities,
+    get_max_completion_tokens,
+    supports_tool_use,
+    supports_json_mode,
+    TOGETHER_MODELS_CONFIG,
+)
+from .list_models import (
+    TogetherListModelsSkill,
+    TogetherListModelsInput,
+    TogetherListModelsOutput,
+)
+from .models import TogetherModel
+
+__all__ = [
+    "TogetherAICredentials", 
+    "TogetherAIChatSkill",
+    "TogetherAIInput",
+    "TogetherAIOutput",
+    "TogetherListModelsSkill",
+    "TogetherListModelsInput",
+    "TogetherListModelsOutput",
+    "TogetherModel",
+    "get_model_config_with_capabilities",
+    "get_max_completion_tokens",
+    "supports_tool_use",
+    "supports_json_mode",
+    "TOGETHER_MODELS_CONFIG",
+]

airtrain/integrations/together/audio_models_config.py

@@ -0,0 +1,34 @@
+from typing import Dict, NamedTuple
+
+
+class AudioModelConfig(NamedTuple):
+    organization: str
+    display_name: str
+
+
+TOGETHER_AUDIO_MODELS: Dict[str, AudioModelConfig] = {
+    "Cartesia/Sonic": AudioModelConfig(
+        organization="Cartesia", display_name="Cartesia/Sonic"
+    )
+}
+
+
+def get_audio_model_config(model_id: str) -> AudioModelConfig:
+    """Get audio model configuration by model ID"""
+    if model_id not in TOGETHER_AUDIO_MODELS:
+        raise ValueError(f"Model {model_id} not found in Together AI audio models")
+    return TOGETHER_AUDIO_MODELS[model_id]
+
+
+def list_audio_models_by_organization(organization: str) -> Dict[str, AudioModelConfig]:
+    """Get all audio models for a specific organization"""
+    return {
+        model_id: config
+        for model_id, config in TOGETHER_AUDIO_MODELS.items()
+        if config.organization.lower() == organization.lower()
+    }
+
+
+def get_default_audio_model() -> str:
+    """Get the default audio model ID"""
+    return "Cartesia/Sonic"

airtrain/integrations/together/credentials.py

@@ -0,0 +1,22 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+import together
+
+
+class TogetherAICredentials(BaseCredentials):
+    """Together AI credentials"""
+
+    together_api_key: SecretStr = Field(..., description="Together AI API key")
+
+    _required_credentials = {"together_api_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Together AI credentials"""
+        try:
+            together.api_key = self.together_api_key.get_secret_value()
+            await together.Models.list()
+            return True
+        except Exception as e:
+            raise CredentialValidationError(
+                f"Invalid Together AI credentials: {str(e)}"
+            )

airtrain/integrations/together/embedding_models_config.py

@@ -0,0 +1,92 @@
+from typing import Dict, NamedTuple
+
+
+class EmbeddingModelConfig(NamedTuple):
+    organization: str
+    display_name: str
+    model_size: str
+    embedding_dimension: int
+    context_window: int
+
+
+TOGETHER_EMBEDDING_MODELS: Dict[str, EmbeddingModelConfig] = {
+    "togethercomputer/m2-bert-80M-2k-retrieval": EmbeddingModelConfig(
+        organization="Together",
+        display_name="M2-BERT-80M-2K-Retrieval",
+        model_size="80M",
+        embedding_dimension=768,
+        context_window=2048,
+    ),
+    "togethercomputer/m2-bert-80M-8k-retrieval": EmbeddingModelConfig(
+        organization="Together",
+        display_name="M2-BERT-80M-8K-Retrieval",
+        model_size="80M",
+        embedding_dimension=768,
+        context_window=8192,
+    ),
+    "togethercomputer/m2-bert-80M-32k-retrieval": EmbeddingModelConfig(
+        organization="Together",
+        display_name="M2-BERT-80M-32K-Retrieval",
+        model_size="80M",
+        embedding_dimension=768,
+        context_window=32768,
+    ),
+    "WhereIsAI/UAE-Large-V1": EmbeddingModelConfig(
+        organization="WhereIsAI",
+        display_name="UAE-Large-v1",
+        model_size="326M",
+        embedding_dimension=1024,
+        context_window=512,
+    ),
+    "BAAI/bge-large-en-v1.5": EmbeddingModelConfig(
+        organization="BAAI",
+        display_name="BGE-Large-EN-v1.5",
+        model_size="326M",
+        embedding_dimension=1024,
+        context_window=512,
+    ),
+    "BAAI/bge-base-en-v1.5": EmbeddingModelConfig(
+        organization="BAAI",
+        display_name="BGE-Base-EN-v1.5",
+        model_size="102M",
+        embedding_dimension=768,
+        context_window=512,
+    ),
+    "sentence-transformers/msmarco-bert-base-dot-v5": EmbeddingModelConfig(
+        organization="sentence-transformers",
+        display_name="Sentence-BERT",
+        model_size="110M",
+        embedding_dimension=768,
+        context_window=512,
+    ),
+    "bert-base-uncased": EmbeddingModelConfig(
+        organization="Hugging Face",
+        display_name="BERT",
+        model_size="110M",
+        embedding_dimension=768,
+        context_window=512,
+    ),
+}
+
+
+def get_embedding_model_config(model_id: str) -> EmbeddingModelConfig:
+    """Get embedding model configuration by model ID"""
+    if model_id not in TOGETHER_EMBEDDING_MODELS:
+        raise ValueError(f"Model {model_id} not found in Together AI embedding models")
+    return TOGETHER_EMBEDDING_MODELS[model_id]
+
+
+def list_embedding_models_by_organization(
+    organization: str,
+) -> Dict[str, EmbeddingModelConfig]:
+    """Get all embedding models for a specific organization"""
+    return {
+        model_id: config
+        for model_id, config in TOGETHER_EMBEDDING_MODELS.items()
+        if config.organization.lower() == organization.lower()
+    }
+
+
+def get_default_embedding_model() -> str:
+    """Get the default embedding model ID"""
+    return "togethercomputer/m2-bert-80M-32k-retrieval"

airtrain/integrations/together/image_models_config.py

@@ -0,0 +1,69 @@
+from typing import Dict, NamedTuple, Optional
+
+
+class ImageModelConfig(NamedTuple):
+    organization: str
+    display_name: str
+    default_steps: Optional[int]
+
+
+TOGETHER_IMAGE_MODELS: Dict[str, ImageModelConfig] = {
+    "black-forest-labs/FLUX.1-schnell-Free": ImageModelConfig(
+        organization="Black Forest Labs",
+        display_name="Flux.1 [schnell] (free)",
+        default_steps=None,
+    ),
+    "black-forest-labs/FLUX.1-schnell": ImageModelConfig(
+        organization="Black Forest Labs",
+        display_name="Flux.1 [schnell] (Turbo)",
+        default_steps=4,
+    ),
+    "black-forest-labs/FLUX.1-dev": ImageModelConfig(
+        organization="Black Forest Labs", display_name="Flux.1 Dev", default_steps=28
+    ),
+    "black-forest-labs/FLUX.1-canny": ImageModelConfig(
+        organization="Black Forest Labs", display_name="Flux.1 Canny", default_steps=28
+    ),
+    "black-forest-labs/FLUX.1-depth": ImageModelConfig(
+        organization="Black Forest Labs", display_name="Flux.1 Depth", default_steps=28
+    ),
+    "black-forest-labs/FLUX.1-redux": ImageModelConfig(
+        organization="Black Forest Labs", display_name="Flux.1 Redux", default_steps=28
+    ),
+    "black-forest-labs/FLUX.1.1-pro": ImageModelConfig(
+        organization="Black Forest Labs",
+        display_name="Flux1.1 [pro]",
+        default_steps=None,
+    ),
+    "black-forest-labs/FLUX.1-pro": ImageModelConfig(
+        organization="Black Forest Labs",
+        display_name="Flux.1 [pro]",
+        default_steps=None,
+    ),
+    "stabilityai/stable-diffusion-xl-base-1.0": ImageModelConfig(
+        organization="Stability AI",
+        display_name="Stable Diffusion XL 1.0",
+        default_steps=None,
+    ),
+}
+
+
+def get_image_model_config(model_id: str) -> ImageModelConfig:
+    """Get image model configuration by model ID"""
+    if model_id not in TOGETHER_IMAGE_MODELS:
+        raise ValueError(f"Model {model_id} not found in Together AI image models")
+    return TOGETHER_IMAGE_MODELS[model_id]
+
+
+def list_image_models_by_organization(organization: str) -> Dict[str, ImageModelConfig]:
+    """Get all image models for a specific organization"""
+    return {
+        model_id: config
+        for model_id, config in TOGETHER_IMAGE_MODELS.items()
+        if config.organization.lower() == organization.lower()
+    }
+
+
+def get_default_image_model() -> str:
+    """Get the default image model ID"""
+    return "black-forest-labs/FLUX.1-schnell-Free"

airtrain/integrations/together/image_skill.py

@@ -0,0 +1,143 @@
+from typing import Optional, List
+from pathlib import Path
+from pydantic import Field
+from together import Together
+import base64
+import time
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import TogetherAICredentials
+from .image_models_config import get_image_model_config, get_default_image_model
+
+
+class TogetherAIImageInput(InputSchema):
+    """Schema for Together AI image generation input"""
+
+    prompt: str = Field(..., description="Text prompt for image generation")
+    model: str = Field(
+        default=get_default_image_model(), description="Together AI image model to use"
+    )
+    steps: int = Field(default=10, description="Number of inference steps", ge=1, le=50)
+    n: int = Field(default=1, description="Number of images to generate", ge=1, le=4)
+    size: str = Field(
+        default="1024x1024", description="Image size in format WIDTHxHEIGHT"
+    )
+    negative_prompt: Optional[str] = Field(
+        default=None, description="Things to exclude from the generation"
+    )
+    seed: Optional[int] = Field(
+        default=None, description="Random seed for reproducibility"
+    )
+
+
+class GeneratedImage(OutputSchema):
+    """Individual generated image data"""
+
+    b64_json: Optional[str] = Field(None, description="Base64 encoded image data")
+    url: str = Field(..., description="URL of the generated image")
+    seed: Optional[int] = Field(None, description="Seed used for this image")
+    finish_reason: Optional[str] = Field(
+        None, description="Reason for finishing generation"
+    )
+
+
+class TogetherAIImageOutput(OutputSchema):
+    """Schema for Together AI image generation output"""
+
+    images: List[GeneratedImage] = Field(..., description="List of generated images")
+    model: str = Field(..., description="Model used for generation")
+    prompt: str = Field(..., description="Original prompt used")
+    total_time: float = Field(..., description="Time taken for generation in seconds")
+    usage: dict = Field(default_factory=dict, description="Usage statistics")
+
+
+class TogetherAIImageSkill(Skill[TogetherAIImageInput, TogetherAIImageOutput]):
+    """Skill for generating images using Together AI"""
+
+    input_schema = TogetherAIImageInput
+    output_schema = TogetherAIImageOutput
+
+    def __init__(self, credentials: Optional[TogetherAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or TogetherAICredentials.from_env()
+        self.client = Together(
+            api_key=self.credentials.together_api_key.get_secret_value()
+        )
+
+    def process(self, input_data: TogetherAIImageInput) -> TogetherAIImageOutput:
+        try:
+            # Validate the model exists in our config
+            get_image_model_config(input_data.model)
+
+            start_time = time.time()
+
+            # Generate images
+            response = self.client.images.generate(
+                prompt=input_data.prompt,
+                model=input_data.model,
+                steps=input_data.steps,
+                n=input_data.n,
+                size=input_data.size,
+                negative_prompt=input_data.negative_prompt,
+                seed=input_data.seed,
+            )
+
+            # Calculate total time
+            total_time = time.time() - start_time
+
+            # Convert response to our output format
+            generated_images = []
+            for img in response.data:
+                if not hasattr(img, "url"):
+                    raise ProcessingError(
+                        f"No URL found in API response. Response structure: {dir(img)}"
+                    )
+
+                generated_images.append(
+                    GeneratedImage(
+                        url=img.url,
+                        seed=getattr(img, "seed", None),
+                        finish_reason=getattr(img, "finish_reason", None),
+                    )
+                )
+
+            return TogetherAIImageOutput(
+                images=generated_images,
+                model=input_data.model,
+                prompt=input_data.prompt,
+                total_time=total_time,
+                usage=getattr(response, "usage", {}),
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Together AI image generation failed: {str(e)}")
+
+    def save_images(
+        self, output: TogetherAIImageOutput, output_dir: Path
+    ) -> List[Path]:
+        """
+        Save generated images to disk
+
+        Args:
+            output (TogetherAIImageOutput): Generation output containing images
+            output_dir (Path): Directory to save images
+
+        Returns:
+            List[Path]: List of paths to saved images
+        """
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        saved_paths = []
+        for i, img in enumerate(output.images):
+            output_path = output_dir / f"image_{i}.png"
+            image_data = base64.b64decode(img.b64_json)
+
+            with open(output_path, "wb") as f:
+                f.write(image_data)
+
+            saved_paths.append(output_path)
+
+        return saved_paths

airtrain/integrations/together/list_models.py

@@ -0,0 +1,76 @@
+from typing import Optional
+import requests
+from pydantic import Field
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import TogetherAICredentials
+from .models import TogetherModel
+
+
+class TogetherListModelsInput(InputSchema):
+    """Schema for Together AI list models input"""
+    pass
+
+
+class TogetherListModelsOutput(OutputSchema):
+    """Schema for Together AI list models output"""
+    
+    data: list[TogetherModel] = Field(
+        default_factory=list,
+        description="List of Together AI models"
+    )
+    object: Optional[str] = Field(
+        default=None, 
+        description="Object type"
+    )
+
+
+class TogetherListModelsSkill(Skill[TogetherListModelsInput, TogetherListModelsOutput]):
+    """Skill for listing Together AI models"""
+
+    input_schema = TogetherListModelsInput
+    output_schema = TogetherListModelsOutput
+    
+    def __init__(self, credentials: Optional[TogetherAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or TogetherAICredentials.from_env()
+        self.base_url = "https://api.together.xyz/v1"
+    
+    def process(
+        self, input_data: TogetherListModelsInput
+    ) -> TogetherListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            # Build the URL
+            url = f"{self.base_url}/models"
+            
+            # Make the request
+            headers = {
+                "Authorization": (
+                    f"Bearer {self.credentials.together_api_key.get_secret_value()}"
+                ),
+                "accept": "application/json"
+            }
+            
+            response = requests.get(url, headers=headers)
+            response.raise_for_status()
+            
+            # Parse the response
+            result = response.json()
+            
+            # Convert the models to TogetherModel objects
+            models = []
+            for model_data in result:
+                models.append(TogetherModel(**model_data))
+            
+            # Return the output
+            return TogetherListModelsOutput(
+                data=models,
+            )
+            
+        except requests.RequestException as e:
+            raise ProcessingError(f"Failed to list Together AI models: {str(e)}")
+        except Exception as e:
+            raise ProcessingError(f"Error listing Together AI models: {str(e)}") 

airtrain/integrations/together/models.py

@@ -0,0 +1,95 @@
+from typing import List, Optional, Dict, Any
+from pydantic import BaseModel, Field, validator
+
+
+class TogetherAIImageInput(BaseModel):
+    """Schema for Together AI image generation input"""
+
+    prompt: str = Field(..., description="Text prompt for image generation")
+    model: str = Field(
+        default="black-forest-labs/FLUX.1-schnell-Free",
+        description="Together AI image model to use",
+    )
+    steps: int = Field(default=10, description="Number of inference steps", ge=1, le=50)
+    n: int = Field(default=1, description="Number of images to generate", ge=1, le=4)
+    size: str = Field(
+        default="1024x1024", description="Image size in format WIDTHxHEIGHT"
+    )
+    negative_prompt: Optional[str] = Field(
+        default=None, description="Things to exclude from the generation"
+    )
+    seed: Optional[int] = Field(
+        default=None, description="Random seed for reproducibility"
+    )
+
+    @validator("size")
+    def validate_size(cls, v):
+        try:
+            width, height = map(int, v.split("x"))
+            if width <= 0 or height <= 0:
+                raise ValueError
+            return v
+        except ValueError:
+            raise ValueError("Size must be in format WIDTHxHEIGHT (e.g., 1024x1024)")
+
+
+class GeneratedImage(BaseModel):
+    """Individual generated image data"""
+
+    b64_json: str = Field(..., description="Base64 encoded image data")
+    seed: Optional[int] = Field(None, description="Seed used for this image")
+    finish_reason: Optional[str] = Field(
+        None, description="Reason for finishing generation"
+    )
+
+
+class TogetherAIImageOutput(BaseModel):
+    """Schema for Together AI image generation output"""
+
+    images: List[GeneratedImage] = Field(..., description="List of generated images")
+    model: str = Field(..., description="Model used for generation")
+    prompt: str = Field(..., description="Original prompt used")
+    total_time: float = Field(..., description="Time taken for generation in seconds")
+    usage: dict = Field(
+        default_factory=dict, description="Usage statistics and billing information"
+    )
+
+
+class TogetherModel(BaseModel):
+    """Schema for Together AI model"""
+
+    id: str = Field(..., description="Model ID")
+    name: Optional[str] = Field(None, description="Model name")
+    object: Optional[str] = Field(None, description="Object type")
+    created: Optional[int] = Field(None, description="Creation timestamp")
+    owned_by: Optional[str] = Field(None, description="Model owner")
+    root: Optional[str] = Field(None, description="Root model identifier")
+    parent: Optional[str] = Field(None, description="Parent model identifier")
+    permission: Optional[List[Dict[str, Any]]] = Field(
+        None, description="Permission details"
+    )
+    metadata: Optional[Dict[str, Any]] = Field(
+        None, description="Additional metadata for the model"
+    )
+    description: Optional[str] = Field(None, description="Model description")
+    pricing: Optional[Dict[str, Any]] = Field(None, description="Pricing information")
+    context_length: Optional[int] = Field(
+        None, description="Maximum context length supported by the model"
+    )
+    capabilities: Optional[List[str]] = Field(
+        None, description="Model capabilities"
+    )
+
+
+class TogetherListModelsInput(BaseModel):
+    """Schema for listing Together AI models input"""
+    pass
+
+
+class TogetherListModelsOutput(BaseModel):
+    """Schema for listing Together AI models output"""
+    
+    data: List[TogetherModel] = Field(
+        ..., description="List of Together AI models"
+    )
+    object: Optional[str] = Field(None, description="Object type")