airtrain 0.1.43__tar.gz → 0.1.44__tar.gz

{airtrain-0.1.43 → airtrain-0.1.44}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: airtrain
-Version: 0.1.43
+Version: 0.1.44
 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 → airtrain-0.1.44}/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.44"
 
 # Core imports
 from .core.skills import Skill, ProcessingError

{airtrain-0.1.43 → airtrain-0.1.44}/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-0.1.44/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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-0.1.44/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-0.1.43 → airtrain-0.1.44}/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-0.1.43 → airtrain-0.1.44}/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 → airtrain-0.1.44}/airtrain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: airtrain
-Version: 0.1.43
+Version: 0.1.44
 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 → airtrain-0.1.44}/airtrain.egg-info/SOURCES.txt

@@ -28,6 +28,7 @@ airtrain/core/skills.py
 airtrain/integrations/__init__.py
 airtrain/integrations/anthropic/__init__.py
 airtrain/integrations/anthropic/credentials.py
+airtrain/integrations/anthropic/list_models.py
 airtrain/integrations/anthropic/models_config.py
 airtrain/integrations/anthropic/skills.py
 airtrain/integrations/aws/__init__.py
@@ -59,6 +60,7 @@ airtrain/integrations/ollama/skills.py
 airtrain/integrations/openai/__init__.py
 airtrain/integrations/openai/chinese_assistant.py
 airtrain/integrations/openai/credentials.py
+airtrain/integrations/openai/list_models.py
 airtrain/integrations/openai/models_config.py
 airtrain/integrations/openai/skills.py
 airtrain/integrations/sambanova/__init__.py

{airtrain-0.1.43 → airtrain-0.1.44}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "airtrain"
-version = "0.1.43"
+version = "0.1.44"
 description = "A platform for building and deploying AI agents with structured skills"
 readme = "README.md"
 requires-python = ">=3.8"