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

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