airtrain 0.1.12__py3-none-any.whl → 0.1.13__py3-none-any.whl

airtrain/__init__.py

@@ -1,6 +1,6 @@
 """Airtrain - A platform for building and deploying AI agents with structured skills"""
 
-__version__ = "0.1.12"
+__version__ = "0.1.13"
 
 # Core imports
 from .core.skills import Skill, ProcessingError

airtrain/contrib/travel/__init__.py

@@ -4,9 +4,9 @@ from .agents import (
     TravelAgentBase,
     ClothingAgent,
     HikingAgent,
-    InternetConnectivityAgent,
-    FoodRecommendationAgent,
-    PersonalizedRecommendationAgent,
+    InternetAgent,
+    FoodAgent,
+    PersonalizedAgent,
 )
 from .models import (
     ClothingRecommendation,
@@ -14,8 +14,8 @@ from .models import (
     InternetAvailability,
     FoodOption,
 )
-from .agents.verification_agent import UserVerificationAgent
-from .models.verification import UserTravelInfo, TravelCompanion, HealthCondition
+from .agentlib.verification_agent import UserVerificationAgent
+from .modellib.verification import UserTravelInfo, TravelCompanion, HealthCondition
 
 __all__ = [
     "TravelAgentBase",

airtrain/integrations/openai/models_config.py

@@ -0,0 +1,119 @@
+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-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-2024-05-13": OpenAIModelConfig(
+        display_name="GPT-4 Optimized (2024-05-13)",
+        base_model="gpt-4o",
+        input_price=Decimal("5.00"),
+        cached_input_price=None,
+        output_price=Decimal("15.00"),
+    ),
+    "gpt-4o-audio-preview-2024-12-17": 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-realtime-preview-2024-12-17": 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-mini-2024-07-18": 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-audio-preview-2024-12-17": 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-realtime-preview-2024-12-17": 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"),
+    ),
+    "o1-2024-12-17": OpenAIModelConfig(
+        display_name="O1",
+        base_model="o1",
+        input_price=Decimal("15.00"),
+        cached_input_price=Decimal("7.50"),
+        output_price=Decimal("60.00"),
+    ),
+    "o3-mini-2025-01-31": 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"),
+    ),
+    "o1-mini-2024-09-12": 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"),
+    ),
+}
+
+
+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/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

@@ -6,14 +6,14 @@ import together
 class TogetherAICredentials(BaseCredentials):
     """Together AI credentials"""
 
-    api_key: SecretStr = Field(..., description="Together AI API key")
+    together_api_key: SecretStr = Field(..., description="Together AI API key")
 
-    _required_credentials = {"api_key"}
+    _required_credentials = {"together_api_key"}
 
     async def validate_credentials(self) -> bool:
         """Validate Together AI credentials"""
         try:
-            together.api_key = self.api_key.get_secret_value()
+            together.api_key = self.together_api_key.get_secret_value()
             await together.Models.list()
             return True
         except Exception as 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,171 @@
