airtrain 0.1.43__py3-none-any.whl → 0.1.45__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.43"
+__version__ = "0.1.45"
 
 # Core imports
 from .core.skills import Skill, ProcessingError

airtrain/integrations/__init__.py

@@ -26,6 +26,13 @@ from .cerebras.skills import CerebrasChatSkill
 from .openai.models_config import OPENAI_MODELS, OpenAIModelConfig
 from .anthropic.models_config import ANTHROPIC_MODELS, AnthropicModelConfig
 
+# Combined modules
+from .combined.list_models_factory import (
+    ListModelsSkillFactory,
+    GenericListModelsInput,
+    GenericListModelsOutput
+)
+
 __all__ = [
     # Credentials
     "OpenAICredentials",
@@ -53,4 +60,8 @@ __all__ = [
     "OpenAIModelConfig",
     "ANTHROPIC_MODELS",
     "AnthropicModelConfig",
+    # Combined modules
+    "ListModelsSkillFactory",
+    "GenericListModelsInput",
+    "GenericListModelsOutput",
 ]

airtrain/integrations/anthropic/__init__.py

@@ -9,6 +9,12 @@ from .models_config import (
     get_default_model,
     calculate_cost,
 )
+from .list_models import (
+    AnthropicListModelsSkill,
+    AnthropicListModelsInput,
+    AnthropicListModelsOutput,
+    AnthropicModel,
+)
 
 __all__ = [
     "AnthropicCredentials",
@@ -20,4 +26,8 @@ __all__ = [
     "get_model_config",
     "get_default_model",
     "calculate_cost",
+    "AnthropicListModelsSkill",
+    "AnthropicListModelsInput",
+    "AnthropicListModelsOutput",
+    "AnthropicModel",
 ]

airtrain/integrations/anthropic/list_models.py

@@ -0,0 +1,110 @@
+from typing import Optional, List, Dict, Any
+from pydantic import Field
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import AnthropicCredentials
+from .models_config import ANTHROPIC_MODELS, AnthropicModelConfig
+
+
+class AnthropicModel:
+    """Class to represent an Anthropic model."""
+    
+    def __init__(self, model_id: str, config: AnthropicModelConfig):
+        """Initialize the Anthropic model."""
+        self.id = model_id
+        self.display_name = config.display_name
+        self.base_model = config.base_model
+        self.input_price = config.input_price
+        self.cached_write_price = config.cached_write_price
+        self.cached_read_price = config.cached_read_price
+        self.output_price = config.output_price
+    
+    def dict(self, exclude_none=False):
+        """Convert the model to a dictionary."""
+        result = {
+            "id": self.id,
+            "display_name": self.display_name,
+            "base_model": self.base_model,
+            "input_price": float(self.input_price),
+            "output_price": float(self.output_price),
+        }
+        
+        if self.cached_write_price is not None:
+            result["cached_write_price"] = float(self.cached_write_price)
+        elif not exclude_none:
+            result["cached_write_price"] = None
+            
+        if self.cached_read_price is not None:
+            result["cached_read_price"] = float(self.cached_read_price)
+        elif not exclude_none:
+            result["cached_read_price"] = None
+            
+        return result
+
+
+class AnthropicListModelsInput(InputSchema):
+    """Schema for Anthropic list models input"""
+    
+    api_models_only: bool = Field(
+        default=False,
+        description=(
+            "If True, fetch models from the API only. If False, use local config."
+        )
+    )
+
+
+class AnthropicListModelsOutput(OutputSchema):
+    """Schema for Anthropic list models output"""
+    
+    models: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="List of Anthropic models"
+    )
+
+
+class AnthropicListModelsSkill(
+    Skill[AnthropicListModelsInput, AnthropicListModelsOutput]
+):
+    """Skill for listing Anthropic models"""
+
+    input_schema = AnthropicListModelsInput
+    output_schema = AnthropicListModelsOutput
+    
+    def __init__(self, credentials: Optional[AnthropicCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials
+    
+    def process(
+        self, input_data: AnthropicListModelsInput
+    ) -> AnthropicListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            models = []
+            
+            if input_data.api_models_only:
+                # Fetch models from Anthropic API
+                # Require credentials if using API models
+                if not self.credentials:
+                    raise ProcessingError(
+                        "Anthropic credentials required for API models"
+                    )
+                
+                # Note: Anthropic doesn't have a public models list endpoint
+                # We'll raise an error instead
+                raise ProcessingError(
+                    "Anthropic API does not provide a models list endpoint. "
+                    "Use api_models_only=False to list models from local config."
+                )
+            else:
+                # Use local model config - no credentials needed
+                for model_id, config in ANTHROPIC_MODELS.items():
+                    model = AnthropicModel(model_id, config)
+                    models.append(model.dict())
+            
+            # Return the output
+            return AnthropicListModelsOutput(models=models)
+            
+        except Exception as e:
+            raise ProcessingError(f"Failed to list Anthropic models: {str(e)}") 

airtrain/integrations/anthropic/skills.py

@@ -26,7 +26,7 @@ class AnthropicInput(InputSchema):
         default="claude-3-opus-20240229", description="Anthropic model to use"
     )
     max_tokens: Optional[int] = Field(
-        default=1024, description="Maximum tokens in response"
+        default=131072, description="Maximum tokens in response"
     )
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1

airtrain/integrations/aws/skills.py

@@ -21,7 +21,7 @@ class AWSBedrockInput(InputSchema):
         default="anthropic.claude-3-sonnet-20240229-v1:0",
         description="AWS Bedrock model to use",
     )
-    max_tokens: int = Field(default=1024, description="Maximum tokens in response")
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )

