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

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.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.12
+Version: 0.1.14
 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.14.dist-info}/RECORD

@@ -1,6 +1,6 @@
-airtrain/__init__.py,sha256=rk7ANjWmvBbBrDNBi_3_UsNXya3ifz8Qa9RLDFcrdrk,2095
+airtrain/__init__.py,sha256=FRI-wOLxVAu3ECtPJGQ8ZLjDSM1PgoiDGvn8ctNtS_8,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
@@ -17,6 +17,10 @@ airtrain/integrations/aws/skills.py,sha256=TQiMXeXRRcJ14fe8Xi7Uk20iS6_INbcznuLGt
 airtrain/integrations/cerebras/__init__.py,sha256=zAD-qV38OzHhMCz1z-NvjjqcYEhURbm8RWTOKHNqbew,174
 airtrain/integrations/cerebras/credentials.py,sha256=IFkn8LxMAaOpvEWXDpb94VQGtqcDxQ7rZHKH-tX4Nuw,884
 airtrain/integrations/cerebras/skills.py,sha256=O9vwFzvv_tUOwFOVE8CszAQEac711eVYVUj_8dVMTpc,1596
+airtrain/integrations/fireworks/__init__.py,sha256=9pJvP0u1FJbNtB0oHa09mHVJLctELf_c27LOYyDk2ZI,271
+airtrain/integrations/fireworks/credentials.py,sha256=UpcwR9V5Hbk5sJbjFDJDbHMRqc90IQSqAvrtJCOvwEo,524
+airtrain/integrations/fireworks/models.py,sha256=F-MddbLCLAsTjwRr1l6IpJxOegyY4pD7jN9ySPiypSo,593
+airtrain/integrations/fireworks/skills.py,sha256=ZykowW8lMbTcZVJ0GO2Ut6E-u2-keXvE4F-_j-3JI4k,4074
 airtrain/integrations/google/__init__.py,sha256=INZFNOcNebz3m-Ggk07ZjmX0kNHIbTe_St9gBlZBki8,176
 airtrain/integrations/google/credentials.py,sha256=Mm4jNWF02rIf0_GuHLcUUPyLHC4NMRdF_iTCoVTQ0Bs,1033
 airtrain/integrations/google/skills.py,sha256=uwmgetl5Ien7fLOA5HIZdqoL6AZnexFDyzfsrGuJ1RU,1606
@@ -29,14 +33,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.14.dist-info/METADATA,sha256=l4IPKLJ7Bf3gmZYSRPVEfz4oe1XGt_lWfvuZg68cNnE,4536
+airtrain-0.1.14.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+airtrain-0.1.14.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.14.dist-info/RECORD,,