+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: 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"
+    )
+    url: Optional[str] = Field(None, description="URL of the generated image")
+
+
+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
+
+            # Debug print the response structure
+            print(f"Response type: {type(response)}")
+            print(f"Response data type: {type(response.data)}")
+            if response.data:
+                print(f"First image type: {type(response.data[0])}")
+                print(f"First image attributes: {dir(response.data[0])}")
+
+            # Convert response to our output format
+            generated_images = []
+            for img in response.data:
+                # Try different possible attribute names for image data
+                b64_data = None
+                for attr in ["b64_json", "image", "base64", "data"]:
+                    if hasattr(img, attr):
+                        b64_data = getattr(img, attr)
+                        if b64_data:
+                            break
+
+                if not b64_data:
+                    # If no direct attribute found, try accessing as dictionary
+                    try:
+                        if hasattr(img, "__dict__"):
+                            img_dict = img.__dict__
+                            for key in ["b64_json", "image", "base64", "data"]:
+                                if key in img_dict and img_dict[key]:
+                                    b64_data = img_dict[key]
+                                    break
+                    except:
+                        pass
+
+                if not b64_data:
+                    raise ProcessingError(
+                        f"No image data found in API response. Response structure: {dir(img)}"
+                    )
+
+                generated_images.append(
+                    GeneratedImage(
+                        b64_json=b64_data,
+                        seed=getattr(img, "seed", None),
+                        finish_reason=getattr(img, "finish_reason", None),
+                        url=getattr(img, "url", 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/models.py

@@ -0,0 +1,56 @@
+from typing import List, Optional
+from pydantic import BaseModel, Field, validator
+from pathlib import Path
+
+
+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:
+            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"
+    )

airtrain/integrations/together/models_config.py

@@ -0,0 +1,277 @@
+from typing import Dict, NamedTuple
+
+
+class ModelConfig(NamedTuple):
+    organization: str
+    display_name: str
+    context_length: int
+    quantization: str
+
+
+TOGETHER_MODELS: Dict[str, ModelConfig] = {
+    # DeepSeek Models
+    "deepseek-ai/DeepSeek-R1": ModelConfig(
+        organization="DeepSeek",
+        display_name="DeepSeek-R1",
+        context_length=32768,
+        quantization="FP8",
+    ),
+    "deepseek-ai/DeepSeek-R1-Distill-Llama-70B": ModelConfig(
+        organization="DeepSeek",
+        display_name="DeepSeek R1 Distill Llama 70B",
+        context_length=131072,
+        quantization="FP16",
+    ),
+    "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B": ModelConfig(
+        organization="DeepSeek",
+        display_name="DeepSeek R1 Distill Qwen 1.5B",
+        context_length=131072,
+        quantization="FP16",
+    ),
+    "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": ModelConfig(
+        organization="DeepSeek",
+        display_name="DeepSeek R1 Distill Qwen 14B",
+        context_length=131072,
+        quantization="FP16",
+    ),
+    "deepseek-ai/DeepSeek-V3": ModelConfig(
+        organization="DeepSeek",
+        display_name="DeepSeek-V3",
+        context_length=16384,
+        quantization="FP8",
+    ),
+    # Meta Models
+    "meta-llama/Llama-3.3-70B-Instruct-Turbo": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3.3 70B Instruct Turbo",
+        context_length=131072,
+        quantization="FP8",
+    ),
+    "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3.1 8B Instruct Turbo",
+        context_length=131072,
+        quantization="FP8",
+    ),
+    "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3.1 70B Instruct Turbo",
+        context_length=131072,
+        quantization="FP8",
+    ),
+    "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3.1 405B Instruct Turbo",
+        context_length=130815,
+        quantization="FP8",
+    ),
+    "meta-llama/Meta-Llama-3-8B-Instruct-Turbo": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3 8B Instruct Turbo",
+        context_length=8192,
+        quantization="FP8",
+    ),
+    "meta-llama/Meta-Llama-3-70B-Instruct-Turbo": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3 70B Instruct Turbo",
+        context_length=8192,
+        quantization="FP8",
+    ),
+    "meta-llama/Llama-3.2-3B-Instruct-Turbo": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3.2 3B Instruct Turbo",
+        context_length=131072,
+        quantization="FP16",
+    ),
+    "meta-llama/Meta-Llama-3-8B-Instruct-Lite": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3 8B Instruct Lite",
+        context_length=8192,
+        quantization="INT4",
+    ),
+    "meta-llama/Meta-Llama-3-70B-Instruct-Lite": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3 70B Instruct Lite",
+        context_length=8192,
+        quantization="INT4",
+    ),
+    "meta-llama/Llama-3-8b-chat-hf": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3 8B Instruct Reference",
+        context_length=8192,
+        quantization="FP16",
+    ),
+    "meta-llama/Llama-3-70b-chat-hf": ModelConfig(
+        organization="Meta",
+        display_name="Llama 3 70B Instruct Reference",
+        context_length=8192,
+        quantization="FP16",
+    ),
+    "meta-llama/Llama-2-13b-chat-hf": ModelConfig(
+        organization="Meta",
+        display_name="LLaMA-2 Chat (13B)",
+        context_length=4096,
+        quantization="FP16",
+    ),
+    # Nvidia Models
+    "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF": ModelConfig(
+        organization="Nvidia",
+        display_name="Llama 3.1 Nemotron 70B",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    # Qwen Models
+    "Qwen/Qwen2.5-Coder-32B-Instruct": ModelConfig(
+        organization="Qwen",
+        display_name="Qwen 2.5 Coder 32B Instruct",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    "Qwen/QwQ-32B-Preview": ModelConfig(
+        organization="Qwen",
+        display_name="QwQ-32B-Preview",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    "Qwen/Qwen2.5-7B-Instruct-Turbo": ModelConfig(
+        organization="Qwen",
+        display_name="Qwen 2.5 7B Instruct Turbo",
+        context_length=32768,
+        quantization="FP8",
+    ),
+    "Qwen/Qwen2.5-72B-Instruct-Turbo": ModelConfig(
+        organization="Qwen",
+        display_name="Qwen 2.5 72B Instruct Turbo",
+        context_length=32768,
+        quantization="FP8",
+    ),
+    "Qwen/Qwen2-72B-Instruct": ModelConfig(
+        organization="Qwen",
+        display_name="Qwen 2 Instruct (72B)",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    "Qwen/Qwen2-VL-72B-Instruct": ModelConfig(
+        organization="Qwen",
+        display_name="Qwen2 VL 72B Instruct",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    # Microsoft Models
+    "microsoft/WizardLM-2-8x22B": ModelConfig(
+        organization="Microsoft",
+        display_name="WizardLM-2 8x22B",
+        context_length=65536,
+        quantization="FP16",
+    ),
+    # Google Models
+    "google/gemma-2-27b-it": ModelConfig(
+        organization="Google",
+        display_name="Gemma 2 27B",
+        context_length=8192,
+        quantization="FP16",
+    ),
+    "google/gemma-2-9b-it": ModelConfig(
+        organization="Google",
+        display_name="Gemma 2 9B",
+        context_length=8192,
+        quantization="FP16",
+    ),
+    "google/gemma-2b-it": ModelConfig(
+        organization="Google",
+        display_name="Gemma Instruct (2B)",
+        context_length=8192,
+        quantization="FP16",
+    ),
+    # Databricks Models
+    "databricks/dbrx-instruct": ModelConfig(
+        organization="databricks",
+        display_name="DBRX Instruct",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    # Gryphe Models
+    "Gryphe/MythoMax-L2-13b": ModelConfig(
+        organization="Gryphe",
+        display_name="MythoMax-L2 (13B)",
+        context_length=4096,
+        quantization="FP16",
+    ),
+    # Mistral AI Models
+    "mistralai/Mistral-Small-24B-Instruct-2501": ModelConfig(
+        organization="mistralai",
+        display_name="Mistral Small 3 Instruct (24B)",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    "mistralai/Mistral-7B-Instruct-v0.1": ModelConfig(
+        organization="mistralai",
+        display_name="Mistral (7B) Instruct",
+        context_length=8192,
+        quantization="FP16",
+    ),
+    "mistralai/Mistral-7B-Instruct-v0.2": ModelConfig(
+        organization="mistralai",
+        display_name="Mistral (7B) Instruct v0.2",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    "mistralai/Mistral-7B-Instruct-v0.3": ModelConfig(
+        organization="mistralai",
+        display_name="Mistral (7B) Instruct v0.3",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    "mistralai/Mixtral-8x7B-Instruct-v0.1": ModelConfig(
+        organization="mistralai",
+        display_name="Mixtral-8x7B Instruct (46.7B)",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    "mistralai/Mixtral-8x22B-Instruct-v0.1": ModelConfig(
+        organization="mistralai",
+        display_name="Mixtral-8x22B Instruct (141B)",
+        context_length=65536,
+        quantization="FP16",
+    ),
+    # NousResearch Models
+    "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO": ModelConfig(
+        organization="NousResearch",
+        display_name="Nous Hermes 2 - Mixtral 8x7B-DPO (46.7B)",
+        context_length=32768,
+        quantization="FP16",
+    ),
+    # Upstage Models
+    "upstage/SOLAR-10.7B-Instruct-v1.0": ModelConfig(
+        organization="upstage",
+        display_name="Upstage SOLAR Instruct v1 (11B)",
+        context_length=4096,
+        quantization="FP16",
+    ),
+}
+
+
+def get_model_config(model_id: str) -> ModelConfig:
+    """Get model configuration by model ID"""
+    if model_id not in TOGETHER_MODELS:
+        raise ValueError(f"Model {model_id} not found in Together AI models")
+    return TOGETHER_MODELS[model_id]
+
+
+def list_models_by_organization(organization: str) -> Dict[str, ModelConfig]:
+    """Get all models for a specific organization"""
+    return {
+        model_id: config
+        for model_id, config in TOGETHER_MODELS.items()
+        if config.organization.lower() == organization.lower()
+    }
+
+
+def get_default_model() -> str:
+    """Get the default model ID"""
+    return "meta-llama/Llama-3.3-70B-Instruct-Turbo"
+
+
+if __name__ == "__main__":
+    print(len(TOGETHER_MODELS))
+    print(get_model_config("meta-llama/Llama-3.3-70B-Instruct-Turbo"))

airtrain/integrations/together/rerank_models_config.py

@@ -0,0 +1,43 @@
+from typing import Dict, NamedTuple
+
+
+class RerankModelConfig(NamedTuple):
+    organization: str
+    display_name: str
+    model_size: str
+    max_doc_size: int
+    max_docs: int
+
+
+TOGETHER_RERANK_MODELS: Dict[str, RerankModelConfig] = {
+    "Salesforce/Llama-Rank-v1": RerankModelConfig(
+        organization="Salesforce",
+        display_name="LlamaRank",
+        model_size="8B",
+        max_doc_size=8192,
+        max_docs=1024,
+    )
+}
+
+
+def get_rerank_model_config(model_id: str) -> RerankModelConfig:
+    """Get rerank model configuration by model ID"""
+    if model_id not in TOGETHER_RERANK_MODELS:
+        raise ValueError(f"Model {model_id} not found in Together AI rerank models")
+    return TOGETHER_RERANK_MODELS[model_id]
+
+
+def list_rerank_models_by_organization(
+    organization: str,
+) -> Dict[str, RerankModelConfig]:
+    """Get all rerank models for a specific organization"""
+    return {
+        model_id: config
+        for model_id, config in TOGETHER_RERANK_MODELS.items()
+        if config.organization.lower() == organization.lower()
+    }
+
+
+def get_default_rerank_model() -> str:
+    """Get the default rerank model ID"""
+    return "Salesforce/Llama-Rank-v1"

airtrain/integrations/together/rerank_skill.py

@@ -0,0 +1,49 @@
+from typing import Optional
+from together import Together
+from airtrain.core.skills import Skill, ProcessingError
+from .credentials import TogetherAICredentials
+from .schemas import TogetherAIRerankInput, TogetherAIRerankOutput, RerankResult
+from .rerank_models_config import get_rerank_model_config
+
+
+class TogetherAIRerankSkill(Skill[TogetherAIRerankInput, TogetherAIRerankOutput]):
+    """Skill for reranking documents using Together AI"""
+
+    input_schema = TogetherAIRerankInput
+    output_schema = TogetherAIRerankOutput
+
+    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: TogetherAIRerankInput) -> TogetherAIRerankOutput:
+        try:
+            # Validate the model exists in our config
+            get_rerank_model_config(input_data.model)
+
+            # Call Together AI rerank API
+            response = self.client.rerank.create(
+                model=input_data.model,
+                query=input_data.query,
+                documents=input_data.documents,
+                top_n=input_data.top_n,
+            )
+
+            # Transform results
+            results = [
+                RerankResult(
+                    index=result.index,
+                    relevance_score=result.relevance_score,
+                    document=input_data.documents[result.index],
+                )
+                for result in response.results
+            ]
+
+            return TogetherAIRerankOutput(results=results, used_model=input_data.model)
+
+        except Exception as e:
+            raise ProcessingError(f"Together AI reranking failed: {str(e)}")

airtrain/integrations/together/schemas.py

@@ -0,0 +1,33 @@
+from typing import List, Optional
+from pydantic import Field, BaseModel
+from airtrain.core.schemas import InputSchema, OutputSchema
+
+
+class RerankResult(BaseModel):
+    """Schema for individual rerank result"""
+
+    index: int = Field(..., description="Index of the document in original list")
+    relevance_score: float = Field(..., description="Relevance score for the document")
+    document: str = Field(..., description="The document content")
+
+
+class TogetherAIRerankInput(InputSchema):
+    """Schema for Together AI rerank input"""
+
+    query: str = Field(..., description="Query to rank documents against")
+    documents: List[str] = Field(..., description="List of documents to rank")
+    model: str = Field(
+        default="Salesforce/Llama-Rank-v1",
+        description="Together AI rerank model to use",
+    )
+    top_n: Optional[int] = Field(
+        default=None,
+        description="Number of top results to return. If None, returns all results",
+    )
+
+
+class TogetherAIRerankOutput(OutputSchema):
+    """Schema for Together AI rerank output"""
+
+    results: List[RerankResult] = Field(..., description="Ranked results")
+    used_model: str = Field(..., description="Model used for ranking")

airtrain/integrations/together/skills.py

@@ -1,8 +1,13 @@
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, List
 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 TogetherAIImageInput, TogetherAIImageOutput, GeneratedImage
+from pathlib import Path
+import base64
+import time
+from together import Together
 
 
 class TogetherAIInput(InputSchema):
@@ -41,3 +46,84 @@ class TogetherAIChatSkill(Skill[TogetherAIInput, TogetherAIOutput]):
 
     def process(self, input_data: TogetherAIInput) -> TogetherAIOutput:
         raise NotImplementedError("TogetherAIChatSkill is not implemented yet")
+
+
+class TogetherAIImageSkill(Skill[TogetherAIImageInput, TogetherAIImageOutput]):
+    """Skill for Together AI image generation"""
+
+    input_schema = TogetherAIImageInput
+    output_schema = TogetherAIImageOutput
+
+    def __init__(self, credentials: Optional[TogetherAICredentials] = None):
+        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:
+            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 = [
+                GeneratedImage(
+                    b64_json=img.b64_json,
+                    seed=getattr(img, "seed", None),
+                    finish_reason=getattr(img, "finish_reason", None),
+                )
+                for img in response.data
+            ]
+
+            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/vision_models_config.py

@@ -0,0 +1,49 @@
+from typing import Dict, NamedTuple
+
+
+class VisionModelConfig(NamedTuple):
+    organization: str
+    display_name: str
+    context_length: int
+
+
+TOGETHER_VISION_MODELS: Dict[str, VisionModelConfig] = {
+    "meta-llama/Llama-Vision-Free": VisionModelConfig(
+        organization="Meta",
+        display_name="(Free) Llama 3.2 11B Vision Instruct Turbo",
+        context_length=131072,
+    ),
+    "meta-llama/Llama-3.2-11B-Vision-Instruct-Turbo": VisionModelConfig(
+        organization="Meta",
+        display_name="Llama 3.2 11B Vision Instruct Turbo",
+        context_length=131072,
+    ),
+    "meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo": VisionModelConfig(
+        organization="Meta",
+        display_name="Llama 3.2 90B Vision Instruct Turbo",
+        context_length=131072,
+    ),
+}
+
+
+def get_vision_model_config(model_id: str) -> VisionModelConfig:
+    """Get vision model configuration by model ID"""
+    if model_id not in TOGETHER_VISION_MODELS:
+        raise ValueError(f"Model {model_id} not found in Together AI vision models")
+    return TOGETHER_VISION_MODELS[model_id]
+
+
+def list_vision_models_by_organization(
+    organization: str,
+) -> Dict[str, VisionModelConfig]:
+    """Get all vision models for a specific organization"""
+    return {
+        model_id: config
+        for model_id, config in TOGETHER_VISION_MODELS.items()
+        if config.organization.lower() == organization.lower()
+    }
+
+
+def get_default_vision_model() -> str:
+    """Get the default vision model ID"""
+    return "meta-llama/Llama-Vision-Free"

{airtrain-0.1.12.dist-info → airtrain-0.1.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.12
+Version: 0.1.13
 Summary: A platform for building and deploying AI agents with structured skills
 Home-page: https://github.com/rosaboyle/airtrain.dev
 Author: Dheeraj Pai

{airtrain-0.1.12.dist-info → airtrain-0.1.13.dist-info}/RECORD

@@ -1,6 +1,6 @@
-airtrain/__init__.py,sha256=rk7ANjWmvBbBrDNBi_3_UsNXya3ifz8Qa9RLDFcrdrk,2095
+airtrain/__init__.py,sha256=BwgzBp8NKID9GrpW2RUmCyV6Aj_wlKI9HJqHTe6blMw,2095
 airtrain/contrib/__init__.py,sha256=pG-7mJ0pBMqp3Q86mIF9bo1PqoBOVSGlnEK1yY1U1ok,641
-airtrain/contrib/travel/__init__.py,sha256=rx_zYv8Qt5-DDiGT3H8IyGmEuBlNe3r71p_erYJxaLc,847
+airtrain/contrib/travel/__init__.py,sha256=clmBodw4nkTA-DsgjVGcXfJGPaWxIpCZDtdO-8RzL0M,811
 airtrain/contrib/travel/agents.py,sha256=tpQtZ0WUiXBuhvZtc2JlEam5TuR5l-Tndi14YyImDBM,8975
 airtrain/contrib/travel/models.py,sha256=E4Mtds5TINmXLXu65aTYqv6wwOh2KclJHZ2eXRrBqiY,1547
 airtrain/core/__init__.py,sha256=9h7iKwTzZocCPc9bU6j8bA02BokteWIOcO1uaqGMcrk,254
@@ -29,14 +29,25 @@ airtrain/integrations/ollama/skills.py,sha256=M_Un8D5VJ5XtPEq9IClzqV3jCPBoFTSm2v
 airtrain/integrations/openai/__init__.py,sha256=K-NY2_T1T6SEOgkpbUA55cWvK2nr2NOJgLCqmmtaCno,371
 airtrain/integrations/openai/chinese_assistant.py,sha256=MMhv4NBOoEQ0O22ZZtP255rd5ajHC9l6FPWIjpqxBOA,1581
 airtrain/integrations/openai/credentials.py,sha256=NfRyp1QgEtgm8cxt2-BOLq-6d0X-Pcm80NnfHM8p0FY,1470
+airtrain/integrations/openai/models_config.py,sha256=bzosqqpDy2AJxu2vGdk2H4voqEGlv7LORR6fpJLhNic,3962
 airtrain/integrations/openai/skills.py,sha256=Olg9-6f_p2XgkVwwcB9tvjAMApmM2EK81i8LP4qVVvs,7676
 airtrain/integrations/sambanova/__init__.py,sha256=dp_263iOckM_J9pOEvyqpf3FrejD6-_x33r0edMCTe0,179
 airtrain/integrations/sambanova/credentials.py,sha256=U36RAEIPNuwo-vTrt3U9kkkj2GfdqSclA1ttOYHxS-w,784
 airtrain/integrations/sambanova/skills.py,sha256=Po1ur_QFwzVIugbkk2mt73WdXDz_Gr9ASlUc9Y12Kok,1614
 airtrain/integrations/together/__init__.py,sha256=we4KXn_pUs6Dxo3QcB-t40BSRraQFdKg2nXw7yi2FjM,185
-airtrain/integrations/together/credentials.py,sha256=y5M6ZQrfYJLJbClxEasq4HaVyZM0l5lFshwVP6jq2E4,720
-airtrain/integrations/together/skills.py,sha256=YMOULyk2TX32rCjhxK29e4ehn8iIzMXpg3xmdYtuyQQ,1664
-airtrain-0.1.12.dist-info/METADATA,sha256=NGt35x8Yel7vDKoe-AzD77E14JK5IG6fw7kR8tiNX4Y,4536
-airtrain-0.1.12.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-airtrain-0.1.12.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
-airtrain-0.1.12.dist-info/RECORD,,
+airtrain/integrations/together/audio_models_config.py,sha256=GtqfmKR1vJ5x4B3kScvEO3x4exvzwNP78vcGVTk_fBE,1004
+airtrain/integrations/together/credentials.py,sha256=cYNhyIwgsxm8LfiFfT-omBvgV3mUP6SZeRSukyzzDlI,747
+airtrain/integrations/together/embedding_models_config.py,sha256=F0ISAXCG_Pcnf-ojkvZwIXacXD8LaU8hQmGHCFzmlds,2927
+airtrain/integrations/together/image_models_config.py,sha256=JlCozrphI9zE4uYpGfj4DCWSN6GZGyr84Tb1HmjNQ28,2455
+airtrain/integrations/together/image_skill.py,sha256=FMO9-TRwfucLRlpvij9VpeVsrTiP9wux_je_pFz6OXs,6508
+airtrain/integrations/together/models.py,sha256=ZW5xfEN9fU18aeltb-sB2O-Bnu5sLkDPZqvUtxgoH-U,2112
+airtrain/integrations/together/models_config.py,sha256=XMKp0Oq1nWWnMMdNAZxkFXmJaURwWrwLE18kFXsMsRw,8829
+airtrain/integrations/together/rerank_models_config.py,sha256=coCg0IOG2tU4L2uc2uPtPdoBwGjSc_zQwxENwdDuwHE,1188
+airtrain/integrations/together/rerank_skill.py,sha256=gjH24hLWCweWKPyyfKZMG3K_g9gWzm80WgiJNjkA9eg,1894
+airtrain/integrations/together/schemas.py,sha256=pBMrbX67oxPCr-sg4K8_Xqu1DWbaC4uLCloVSascROg,1210
+airtrain/integrations/together/skills.py,sha256=UfLHnseZbA7R7q5dDco6mpV546Zfd3DTliZSrNkCL6Q,4518
+airtrain/integrations/together/vision_models_config.py,sha256=m28HwYDk2Kup_J-a1FtynIa2ZVcbl37kltfoHnK8zxs,1544
+airtrain-0.1.13.dist-info/METADATA,sha256=kvTff3-jQ_th6KyfKrSQ8EOjjq5SiyHwq802VFYKqCg,4536
+airtrain-0.1.13.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+airtrain-0.1.13.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.13.dist-info/RECORD,,