airtrain 0.1.12__tar.gz → 0.1.14__tar.gz

{airtrain-0.1.12/airtrain.egg-info → airtrain-0.1.14}/PKG-INFO

@@ -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 → airtrain-0.1.14}/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-0.1.12 → airtrain-0.1.14}/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-0.1.14/airtrain/contrib/travel/agentlib/verification_agent.py

@@ -0,0 +1,96 @@
+from typing import Optional, List, Tuple
+from datetime import date
+from pydantic import Field
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.integrations.openai.skills import OpenAIParserSkill, OpenAIParserInput
+from ..modellib.verification import UserTravelInfo
+from pydantic import BaseModel
+
+
+class VerificationInput(OpenAIParserInput):
+    conversation_history: List[str] = Field(
+        default_factory=list, description="History of conversation with user"
+    )
+    ask_followup: bool = Field(
+        default=True, description="Whether to ask follow-up questions"
+    )
+    followup_question: Optional[str] = Field(
+        default=None, description="Specific follow-up question to ask"
+    )
+
+
+class VerificationOutput(BaseModel):
+    travel_info: UserTravelInfo
+    needs_followup: bool
+    next_question: Optional[str] = None
+    missing_fields: List[str] = Field(default_factory=list)
+
+
+class UserVerificationAgent(OpenAIParserSkill):
+    """Agent for verifying and collecting user travel information"""
+
+    input_schema = VerificationInput
+    output_schema = VerificationOutput
+
+    def __init__(self, credentials=None):
+        super().__init__(credentials)
+        self.model = "gpt-4o"
+        self.temperature = 0.2
+
+    def process(self, input_data: VerificationInput) -> VerificationOutput:
+        system_prompt = """
+        You are a travel information verification agent. Your role is to:
+        1. Extract travel information from the conversation
+        2. Identify missing required information
+        3. Generate appropriate follow-up questions
+        4. Ensure all necessary details are collected for safe travel planning
+        
+        Required fields:
+        - Origin location
+        - Destination
+        - Travel dates
+        - Companions (if any)
+        - Preferred outdoor activities
+        - Health conditions (if any)
+        
+        Provide structured output and indicate if follow-up questions are needed.
+        """
+
+        # Combine conversation history into a single string
+        conversation = "\n".join(input_data.conversation_history)
+
+        # Add follow-up question if present
+        if input_data.followup_question:
+            conversation += f"\nFollow-up question: {input_data.followup_question}"
+
+        input_data = OpenAIParserInput(
+            user_input=conversation,
+            system_prompt=system_prompt,
+            response_model=VerificationOutput,
+            model=self.model,
+            temperature=self.temperature,
+        )
+
+        try:
+            result = self.process(input_data)
+            return result.parsed_response
+        except Exception as e:
+            raise ProcessingError(f"Failed to process verification: {str(e)}")
+
+    def get_next_question(self, missing_fields: List[str]) -> str:
+        """Generate appropriate follow-up question based on missing fields"""
+        questions = {
+            "origin": "What is your starting location?",
+            "destination": "Where would you like to travel to?",
+            "start_date": "When do you plan to start your journey?",
+            "end_date": "When do you plan to return?",
+            "companions": "Will anyone be traveling with you (children, pets, other adults)?",
+            "outdoor_activities": "What types of outdoor activities are you interested in?",
+            "health_conditions": "Do you or your companions have any health conditions we should be aware of?",
+        }
+
+        for field in missing_fields:
+            if field in questions:
+                return questions[field]
+
+        return "Is there anything else you'd like to share about your travel plans?"

airtrain-0.1.14/airtrain/contrib/travel/modellib/verification.py

@@ -0,0 +1,32 @@
+from typing import Optional, List, Dict
+from pydantic import BaseModel, Field
+from datetime import date
+
+
+class TravelCompanion(BaseModel):
+    type: str = Field(..., description="Type of companion (kid/pet/adult)")
+    count: int = Field(..., description="Number of companions of this type")
+    details: Optional[Dict[str, str]] = Field(
+        default=None, description="Additional details like ages, special needs"
+    )
+
+
+class HealthCondition(BaseModel):
+    condition: str = Field(..., description="Name of health condition")
+    severity: str = Field(..., description="Severity level (mild/moderate/severe)")
+    requirements: List[str] = Field(
+        ..., description="Special requirements or precautions"
+    )
+
+
+class UserTravelInfo(BaseModel):
+    origin: str = Field(..., description="Starting location")
+    destination: str = Field(..., description="Travel destination")
+    start_date: date = Field(..., description="Travel start date")
+    end_date: date = Field(..., description="Travel end date")
+    companions: List[TravelCompanion] = Field(default_factory=list)
+    outdoor_activities: List[str] = Field(default_factory=list)
+    health_conditions: List[HealthCondition] = Field(default_factory=list)
+    complete: bool = Field(
+        default=False, description="Whether all required info is collected"
+    )

airtrain-0.1.14/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-0.1.14/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-0.1.14/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-0.1.14/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-0.1.14/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-0.1.14/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-0.1.12 → airtrain-0.1.14}/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-0.1.14/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-0.1.14/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"