airtrain 0.1.12__py3-none-any.whl → 0.1.14__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.14"
 
 # 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/fireworks/__init__.py

@@ -0,0 +1,11 @@
+"""Fireworks AI integration module"""
+
+from .credentials import FireworksCredentials
+from .skills import FireworksChatSkill, FireworksInput, FireworksOutput
+
+__all__ = [
+    "FireworksCredentials",
+    "FireworksChatSkill",
+    "FireworksInput",
+    "FireworksOutput",
+]

airtrain/integrations/fireworks/credentials.py

@@ -0,0 +1,18 @@
+from pydantic import SecretStr, BaseModel
+from typing import Optional
+import os
+
+
+class FireworksCredentials(BaseModel):
+    """Credentials for Fireworks AI API"""
+
+    fireworks_api_key: SecretStr
+
+    @classmethod
+    def from_env(cls) -> "FireworksCredentials":
+        """Create credentials from environment variables"""
+        api_key = os.getenv("FIREWORKS_API_KEY")
+        if not api_key:
+            raise ValueError("FIREWORKS_API_KEY environment variable not set")
+
+        return cls(fireworks_api_key=api_key)

airtrain/integrations/fireworks/models.py

@@ -0,0 +1,27 @@
+from typing import List, Optional, Dict, Any
+from pydantic import Field, BaseModel
+
+
+class FireworksMessage(BaseModel):
+    """Schema for Fireworks chat message"""
+
+    content: str
+    role: str = Field(..., pattern="^(system|user|assistant)$")
+
+
+class FireworksUsage(BaseModel):
+    """Schema for Fireworks API usage statistics"""
+
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class FireworksResponse(BaseModel):
+    """Schema for Fireworks API response"""
+
+    id: str
+    choices: List[Dict[str, Any]]
+    created: int
+    model: str
+    usage: FireworksUsage

airtrain/integrations/fireworks/skills.py

@@ -0,0 +1,107 @@
+from typing import List, Optional, Dict, Any
+from pydantic import Field
+import requests
+from loguru import logger
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+from .models import FireworksMessage, FireworksResponse
+
+
+class FireworksInput(InputSchema):
+    """Schema for Fireworks AI chat input"""
+
+    user_input: str = Field(..., description="User's input text")
+    system_prompt: str = Field(
+        default="You are a helpful assistant.",
+        description="System prompt to guide the model's behavior",
+    )
+    model: str = Field(
+        default="accounts/fireworks/models/deepseek-r1",
+        description="Fireworks AI model to use",
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+    max_tokens: Optional[int] = Field(
+        default=None, description="Maximum tokens in response"
+    )
+    context_length_exceeded_behavior: str = Field(
+        default="truncate", description="Behavior when context length is exceeded"
+    )
+
+
+class FireworksOutput(OutputSchema):
+    """Schema for Fireworks AI output"""
+
+    response: str = Field(..., description="Model's response text")
+    used_model: str = Field(..., description="Model used for generation")
+    usage: Dict[str, int] = Field(default_factory=dict, description="Usage statistics")
+
+
+class FireworksChatSkill(Skill[FireworksInput, FireworksOutput]):
+    """Skill for interacting with Fireworks AI models"""
+
+    input_schema = FireworksInput
+    output_schema = FireworksOutput
+
+    def __init__(self, credentials: Optional[FireworksCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or FireworksCredentials.from_env()
+        self.base_url = "https://api.fireworks.ai/inference/v1"
+
+    def process(self, input_data: FireworksInput) -> FireworksOutput:
+        """Process the input using Fireworks AI API"""
+        try:
+            logger.info(f"Processing request with model {input_data.model}")
+
+            # Prepare messages
+            messages = [
+                {"role": "system", "content": input_data.system_prompt},
+                {"role": "user", "content": input_data.user_input},
+            ]
+
+            # Prepare request payload
+            payload = {
+                "messages": messages,
+                "model": input_data.model,
+                "context_length_exceeded_behavior": input_data.context_length_exceeded_behavior,
+                "temperature": input_data.temperature,
+                "n": 1,
+                "response_format": {"type": "text"},
+                "stream": False,
+            }
+
+            if input_data.max_tokens:
+                payload["max_tokens"] = input_data.max_tokens
+
+            # Make API request
+            response = requests.post(
+                f"{self.base_url}/chat/completions",
+                json=payload,
+                headers={
+                    "Authorization": f"Bearer {self.credentials.fireworks_api_key.get_secret_value()}",
+                    "Content-Type": "application/json",
+                },
+            )
+
+            response.raise_for_status()
+            response_data = FireworksResponse(**response.json())
+
+            logger.success("Successfully processed Fireworks AI request")
+
+            return FireworksOutput(
+                response=response_data.choices[0]["message"]["content"],
+                used_model=response_data.model,
+                usage={
+                    "prompt_tokens": response_data.usage.prompt_tokens,
+                    "completion_tokens": response_data.usage.completion_tokens,
+                    "total_tokens": response_data.usage.total_tokens,
+                },
+            )
+
+        except Exception as e:
+            logger.exception(f"Fireworks AI processing failed: {str(e)}")
+            raise ProcessingError(f"Fireworks AI processing failed: {str(e)}")

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"