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

airtrain/integrations/__init__.py

@@ -22,6 +22,10 @@ from .ollama.skills import OllamaChatSkill
 from .sambanova.skills import SambanovaChatSkill
 from .cerebras.skills import CerebrasChatSkill
 
+# Model configurations
+from .openai.models_config import OPENAI_MODELS, OpenAIModelConfig
+from .anthropic.models_config import ANTHROPIC_MODELS, AnthropicModelConfig
+
 __all__ = [
     # Credentials
     "OpenAICredentials",
@@ -38,10 +42,15 @@ __all__ = [
     "OpenAIParserSkill",
     "AnthropicChatSkill",
     "AWSBedrockSkill",
-    "VertexAISkill",
+    "GoogleChatSkill",
     "GroqChatSkill",
     "TogetherAIChatSkill",
     "OllamaChatSkill",
     "SambanovaChatSkill",
     "CerebrasChatSkill",
+    # Model configurations
+    "OPENAI_MODELS",
+    "OpenAIModelConfig",
+    "ANTHROPIC_MODELS",
+    "AnthropicModelConfig",
 ]

airtrain/integrations/anthropic/__init__.py

@@ -2,10 +2,22 @@
 
 from .credentials import AnthropicCredentials
 from .skills import AnthropicChatSkill, AnthropicInput, AnthropicOutput
+from .models_config import (
+    ANTHROPIC_MODELS,
+    AnthropicModelConfig,
+    get_model_config,
+    get_default_model,
+    calculate_cost,
+)
 
 __all__ = [
     "AnthropicCredentials",
     "AnthropicChatSkill",
     "AnthropicInput",
     "AnthropicOutput",
+    "ANTHROPIC_MODELS",
+    "AnthropicModelConfig",
+    "get_model_config",
+    "get_default_model",
+    "calculate_cost",
 ]

airtrain/integrations/anthropic/models_config.py

@@ -0,0 +1,100 @@
+from typing import Dict, NamedTuple, Optional
+from decimal import Decimal
+
+
+class AnthropicModelConfig(NamedTuple):
+    display_name: str
+    base_model: str
+    input_price: Decimal
+    cached_write_price: Optional[Decimal]
+    cached_read_price: Optional[Decimal]
+    output_price: Decimal
+
+
+ANTHROPIC_MODELS: Dict[str, AnthropicModelConfig] = {
+    "claude-3-7-sonnet": AnthropicModelConfig(
+        display_name="Claude 3.7 Sonnet",
+        base_model="claude-3-7-sonnet",
+        input_price=Decimal("3.00"),
+        cached_write_price=Decimal("3.75"),
+        cached_read_price=Decimal("0.30"),
+        output_price=Decimal("15.00"),
+    ),
+    "claude-3-5-haiku": AnthropicModelConfig(
+        display_name="Claude 3.5 Haiku",
+        base_model="claude-3-5-haiku",
+        input_price=Decimal("0.80"),
+        cached_write_price=Decimal("1.00"),
+        cached_read_price=Decimal("0.08"),
+        output_price=Decimal("4.00"),
+    ),
+    "claude-3-opus": AnthropicModelConfig(
+        display_name="Claude 3 Opus",
+        base_model="claude-3-opus",
+        input_price=Decimal("15.00"),
+        cached_write_price=Decimal("18.75"),
+        cached_read_price=Decimal("1.50"),
+        output_price=Decimal("75.00"),
+    ),
+    "claude-3-sonnet": AnthropicModelConfig(
+        display_name="Claude 3 Sonnet",
+        base_model="claude-3-sonnet",
+        input_price=Decimal("3.00"),
+        cached_write_price=Decimal("3.75"),
+        cached_read_price=Decimal("0.30"),
+        output_price=Decimal("15.00"),
+    ),
+    "claude-3-haiku": AnthropicModelConfig(
+        display_name="Claude 3 Haiku",
+        base_model="claude-3-haiku",
+        input_price=Decimal("0.25"),
+        cached_write_price=Decimal("0.31"),
+        cached_read_price=Decimal("0.025"),
+        output_price=Decimal("1.25"),
+    ),
+}
+
+
+def get_model_config(model_id: str) -> AnthropicModelConfig:
+    """Get model configuration by model ID"""
+    if model_id not in ANTHROPIC_MODELS:
+        raise ValueError(f"Model {model_id} not found in Anthropic models")
+    return ANTHROPIC_MODELS[model_id]
+
+
+def get_default_model() -> str:
+    """Get the default model ID"""
+    return "claude-3-sonnet"
+
+
+def calculate_cost(
+    model_id: str, 
+    input_tokens: int, 
+    output_tokens: int, 
+    use_cached: bool = False,
+    cache_type: str = "read"
+) -> Decimal:
+    """Calculate cost for token usage
+    
+    Args:
+        model_id: The model ID to calculate costs for
+        input_tokens: Number of input tokens
+        output_tokens: Number of output tokens
+        use_cached: Whether to use cached pricing
+        cache_type: Either "read" or "write" for cached pricing type
+    """
+    config = get_model_config(model_id)
+    
+    if not use_cached:
+        input_cost = config.input_price * Decimal(str(input_tokens))
+    else:
+        if cache_type == "read" and config.cached_read_price is not None:
+            input_cost = config.cached_read_price * Decimal(str(input_tokens))
+        elif cache_type == "write" and config.cached_write_price is not None:
+            input_cost = config.cached_write_price * Decimal(str(input_tokens))
+        else:
+            input_cost = config.input_price * Decimal(str(input_tokens))
+    
+    output_cost = config.output_price * Decimal(str(output_tokens))
+    
+    return (input_cost + output_cost) / Decimal("1000") 

airtrain/integrations/fireworks/__init__.py

@@ -2,10 +2,20 @@
 
 from .credentials import FireworksCredentials
 from .skills import FireworksChatSkill, FireworksInput, FireworksOutput
+from .list_models import (
+    FireworksListModelsSkill,
+    FireworksListModelsInput,
+    FireworksListModelsOutput,
+)
+from .models import FireworksModel
 
 __all__ = [
     "FireworksCredentials",
     "FireworksChatSkill",
     "FireworksInput",
     "FireworksOutput",
+    "FireworksListModelsSkill",
+    "FireworksListModelsInput",
+    "FireworksListModelsOutput",
+    "FireworksModel",
 ]

airtrain/integrations/fireworks/credentials.py

@@ -1,4 +1,4 @@
-from pydantic import SecretStr, BaseModel
+from pydantic import SecretStr, BaseModel, Field
 from typing import Optional
 import os
 
@@ -6,7 +6,15 @@ import os
 class FireworksCredentials(BaseModel):
     """Credentials for Fireworks AI API"""
 
-    fireworks_api_key: SecretStr
+    fireworks_api_key: SecretStr = Field(..., min_length=1)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the credentials."""
+        return f"FireworksCredentials(fireworks_api_key=SecretStr('**********'))"
+
+    def __str__(self) -> str:
+        """Return a string representation of the credentials."""
+        return self.__repr__()
 
     @classmethod
     def from_env(cls) -> "FireworksCredentials":

airtrain/integrations/fireworks/list_models.py

@@ -0,0 +1,128 @@
+from typing import Optional, List
+import requests
+from pydantic import Field
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+from .models import FireworksModel
+
+
+class FireworksListModelsInput(InputSchema):
+    """Schema for Fireworks AI list models input"""
+
+    account_id: str = Field(..., description="The Account Id")
+    page_size: Optional[int] = Field(
+        default=50, 
+        description=(
+            "The maximum number of models to return. The maximum page_size is 200, "
+            "values above 200 will be coerced to 200."
+        ),
+        le=200
+    )
+    page_token: Optional[str] = Field(
+        default=None,
+        description=(
+            "A page token, received from a previous ListModels call. Provide this "
+            "to retrieve the subsequent page. When paginating, all other parameters "
+            "provided to ListModels must match the call that provided the page token."
+        )
+    )
+    filter: Optional[str] = Field(
+        default=None,
+        description=(
+            "Only model satisfying the provided filter (if specified) will be "
+            "returned. See https://google.aip.dev/160 for the filter grammar."
+        )
+    )
+    order_by: Optional[str] = Field(
+        default=None,
+        description=(
+            "A comma-separated list of fields to order by. e.g. \"foo,bar\" "
+            "The default sort order is ascending. To specify a descending order for a "
+            "field, append a \" desc\" suffix. e.g. \"foo desc,bar\" "
+            "Subfields are specified with a \".\" character. e.g. \"foo.bar\" "
+            "If not specified, the default order is by \"name\"."
+        )
+    )
+
+
+class FireworksListModelsOutput(OutputSchema):
+    """Schema for Fireworks AI list models output"""
+
+    models: List[FireworksModel] = Field(
+        default_factory=list, 
+        description="List of Fireworks models"
+    )
+    next_page_token: Optional[str] = Field(
+        default=None, 
+        description="Token for retrieving the next page of results"
+    )
+    total_size: Optional[int] = Field(
+        default=None, 
+        description="Total number of models available"
+    )
+
+
+class FireworksListModelsSkill(
+    Skill[FireworksListModelsInput, FireworksListModelsOutput]
+):
+    """Skill for listing Fireworks AI models"""
+
+    input_schema = FireworksListModelsInput
+    output_schema = FireworksListModelsOutput
+    
+    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/v1"
+    
+    def process(
+        self, input_data: FireworksListModelsInput
+    ) -> FireworksListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            # Build the URL
+            url = f"{self.base_url}/accounts/{input_data.account_id}/models"
+            
+            # Prepare query parameters
+            params = {}
+            if input_data.page_size:
+                params["pageSize"] = input_data.page_size
+            if input_data.page_token:
+                params["pageToken"] = input_data.page_token
+            if input_data.filter:
+                params["filter"] = input_data.filter
+            if input_data.order_by:
+                params["orderBy"] = input_data.order_by
+            
+            # Make the request
+            headers = {
+                "Authorization": (
+                    f"Bearer {self.credentials.fireworks_api_key.get_secret_value()}"
+                )
+            }
+            
+            response = requests.get(url, headers=headers, params=params)
+            response.raise_for_status()
+            
+            # Parse the response
+            result = response.json()
+            
+            # Convert the models to FireworksModel objects
+            models = []
+            for model_data in result.get("models", []):
+                models.append(FireworksModel(**model_data))
+            
+            # Return the output
+            return FireworksListModelsOutput(
+                models=models,
+                next_page_token=result.get("nextPageToken"),
+                total_size=result.get("totalSize")
+            )
+            
+        except requests.RequestException as e:
+            raise ProcessingError(f"Failed to list Fireworks models: {str(e)}")
+        except Exception as e:
+            raise ProcessingError(f"Error listing Fireworks models: {str(e)}") 

airtrain/integrations/fireworks/models.py

@@ -25,3 +25,115 @@ class FireworksResponse(BaseModel):
     created: int
     model: str
     usage: FireworksUsage
+
+
+class FireworksModelStatus(BaseModel):
+    """Schema for Fireworks model status"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksModelBaseDetails(BaseModel):
+    """Schema for Fireworks base model details"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksPeftDetails(BaseModel):
+    """Schema for Fireworks PEFT details"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksConversationConfig(BaseModel):
+    """Schema for Fireworks conversation configuration"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksModelDeployedRef(BaseModel):
+    """Schema for Fireworks deployed model reference"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksDeprecationDate(BaseModel):
+    """Schema for Fireworks deprecation date"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksModel(BaseModel):
+    """Schema for a Fireworks model"""
+    
+    name: str
+    display_name: Optional[str] = None
+    description: Optional[str] = None
+    create_time: Optional[str] = None
+    created_by: Optional[str] = None
+    state: Optional[str] = None
+    status: Optional[Dict[str, Any]] = None
+    kind: Optional[str] = None
+    github_url: Optional[str] = None
+    hugging_face_url: Optional[str] = None
+    base_model_details: Optional[Dict[str, Any]] = None
+    peft_details: Optional[Dict[str, Any]] = None
+    teft_details: Optional[Dict[str, Any]] = None
+    public: Optional[bool] = None
+    conversation_config: Optional[Dict[str, Any]] = None
+    context_length: Optional[int] = None
+    supports_image_input: Optional[bool] = None
+    supports_tools: Optional[bool] = None
+    imported_from: Optional[str] = None
+    fine_tuning_job: Optional[str] = None
+    default_draft_model: Optional[str] = None
+    default_draft_token_count: Optional[int] = None
+    precisions: Optional[List[str]] = None
+    deployed_model_refs: Optional[List[Dict[str, Any]]] = None
+    cluster: Optional[str] = None
+    deprecation_date: Optional[Dict[str, Any]] = None
+    calibrated: Optional[bool] = None
+    tunable: Optional[bool] = None
+    supports_lora: Optional[bool] = None
+    use_hf_apply_chat_template: Optional[bool] = None
+
+
+class ListModelsInput(BaseModel):
+    """Schema for listing Fireworks models input"""
+    
+    account_id: str = Field(..., description="The Account Id")
+    page_size: Optional[int] = Field(
+        default=50, 
+        description=(
+            "The maximum number of models to return. The maximum page_size is 200, "
+            "values above 200 will be coerced to 200."
+        ),
+        le=200
+    )
+    page_token: Optional[str] = Field(
+        default=None,
+        description=(
+            "A page token, received from a previous ListModels call. Provide this "
+            "to retrieve the subsequent page. When paginating, all other parameters "
+            "provided to ListModels must match the call that provided the page token."
+        )
+    )
+    filter: Optional[str] = Field(
+        default=None,
+        description=(
+            "Only model satisfying the provided filter (if specified) will be "
+            "returned. See https://google.aip.dev/160 for the filter grammar."
+        )
+    )
+    order_by: Optional[str] = Field(
+        default=None,
+        description=(
+            "A comma-separated list of fields to order by. e.g. \"foo,bar\" "
+            "The default sort order is ascending. To specify a descending order for a "
+            "field, append a \" desc\" suffix. e.g. \"foo desc,bar\" "
+            "Subfields are specified with a \".\" character. e.g. \"foo.bar\" "
+            "If not specified, the default order is by \"name\"."
+        )
+    )
+
+
+class ListModelsOutput(BaseModel):
+    """Schema for listing Fireworks models output"""
+    
+    models: List[FireworksModel]
+    next_page_token: Optional[str] = None
+    total_size: Optional[int] = None

airtrain/integrations/fireworks/skills.py

@@ -1,14 +1,10 @@
-from typing import List, Optional, Dict, Any, Generator
+from typing import List, Optional, Dict, Any, Generator, Union
 from pydantic import Field
-import requests
-from loguru import logger
 from openai import OpenAI
-from openai.types.chat import ChatCompletionChunk
 
 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):
@@ -19,9 +15,9 @@ class FireworksInput(InputSchema):
         default="You are a helpful assistant.",
         description="System prompt to guide the model's behavior",
     )
-    conversation_history: List[Dict[str, str]] = Field(
+    conversation_history: List[Dict[str, Any]] = Field(
         default_factory=list,
-        description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
+        description="List of previous conversation messages",
     )
     model: str = Field(
         default="accounts/fireworks/models/deepseek-r1",
@@ -40,6 +36,20 @@ class FireworksInput(InputSchema):
         default=False,
         description="Whether to stream the response token by token",
     )
+    tools: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description=(
+            "A list of tools the model may use. "
+            "Currently only functions supported."
+        ),
+    )
+    tool_choice: Optional[Union[str, Dict[str, Any]]] = Field(
+        default=None,
+        description=(
+            "Controls which tool is called by the model. "
+            "'none', 'auto', or specific tool."
+        ),
+    )
 
 
 class FireworksOutput(OutputSchema):
@@ -48,6 +58,9 @@ class FireworksOutput(OutputSchema):
     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")
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        default=None, description="Tool calls generated by the model"
+    )
 
 
 class FireworksChatSkill(Skill[FireworksInput, FireworksOutput]):
@@ -65,7 +78,7 @@ class FireworksChatSkill(Skill[FireworksInput, FireworksOutput]):
             api_key=self.credentials.fireworks_api_key.get_secret_value(),
         )
 
-    def _build_messages(self, input_data: FireworksInput) -> List[Dict[str, str]]:
+    def _build_messages(self, input_data: FireworksInput) -> List[Dict[str, Any]]:
         """Build messages list from input data including conversation history."""
         messages = [{"role": "system", "content": input_data.system_prompt}]
 
@@ -104,8 +117,8 @@ class FireworksChatSkill(Skill[FireworksInput, FireworksOutput]):
                 for chunk in self.process_stream(input_data):
                     response_chunks.append(chunk)
                 response = "".join(response_chunks)
-            else:
-                # For non-streaming, use regular completion
+                
+                # Create completion object for usage stats
                 messages = self._build_messages(input_data)
                 completion = self.client.chat.completions.create(
                     model=input_data.model,
@@ -114,7 +127,44 @@ class FireworksChatSkill(Skill[FireworksInput, FireworksOutput]):
                     max_tokens=input_data.max_tokens,
                     stream=False,
                 )
-                response = completion.choices[0].message.content
+            else:
+                # For non-streaming, use regular completion
+                messages = self._build_messages(input_data)
+                
+                # Prepare API call parameters
+                api_params = {
+                    "model": input_data.model,
+                    "messages": messages,
+                    "temperature": input_data.temperature,
+                    "max_tokens": input_data.max_tokens,
+                    "stream": False,
+                }
+                
+                # Add tools and tool_choice if provided
+                if input_data.tools:
+                    api_params["tools"] = input_data.tools
+                
+                if input_data.tool_choice:
+                    api_params["tool_choice"] = input_data.tool_choice
+                
+                completion = self.client.chat.completions.create(**api_params)
+                response = completion.choices[0].message.content or ""
+            
+            # Check for tool calls in the response
+            tool_calls = None
+            if (hasattr(completion.choices[0].message, "tool_calls") and
+                    completion.choices[0].message.tool_calls):
+                tool_calls = [
+                    {
+                        "id": tool_call.id,
+                        "type": tool_call.type,
+                        "function": {
+                            "name": tool_call.function.name,
+                            "arguments": tool_call.function.arguments
+                        }
+                    }
+                    for tool_call in completion.choices[0].message.tool_calls
+                ]
 
             return FireworksOutput(
                 response=response,
@@ -124,6 +174,7 @@ class FireworksChatSkill(Skill[FireworksInput, FireworksOutput]):
                     "prompt_tokens": completion.usage.prompt_tokens,
                     "completion_tokens": completion.usage.completion_tokens,
                 },
+                tool_calls=tool_calls
             )
 
         except Exception as e:

airtrain/integrations/fireworks/structured_completion_skills.py

@@ -1,14 +1,13 @@
-from typing import Type, TypeVar, Optional, List, Dict, Any, Generator
+from typing import Any, Dict, Generator, List, Optional, Type, TypeVar
 from pydantic import BaseModel, Field
 import requests
 import json
-from loguru import logger
 
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
 from .credentials import FireworksCredentials
 
-ResponseT = TypeVar("ResponseT", bound=BaseModel)
+ResponseT = TypeVar("ResponseT")
 
 
 class FireworksStructuredCompletionInput(InputSchema):
@@ -26,7 +25,7 @@ class FireworksStructuredCompletionInput(InputSchema):
     response_model: Type[ResponseT]
     stream: bool = Field(
         default=False,
-        description="Whether to stream the response",
+        description="Whether to stream the response token by token",
     )
 
     class Config:
@@ -39,6 +38,13 @@ class FireworksStructuredCompletionOutput(OutputSchema):
     parsed_response: Any
     used_model: str
     usage: Dict[str, int]
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        default=None, 
+        description=(
+            "Tool calls are not applicable for completions, "
+            "included for compatibility"
+        )
+    )
 
 
 class FireworksStructuredCompletionSkill(