airtrain/integrations/cerebras/skills.py

@@ -22,7 +22,7 @@ class CerebrasInput(InputSchema):
     )
     model: str = Field(default="llama3.1-8b", description="Cerebras model to use")
     max_tokens: Optional[int] = Field(
-        default=1024, description="Maximum tokens in response"
+        default=131072, description="Maximum tokens in response"
     )
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1

airtrain/integrations/combined/__init__.py

@@ -0,0 +1,21 @@
+"""Combined integration modules for Airtrain"""
+
+from .groq_fireworks_skills import (
+    GroqFireworksSkill,
+    GroqFireworksInput,
+    GroqFireworksOutput
+)
+from .list_models_factory import (
+    ListModelsSkillFactory, 
+    GenericListModelsInput, 
+    GenericListModelsOutput
+)
+
+__all__ = [
+    "GroqFireworksSkill",
+    "GroqFireworksInput",
+    "GroqFireworksOutput",
+    "ListModelsSkillFactory",
+    "GenericListModelsInput",
+    "GenericListModelsOutput"
+] 

airtrain/integrations/combined/groq_fireworks_skills.py

@@ -0,0 +1,126 @@
+from typing import Optional, Dict, Any, List
+from pydantic import Field
+import requests
+from groq import Groq
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from airtrain.integrations.fireworks.completion_skills import (
+    FireworksCompletionSkill,
+    FireworksCompletionInput,
+)
+
+
+class GroqFireworksInput(InputSchema):
+    """Schema for combined Groq and Fireworks input"""
+
+    user_input: str = Field(..., description="User's input text")
+    groq_model: str = Field(
+        default="mixtral-8x7b-32768", description="Groq model to use"
+    )
+    fireworks_model: str = Field(
+        default="accounts/fireworks/models/deepseek-r1",
+        description="Fireworks model to use",
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation"
+    )
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
+
+
+class GroqFireworksOutput(OutputSchema):
+    """Schema for combined Groq and Fireworks output"""
+
+    combined_response: str
+    groq_response: str
+    fireworks_response: str
+    used_models: Dict[str, str]
+    usage: Dict[str, Dict[str, int]]
+
+
+class GroqFireworksSkill(Skill[GroqFireworksInput, GroqFireworksOutput]):
+    """Skill combining Groq and Fireworks responses"""
+
+    input_schema = GroqFireworksInput
+    output_schema = GroqFireworksOutput
+
+    def __init__(
+        self,
+        groq_api_key: Optional[str] = None,
+        fireworks_skill: Optional[FireworksCompletionSkill] = None,
+    ):
+        """Initialize the skill with optional API keys"""
+        super().__init__()
+        self.groq_client = Groq(api_key=groq_api_key)
+        self.fireworks_skill = fireworks_skill or FireworksCompletionSkill()
+
+    def _get_groq_response(self, input_data: GroqFireworksInput) -> Dict[str, Any]:
+        """Get response from Groq"""
+        try:
+            completion = self.groq_client.chat.completions.create(
+                model=input_data.groq_model,
+                messages=[{"role": "user", "content": input_data.user_input}],
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+            return {
+                "response": completion.choices[0].message.content,
+                "usage": completion.usage.model_dump(),
+            }
+        except Exception as e:
+            raise ProcessingError(f"Groq request failed: {str(e)}")
+
+    def _get_fireworks_response(
+        self, groq_response: str, input_data: GroqFireworksInput
+    ) -> Dict[str, Any]:
+        """Get response from Fireworks"""
+        try:
+            formatted_prompt = (
+                f"<USER>{input_data.user_input}</USER>\n<ASSISTANT>{groq_response}"
+            )
+
+            fireworks_input = FireworksCompletionInput(
+                prompt=formatted_prompt,
+                model=input_data.fireworks_model,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+
+            result = self.fireworks_skill.process(fireworks_input)
+            return {"response": result.response, "usage": result.usage}
+        except Exception as e:
+            raise ProcessingError(f"Fireworks request failed: {str(e)}")
+
+    def process(self, input_data: GroqFireworksInput) -> GroqFireworksOutput:
+        """Process the input using both Groq and Fireworks"""
+        try:
+            # Get Groq response
+            groq_result = self._get_groq_response(input_data)
+
+            # Get Fireworks response
+            fireworks_result = self._get_fireworks_response(
+                groq_result["response"], input_data
+            )
+
+            # Combine responses in the required format
+            combined_response = (
+                f"<USER>{input_data.user_input}</USER>\n"
+                f"<ASSISTANT>{groq_result['response']} {fireworks_result['response']}"
+            )
+
+            return GroqFireworksOutput(
+                combined_response=combined_response,
+                groq_response=groq_result["response"],
+                fireworks_response=fireworks_result["response"],
+                used_models={
+                    "groq": input_data.groq_model,
+                    "fireworks": input_data.fireworks_model,
+                },
+                usage={
+                    "groq": groq_result["usage"],
+                    "fireworks": fireworks_result["usage"],
+                },
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Combined processing failed: {str(e)}")

airtrain/integrations/combined/list_models_factory.py

@@ -0,0 +1,172 @@
+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 airtrain.core.credentials import BaseCredentials
+
+# Import existing list models skills
+from airtrain.integrations.openai.list_models import OpenAIListModelsSkill
+from airtrain.integrations.anthropic.list_models import AnthropicListModelsSkill
+from airtrain.integrations.together.list_models import TogetherListModelsSkill
+from airtrain.integrations.fireworks.list_models import FireworksListModelsSkill
+
+# Import credentials
+from airtrain.integrations.groq.credentials import GroqCredentials
+from airtrain.integrations.cerebras.credentials import CerebrasCredentials
+from airtrain.integrations.sambanova.credentials import SambanovaCredentials
+
+
+# Generic list models input schema
+class GenericListModelsInput(InputSchema):
+    """Generic schema for listing models from any provider"""
+    
+    api_models_only: bool = Field(
+        default=False,
+        description=(
+            "If True, fetch models from the API only. If False, use local config."
+        )
+    )
+
+
+# Generic list models output schema
+class GenericListModelsOutput(OutputSchema):
+    """Generic schema for list models output from any provider"""
+    
+    models: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="List of models"
+    )
+    provider: str = Field(
+        ...,
+        description="Provider name"
+    )
+
+
+# Base class for stub implementations
+class BaseListModelsSkill(Skill[GenericListModelsInput, GenericListModelsOutput]):
+    """Base skill for listing models"""
+    
+    input_schema = GenericListModelsInput
+    output_schema = GenericListModelsOutput
+    
+    def __init__(self, provider: str, credentials: Optional[BaseCredentials] = None):
+        """Initialize the skill with provider name and optional credentials"""
+        super().__init__()
+        self.provider = provider
+        self.credentials = credentials
+        
+    def get_models(self) -> List[Dict[str, Any]]:
+        """Return list of models. To be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement get_models()")
+    
+    def process(self, input_data: GenericListModelsInput) -> GenericListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            models = self.get_models()
+            return GenericListModelsOutput(models=models, provider=self.provider)
+        except Exception as e:
+            raise ProcessingError(f"Failed to list {self.provider} models: {str(e)}")
+
+
+# Groq implementation
+class GroqListModelsSkill(BaseListModelsSkill):
+    """Skill for listing Groq models"""
+    
+    def __init__(self, credentials: Optional[GroqCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__(provider="groq", credentials=credentials)
+    
+    def get_models(self) -> List[Dict[str, Any]]:
+        """Return list of Groq models."""
+        # Default Groq models from trmx_agent config
+        models = [
+            {"id": "llama-3-70b-8192", "display_name": "Llama 3 70B (8K)"},
+            {"id": "mixtral-8x7b-32768", "display_name": "Mixtral 8x7B (32K)"},
+            {"id": "gemma-7b-it", "display_name": "Gemma 7B Instruct"}
+        ]
+        return models
+
+
+# Cerebras implementation
+class CerebrasListModelsSkill(BaseListModelsSkill):
+    """Skill for listing Cerebras models"""
+    
+    def __init__(self, credentials: Optional[CerebrasCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__(provider="cerebras", credentials=credentials)
+    
+    def get_models(self) -> List[Dict[str, Any]]:
+        """Return list of Cerebras models."""
+        # Default Cerebras models from trmx_agent config
+        models = [
+            {"id": "cerebras/Cerebras-GPT-13B-v0.1", "display_name": "Cerebras GPT 13B v0.1"},
+            {"id": "cerebras/Cerebras-GPT-111M-v0.9", "display_name": "Cerebras GPT 111M v0.9"},
+            {"id": "cerebras/Cerebras-GPT-590M-v0.7", "display_name": "Cerebras GPT 590M v0.7"}
+        ]
+        return models
+
+
+# Sambanova implementation
+class SambanovaListModelsSkill(BaseListModelsSkill):
+    """Skill for listing Sambanova models"""
+    
+    def __init__(self, credentials: Optional[SambanovaCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__(provider="sambanova", credentials=credentials)
+    
+    def get_models(self) -> List[Dict[str, Any]]:
+        """Return list of Sambanova models."""
+        # Limited Sambanova model information
+        models = [
+            {"id": "sambanova/samba-1", "display_name": "Samba-1"},
+            {"id": "sambanova/samba-2", "display_name": "Samba-2"}
+        ]
+        return models
+
+
+# Factory class
+class ListModelsSkillFactory:
+    """Factory for creating list models skills for different providers"""
+    
+    # Map provider names to their corresponding list models skills
+    _PROVIDER_MAP = {
+        "openai": OpenAIListModelsSkill,
+        "anthropic": AnthropicListModelsSkill,
+        "together": TogetherListModelsSkill,
+        "fireworks": FireworksListModelsSkill,
+        "groq": GroqListModelsSkill,
+        "cerebras": CerebrasListModelsSkill,
+        "sambanova": SambanovaListModelsSkill
+    }
+    
+    @classmethod
+    def get_skill(cls, provider: str, credentials=None):
+        """Return a list models skill for the specified provider
+        
+        Args:
+            provider (str): The provider name (case-insensitive)
+            credentials: Optional credentials for the provider
+            
+        Returns:
+            A ListModelsSkill instance for the specified provider
+            
+        Raises:
+            ValueError: If the provider is not supported
+        """
+        provider = provider.lower()
+        
+        if provider not in cls._PROVIDER_MAP:
+            supported = ", ".join(cls.get_supported_providers())
+            raise ValueError(
+                f"Unsupported provider: {provider}. "
+                f"Supported providers are: {supported}"
+            )
+        
+        skill_class = cls._PROVIDER_MAP[provider]
+        return skill_class(credentials=credentials)
+    
+    @classmethod
+    def get_supported_providers(cls):
+        """Return a list of supported provider names"""
+        return list(cls._PROVIDER_MAP.keys()) 

airtrain/integrations/fireworks/completion_skills.py

@@ -17,7 +17,7 @@ class FireworksCompletionInput(InputSchema):
         default="accounts/fireworks/models/deepseek-r1",
         description="Fireworks AI model to use",
     )
-    max_tokens: int = Field(default=4096, description="Maximum tokens in response")
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )

airtrain/integrations/fireworks/conversation_manager.py

@@ -20,7 +20,7 @@ class ConversationState(BaseModel):
         description="Model being used for the conversation",
     )
     temperature: float = Field(default=0.7, description="Temperature setting")
-    max_tokens: Optional[int] = Field(default=None, description="Max tokens setting")
+    max_tokens: Optional[int] = Field(default=131072, description="Max tokens setting")
 
 
 class FireworksConversationManager:

airtrain/integrations/fireworks/requests_skills.py

@@ -29,7 +29,7 @@ class FireworksRequestInput(InputSchema):
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )
-    max_tokens: int = Field(default=4096, description="Maximum tokens in response")
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
     top_p: float = Field(
         default=1.0, description="Top p sampling parameter", ge=0, le=1
     )

airtrain/integrations/fireworks/skills.py

@@ -27,7 +27,7 @@ class FireworksInput(InputSchema):
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )
     max_tokens: Optional[int] = Field(
-        default=None, description="Maximum tokens in response"
+        default=131072, description="Maximum tokens in response"
     )
     context_length_exceeded_behavior: str = Field(
         default="truncate", description="Behavior when context length is exceeded"

airtrain/integrations/fireworks/structured_completion_skills.py

@@ -21,7 +21,7 @@ class FireworksStructuredCompletionInput(InputSchema):
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )
-    max_tokens: int = Field(default=4096, description="Maximum tokens in response")
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
     response_model: Type[ResponseT]
     stream: bool = Field(
         default=False,

airtrain/integrations/fireworks/structured_requests_skills.py

@@ -31,7 +31,7 @@ class FireworksStructuredRequestInput(InputSchema):
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )
-    max_tokens: int = Field(default=4096, description="Maximum tokens in response")
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
     response_model: Type[ResponseT]
     stream: bool = Field(
         default=False, description="Whether to stream the response token by token"

airtrain/integrations/fireworks/structured_skills.py

@@ -17,7 +17,7 @@ class FireworksParserInput(InputSchema):
     system_prompt: str = "You are a helpful assistant that provides structured data."
     model: str = "accounts/fireworks/models/deepseek-r1"
     temperature: float = 0.7
-    max_tokens: Optional[int] = None
+    max_tokens: Optional[int] = 131072
     response_model: Type[ResponseT]
     conversation_history: List[Dict[str, str]] = Field(
         default_factory=list,

airtrain/integrations/groq/skills.py

@@ -21,7 +21,7 @@ class GroqInput(InputSchema):
     model: str = Field(
         default="deepseek-r1-distill-llama-70b-specdec", description="Groq model to use"
     )
-    max_tokens: int = Field(default=1024, description="Maximum tokens in response")
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )

airtrain/integrations/ollama/skills.py

@@ -14,7 +14,7 @@ class OllamaInput(InputSchema):
         description="System prompt to guide the model's behavior",
     )
     model: str = Field(default="llama2", description="Ollama model to use")
-    max_tokens: int = Field(default=1024, description="Maximum tokens in response")
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )

airtrain/integrations/openai/__init__.py

@@ -1,3 +1,5 @@
+"""OpenAI API integration."""
+
 from .skills import (
     OpenAIChatSkill,
     OpenAIInput,
@@ -10,6 +12,12 @@ from .skills import (
     OpenAIEmbeddingsOutput,
 )
 from .credentials import OpenAICredentials
+from .list_models import (
+    OpenAIListModelsSkill,
+    OpenAIListModelsInput,
+    OpenAIListModelsOutput,
+    OpenAIModel,
+)
 
 __all__ = [
     "OpenAIChatSkill",
@@ -22,4 +30,8 @@ __all__ = [
     "OpenAIEmbeddingsSkill",
     "OpenAIEmbeddingsInput",
     "OpenAIEmbeddingsOutput",
+    "OpenAIListModelsSkill",
+    "OpenAIListModelsInput",
+    "OpenAIListModelsOutput",
+    "OpenAIModel",
 ]

airtrain/integrations/openai/chinese_assistant.py

@@ -17,7 +17,7 @@ class ChineseAssistantInput(OpenAIInput):
         description="System prompt in Chinese",
     )
     model: str = Field(default="gpt-4o", description="OpenAI model to use")
-    max_tokens: int = Field(default=8096, description="Maximum tokens in response")
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )

airtrain/integrations/openai/list_models.py

@@ -0,0 +1,112 @@
+from typing import Optional, List, Dict, Any
+from pydantic import Field
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import OpenAICredentials
+from .models_config import OPENAI_MODELS, OpenAIModelConfig
+
+
+class OpenAIModel:
+    """Class to represent an OpenAI model."""
+    
+    def __init__(self, model_id: str, config: OpenAIModelConfig):
+        """Initialize the OpenAI model."""
+        self.id = model_id
+        self.display_name = config.display_name
+        self.base_model = config.base_model
+        self.input_price = config.input_price
+        self.cached_input_price = config.cached_input_price
+        self.output_price = config.output_price
+    
+    def dict(self, exclude_none=False):
+        """Convert the model to a dictionary."""
+        result = {
+            "id": self.id,
+            "display_name": self.display_name,
+            "base_model": self.base_model,
+            "input_price": float(self.input_price),
+            "output_price": float(self.output_price),
+        }
+        if self.cached_input_price is not None:
+            result["cached_input_price"] = float(self.cached_input_price)
+        elif not exclude_none:
+            result["cached_input_price"] = None
+        return result
+
+
+class OpenAIListModelsInput(InputSchema):
+    """Schema for OpenAI list models input"""
+    
+    api_models_only: bool = Field(
+        default=False,
+        description=(
+            "If True, fetch models from the API only. If False, use local config."
+        )
+    )
+
+
+class OpenAIListModelsOutput(OutputSchema):
+    """Schema for OpenAI list models output"""
+    
+    models: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="List of OpenAI models"
+    )
+
+
+class OpenAIListModelsSkill(Skill[OpenAIListModelsInput, OpenAIListModelsOutput]):
+    """Skill for listing OpenAI models"""
+
+    input_schema = OpenAIListModelsInput
+    output_schema = OpenAIListModelsOutput
+    
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials
+    
+    def process(
+        self, input_data: OpenAIListModelsInput
+    ) -> OpenAIListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            models = []
+            
+            if input_data.api_models_only:
+                # Fetch models from OpenAI API - requires credentials
+                if not self.credentials:
+                    raise ProcessingError(
+                        "OpenAI credentials required for API models"
+                    )
+                
+                from openai import OpenAI
+                client = OpenAI(
+                    api_key=self.credentials.openai_api_key.get_secret_value(),
+                    organization=self.credentials.openai_organization_id,
+                )
+                
+                # Make API call to get models
+                response = client.models.list()
+                
+                # Convert response to our format
+                for model in response.data:
+                    models.append({
+                        "id": model.id,
+                        "display_name": model.id,  # API doesn't provide display_name
+                        "base_model": model.id,    # API doesn't provide base_model
+                        "created": model.created,
+                        "owned_by": model.owned_by,
+                        # Pricing info not available from API
+                    })
+            else:
+                # Use local model config - no credentials needed
+                for model_id, config in OPENAI_MODELS.items():
+                    model = OpenAIModel(model_id, config)
+                    models.append(model.dict())
+            
+            # Return the output
+            return OpenAIListModelsOutput(models=models)
+            
+        except Exception as e:
+            raise ProcessingError(f"Failed to list OpenAI models: {str(e)}") 

airtrain/integrations/openai/skills.py

@@ -29,7 +29,7 @@ class OpenAIInput(InputSchema):
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )
     max_tokens: Optional[int] = Field(
-        default=None, description="Maximum tokens in response"
+        default=131072, description="Maximum tokens in response"
     )
     stream: bool = Field(
         default=False,

airtrain/integrations/together/models_config.py

@@ -13,7 +13,7 @@ TOGETHER_MODELS: Dict[str, ModelConfig] = {
     "deepseek-ai/DeepSeek-R1": ModelConfig(
         organization="DeepSeek",
         display_name="DeepSeek-R1",
-        context_length=32768,
+        context_length=131072,
         quantization="FP8",
     ),
     "deepseek-ai/DeepSeek-R1-Distill-Llama-70B": ModelConfig(
@@ -37,7 +37,7 @@ TOGETHER_MODELS: Dict[str, ModelConfig] = {
     "deepseek-ai/DeepSeek-V3": ModelConfig(
         organization="DeepSeek",
         display_name="DeepSeek-V3",
-        context_length=16384,
+        context_length=131072,
         quantization="FP8",
     ),
     # Meta Models

{airtrain-0.1.43.dist-info → airtrain-0.1.45.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: airtrain
-Version: 0.1.43
+Version: 0.1.45
 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.43.dist-info → airtrain-0.1.45.dist-info}/RECORD

@@ -1,4 +1,4 @@
-airtrain/__init__.py,sha256=f32eUcBlLHtZ3NDoxvRR_a36tp6PkovEYV-G9IG__mw,2099
+airtrain/__init__.py,sha256=r8knBRelQBR1dv6YknyXNm00o3Eo7UmfPDuHYDxLtL4,2099
 airtrain/__main__.py,sha256=EU8ffFmCdC1G-UcHHt0Oo3lB1PGqfC6kwzH39CnYSwU,72
 airtrain/builder/__init__.py,sha256=D33sr0k_WAe6FAJkk8rUaivEzFaeVqLXkQgyFWEhfPU,110
 airtrain/builder/agent_builder.py,sha256=3XnGUAcK_6lWoUDtL0TanliQZuh7u0unhNbnrz1z2-I,5018
@@ -13,42 +13,47 @@ airtrain/core/__init__.py,sha256=9h7iKwTzZocCPc9bU6j8bA02BokteWIOcO1uaqGMcrk,254
 airtrain/core/credentials.py,sha256=PgQotrQc46J5djidKnkK1znUv3fyNkUFDO-m2Kn_Gzo,4006
 airtrain/core/schemas.py,sha256=MMXrDviC4gRea_QaPpbjgO--B_UKxnD7YrxqZOLJZZU,7003
 airtrain/core/skills.py,sha256=LljalzeSHK5eQPTAOEAYc5D8Qn1kVSfiz9WgziTD5UM,4688
-airtrain/integrations/__init__.py,sha256=rk9QFl0Dd7Qp4rULhi_u4smwsJwk69Kg_-fv0GQ43iw,1782
-airtrain/integrations/anthropic/__init__.py,sha256=F4kB5fuj7nYgTVcgzeHGc91LT96FZfsCJVBVCnTRh-k,541
+airtrain/integrations/__init__.py,sha256=Y0yJt6c6uKOkMXdGPELSc6NYiLrEYr_38acViNPXWNk,2046
+airtrain/integrations/anthropic/__init__.py,sha256=K741w3v7fWsCknTo38ARqDL0D3HPlwDIvDuuBao9Tto,800
 airtrain/integrations/anthropic/credentials.py,sha256=hlTSw9HX66kYNaeQUtn0JjdZQBMNkzzFOJOoLOOzvcY,1246
+airtrain/integrations/anthropic/list_models.py,sha256=o7FABp0Cq3gs76zOF-CM9ohmYslWT6vK9qtQabV9XzI,3973
 airtrain/integrations/anthropic/models_config.py,sha256=TZt31hLcT-9YK-NxqiarMyOwvUWMgXAzAcPfSwzDSiQ,3347
-airtrain/integrations/anthropic/skills.py,sha256=WV-9254H2VqUAq_7Zr1xG5IhejeC_gQSqyH0hwW1_tY,5870
+airtrain/integrations/anthropic/skills.py,sha256=bs9OYHeY1ETsvXt60vq6xNj2PQK_sBULnaVGvSQ9vXQ,5872
 airtrain/integrations/aws/__init__.py,sha256=3x7v2NxpAfI-U-YgwQeH5PtsmUrNLPMfLyUGFLiBjbs,155
 airtrain/integrations/aws/credentials.py,sha256=nN-daKAl7qOb_VdRpsThG8gN5GeSUkx-ji5E_gF_vYw,1444
-airtrain/integrations/aws/skills.py,sha256=TQiMXeXRRcJ14fe8Xi7Uk20iS6_INbcznuLGtMorcKY,3870
+airtrain/integrations/aws/skills.py,sha256=2l16Y5zYeNd9trrPca6Rbhvl6a-GJBuCQMu7RqX9txo,3872
 airtrain/integrations/cerebras/__init__.py,sha256=zAD-qV38OzHhMCz1z-NvjjqcYEhURbm8RWTOKHNqbew,174
 airtrain/integrations/cerebras/credentials.py,sha256=KDEH4r8FGT68L9p34MLZWK65wq_a703pqIF3ODaSbts,694
-airtrain/integrations/cerebras/skills.py,sha256=hmqcnF-nkFk5YJVf8f-TiKBfb8kYCfnC30W67VZ7CKU,4922
+airtrain/integrations/cerebras/skills.py,sha256=BJEb_7TglCYAukD3kcx37R8ibnJWdxVrBrwf3ZTYP-4,4924
+airtrain/integrations/combined/__init__.py,sha256=EL_uZs8436abeZA6NbM-gLdur7kJr4YfGGOYybKOhjc,467
+airtrain/integrations/combined/groq_fireworks_skills.py,sha256=Kz8UDU4-Rl71znz3ml9qVMRz669iWx4sUl37iafW0NU,4612
+airtrain/integrations/combined/list_models_factory.py,sha256=tDhfu7-MdAgPXCDkASNezSu4gfNjIC7Zf4XLw_Rx0sk,6500
 airtrain/integrations/fireworks/__init__.py,sha256=GstUg0rYC-7Pg0DVbDXwL5eO1hp3WCSfroWazbGpfi0,545
-airtrain/integrations/fireworks/completion_skills.py,sha256=G657xWd7izLOxXq7RmqdupBF4DHqXQgXuhQ-MW7mtqc,5613
-airtrain/integrations/fireworks/conversation_manager.py,sha256=m6VEHijqpYEYawkKhuHtb8RQxw4kxGWFWdbSK6zGuro,3704
+airtrain/integrations/fireworks/completion_skills.py,sha256=zxx7aNlum9scQMri5Ek0qN8VfAomhyUp3u8JJo_AFWM,5615
+airtrain/integrations/fireworks/conversation_manager.py,sha256=ifscKHYKWM_NDElin-oTzpRhyoh6pzBnklmMuH5geOY,3706
 airtrain/integrations/fireworks/credentials.py,sha256=eeV9y_4pTe8LZX02I7kfA_YNY2D7MSbFl7JEZVn22zQ,864
 airtrain/integrations/fireworks/list_models.py,sha256=o4fP0K3qstBopO7va2LysLp4_KUf5Iz_YROrYkaNtVs,4686
 airtrain/integrations/fireworks/models.py,sha256=yo4xtweSi4qQftg04r4naRddx3KjU9Jluzqf5C7V9f4,4626
-airtrain/integrations/fireworks/requests_skills.py,sha256=c84Vy_4EcBrwJfp3jqizzlcja_LsEtvWh59qiaIjukg,8233
-airtrain/integrations/fireworks/skills.py,sha256=o9OY69cC10P8BtBBYRYLCyR_GwxmNlF6YhnrXiNS53o,7154
-airtrain/integrations/fireworks/structured_completion_skills.py,sha256=-AJTaOFC8vkFiEjHW24VL8ymcNSVbhZp6xb4enkL95U,6620
-airtrain/integrations/fireworks/structured_requests_skills.py,sha256=FgUdWb6_GI2ZBWhK2wp-WqKZUkwCkKNBBjYcRkHtjog,11850
-airtrain/integrations/fireworks/structured_skills.py,sha256=BZaLqSOTC11QdZ4kDORS4JnwF_YXBAa-IiwQ5dJiHXw,3895
+airtrain/integrations/fireworks/requests_skills.py,sha256=h6HRV5dGvV7t3zyjD-awW47RyeDbu8onNevhcgSSy94,8235
+airtrain/integrations/fireworks/skills.py,sha256=Ns1tXXTVtTeeVYadzm4dnmmOboo430WTMu2o56oWTDc,7156
+airtrain/integrations/fireworks/structured_completion_skills.py,sha256=airYakYWXzYRS9nfNfrH90N3eeN8YW7GaY3ygLSiBO8,6622
+airtrain/integrations/fireworks/structured_requests_skills.py,sha256=uQR-nygtWmdGTwvU-aUdMNOMit_PiBVPYRa80ZloHLs,11852
+airtrain/integrations/fireworks/structured_skills.py,sha256=1wZ_7QDUhKWCSv_1lSEF6VnAqEeEA3jWHq7n0fWicgw,3897
 airtrain/integrations/google/__init__.py,sha256=ElwgcXfbg_gGMm6zbkMXCQPFKZUb-yTJk986o19A7Cs,214
 airtrain/integrations/google/credentials.py,sha256=KSvWNqW8Mjr4MkysRvUqlrOSGdShNIe5u2OPO6vRrWY,2047
 airtrain/integrations/google/skills.py,sha256=ytsoksCY4qbfRO9Brnxhc2694fAj0ytnHX20SXS_FOM,4547
 airtrain/integrations/groq/__init__.py,sha256=B_X2fXbsJfFD6GquKeVCsEJjwd9Ygbq1uEHlV4Jy7YE,154
 airtrain/integrations/groq/credentials.py,sha256=bdTHykcIeaQ7td8KZlQBPfEFAkvJuxk2f_cbTLPD_I4,844
-airtrain/integrations/groq/skills.py,sha256=qFyxC_2xZYnByAPo5p2aHbrqhdHYCoIdvDRAauSfnjk,4821
+airtrain/integrations/groq/skills.py,sha256=XNwGE2fjb9MDth3NI5uqSiEQYLsLBuCFF7YTu_xoTug,4823
 airtrain/integrations/ollama/__init__.py,sha256=zMHBsGzViVrvxAeJmfq6r-ZfSE6Dy5QcKLhe4d5fEcM,164
 airtrain/integrations/ollama/credentials.py,sha256=D7O4kUAb_VHs5s1ncUN9Ezhu5PvLfgj3RifAkB9sEZk,940
-airtrain/integrations/ollama/skills.py,sha256=M_Un8D5VJ5XtPEq9IClzqV3jCPBoFTSm2ve6EO8W2JU,1556
-airtrain/integrations/openai/__init__.py,sha256=w5V7lxvrKtrrjyqGoppEKg9ORKKQ2cxaLOpgCZdm_H8,541
-airtrain/integrations/openai/chinese_assistant.py,sha256=MMhv4NBOoEQ0O22ZZtP255rd5ajHC9l6FPWIjpqxBOA,1581
+airtrain/integrations/ollama/skills.py,sha256=QHEvIrFmuwFuC3ZAmy6xL2hNNGZWii1z9Y884JuOvnI,1558
+airtrain/integrations/openai/__init__.py,sha256=LYLxgDOAMUhw0ChBqj7XAJSTMNt9JiS2hHJDnJWS6IE,807
+airtrain/integrations/openai/chinese_assistant.py,sha256=F8bMeUUDly7BYG6wO648cAEIj5S_frVE5tm1Xno63Gc,1583
 airtrain/integrations/openai/credentials.py,sha256=NfRyp1QgEtgm8cxt2-BOLq-6d0X-Pcm80NnfHM8p0FY,1470
+airtrain/integrations/openai/list_models.py,sha256=vg8pZwLZ3F2Fx42X18WykpJOzZD9JG-2KJi49XWgSKo,4121
 airtrain/integrations/openai/models_config.py,sha256=W9mu_z9tCC4ZUKHSJ6Hk4X09TRZLqEhT7TtRY5JEk5g,8007
-airtrain/integrations/openai/skills.py,sha256=1dvRJYrnU2hOmGRlkHBtyR6P8D7aIwHZfUKmjlReWrQ,12821
+airtrain/integrations/openai/skills.py,sha256=kEDe5q0Zlv_X-JGOYtb552ktb3aQQYVUYczVwMH0jxA,12823
 airtrain/integrations/sambanova/__init__.py,sha256=dp_263iOckM_J9pOEvyqpf3FrejD6-_x33r0edMCTe0,179
 airtrain/integrations/sambanova/credentials.py,sha256=JyN8sbMCoXuXAjim46aI3LTicBijoemS7Ao0rn4yBJU,824
 airtrain/integrations/sambanova/skills.py,sha256=SZ_GAimMiOCILiNkzyhNflyRR6bdC5r0Tnog19K8geU,4997
@@ -60,14 +65,14 @@ airtrain/integrations/together/image_models_config.py,sha256=JlCozrphI9zE4uYpGfj
 airtrain/integrations/together/image_skill.py,sha256=wQ8wSzfL-QHpM_esYGLNXf8ciOPPsz-QJw6zSrxZT68,5214
 airtrain/integrations/together/list_models.py,sha256=_QLGqweBiK6saz3n4xTQRmXSSs-qGFnV9kma-MSaE9o,2520
 airtrain/integrations/together/models.py,sha256=q5KsouOK7IvyzGZ7nhSjTpZw-CcLfPghJr6o_UU9uMo,3652
-airtrain/integrations/together/models_config.py,sha256=XMKp0Oq1nWWnMMdNAZxkFXmJaURwWrwLE18kFXsMsRw,8829
+airtrain/integrations/together/models_config.py,sha256=lCPouXjKa49lebGbMaaL2SU-CMYxOc-dJviUOir2e_w,8831
 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=8DwkexMJu1Gm6QmNDfNasYStQ31QsXBbFP99zR-YCf0,7598
 airtrain/integrations/together/vision_models_config.py,sha256=m28HwYDk2Kup_J-a1FtynIa2ZVcbl37kltfoHnK8zxs,1544
-airtrain-0.1.43.dist-info/METADATA,sha256=FM114Wh1dKSK_-PH7H2b9vHIXdc7-EoHvSsOClJWDkk,5375
-airtrain-0.1.43.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-airtrain-0.1.43.dist-info/entry_points.txt,sha256=rrJ36IUsyq6n1dSfTWXqVAgpQLPRWDfCqwd6_3B-G0U,52
-airtrain-0.1.43.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
-airtrain-0.1.43.dist-info/RECORD,,
+airtrain-0.1.45.dist-info/METADATA,sha256=PXbwhd_qiTdj56DAAG1iwZ9CBgThiTRaaf4xw6A_G7o,5375
+airtrain-0.1.45.dist-info/WHEEL,sha256=tTnHoFhvKQHCh4jz3yCn0WPTYIy7wXx3CJtJ7SJGV7c,91
+airtrain-0.1.45.dist-info/entry_points.txt,sha256=rrJ36IUsyq6n1dSfTWXqVAgpQLPRWDfCqwd6_3B-G0U,52
+airtrain-0.1.45.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.45.dist-info/RECORD,,

{airtrain-0.1.43.dist-info → airtrain-0.1.45.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (76.1.0)
+Generator: setuptools (77.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any