airtrain 0.1.39__py3-none-any.whl → 0.1.40__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.40"
 
 # 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(

airtrain/integrations/fireworks/structured_requests_skills.py

@@ -1,5 +1,5 @@
-from typing import Type, TypeVar, Optional, List, Dict, Any, Generator
-from pydantic import BaseModel, Field
+from typing import Type, TypeVar, Optional, List, Dict, Any, Generator, Union
+from pydantic import BaseModel, Field, create_model
 import requests
 import json
 from loguru import logger
@@ -20,7 +20,7 @@ class FireworksStructuredRequestInput(InputSchema):
         default="You are a helpful assistant that provides structured data.",
         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",
     )
@@ -34,8 +34,21 @@ class FireworksStructuredRequestInput(InputSchema):
     max_tokens: int = Field(default=4096, description="Maximum tokens in response")
     response_model: Type[ResponseT]
     stream: bool = Field(
-        default=False,
-        description="Whether to stream the response",
+        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 Config:
@@ -49,6 +62,9 @@ class FireworksStructuredRequestOutput(OutputSchema):
     used_model: str
     usage: Dict[str, int]
     reasoning: Optional[str] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        default=None, description="Tool calls generated by the model"
+    )
 
 
 class FireworksStructuredRequestSkill(
@@ -72,7 +88,7 @@ class FireworksStructuredRequestSkill(
 
     def _build_messages(
         self, input_data: FireworksStructuredRequestInput
-    ) -> List[Dict[str, str]]:
+    ) -> List[Dict[str, Any]]:
         """Build messages list from input data including conversation history."""
         messages = [{"role": "system", "content": input_data.system_prompt}]
 
@@ -86,24 +102,24 @@ class FireworksStructuredRequestSkill(
         self, input_data: FireworksStructuredRequestInput
     ) -> Dict[str, Any]:
         """Build the request payload."""
-        return {
+        payload = {
             "model": input_data.model,
             "messages": self._build_messages(input_data),
             "temperature": input_data.temperature,
             "max_tokens": input_data.max_tokens,
             "stream": input_data.stream,
-            "response_format": {
-                "type": "json_object",
-                "schema": {
-                    **input_data.response_model.model_json_schema(),
-                    "required": [
-                        field
-                        for field, _ in input_data.response_model.model_fields.items()
-                    ],
-                },
-            },
+            "response_format": {"type": "json_object"},
         }
 
+        # Add tool-related parameters if provided
+        if input_data.tools:
+            payload["tools"] = input_data.tools
+        
+        if input_data.tool_choice:
+            payload["tool_choice"] = input_data.tool_choice
+
+        return payload
+
     def process_stream(
         self, input_data: FireworksStructuredRequestInput
     ) -> Generator[Dict[str, Any], None, None]:
@@ -131,6 +147,10 @@ class FireworksStructuredRequestSkill(
                         continue
 
             # Once complete, parse the full response with think tags
+            if not json_buffer:
+                # If no data was collected, raise error
+                raise ProcessingError("No data received from Fireworks API")
+
             complete_response = "".join(json_buffer)
             reasoning, json_str = self._parse_response_content(complete_response)
 
@@ -177,37 +197,94 @@ class FireworksStructuredRequestSkill(
 
                 if parsed_response is None:
                     raise ProcessingError("Failed to parse streamed response")
+                
+                # Make a non-streaming call to get tool calls if tools were provided
+                tool_calls = None
+                if input_data.tools:
+                    # Create a non-streaming request to get tool calls
+                    non_stream_payload = self._build_payload(input_data)
+                    non_stream_payload["stream"] = False
+                    
+                    response = requests.post(
+                        self.BASE_URL,
+                        headers=self.headers,
+                        data=json.dumps(non_stream_payload),
+                    )
+                    response.raise_for_status()
+                    result = response.json()
+                    
+                    # Check for tool calls
+                    if (result["choices"][0]["message"].get("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 result["choices"][0]["message"]["tool_calls"]
+                        ]
 
                 return FireworksStructuredRequestOutput(
                     parsed_response=parsed_response,
                     used_model=input_data.model,
-                    usage={},  # Usage stats not available in streaming mode
+                    usage={"total_tokens": 0},  # Can't get usage stats from streaming
                     reasoning=reasoning,
+                    tool_calls=tool_calls,
                 )
             else:
                 # For non-streaming, use regular request
                 payload = self._build_payload(input_data)
+                payload["stream"] = False  # Ensure it's not streaming
+
                 response = requests.post(
                     self.BASE_URL, headers=self.headers, data=json.dumps(payload)
                 )
                 response.raise_for_status()
-                data = response.json()
-
-                response_content = data["choices"][0]["message"]["content"]
-
-                # Parse the response content to extract reasoning and JSON
-                reasoning, json_str = self._parse_response_content(response_content)
-
-                # Parse the JSON string into the specified model
-                parsed_response = input_data.response_model.model_validate_json(
-                    json_str
-                )
+                result = response.json()
+
+                # Get the content from the response
+                if "choices" not in result or not result["choices"]:
+                    raise ProcessingError("Invalid response format from Fireworks API")
+
+                content = result["choices"][0]["message"].get("content", "")
+                
+                # Check for tool calls
+                tool_calls = None
+                if (result["choices"][0]["message"].get("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 result["choices"][0]["message"]["tool_calls"]
+                    ]
+
+                # Parse the response content
+                reasoning, json_str = self._parse_response_content(content)
+                try:
+                    parsed_response = input_data.response_model.model_validate_json(
+                        json_str
+                    )
+                except Exception as e:
+                    raise ProcessingError(f"Failed to parse JSON response: {str(e)}")
 
                 return FireworksStructuredRequestOutput(
                     parsed_response=parsed_response,
                     used_model=input_data.model,
-                    usage=data["usage"],
-                    reasoning=reasoning,  # Add reasoning to output if present
+                    usage={
+                        "total_tokens": result["usage"]["total_tokens"],
+                        "prompt_tokens": result["usage"]["prompt_tokens"],
+                        "completion_tokens": result["usage"]["completion_tokens"],
+                    },
+                    reasoning=reasoning,
+                    tool_calls=tool_calls,
                 )
 
         except Exception as e:

airtrain/integrations/openai/__init__.py

@@ -5,6 +5,9 @@ from .skills import (
     OpenAIOutput,
     OpenAIParserInput,
     OpenAIParserOutput,
+    OpenAIEmbeddingsSkill,
+    OpenAIEmbeddingsInput,
+    OpenAIEmbeddingsOutput,
 )
 from .credentials import OpenAICredentials
 
@@ -16,4 +19,7 @@ __all__ = [
     "OpenAIParserOutput",
     "OpenAICredentials",
     "OpenAIOutput",
+    "OpenAIEmbeddingsSkill",
+    "OpenAIEmbeddingsInput",
+    "OpenAIEmbeddingsOutput",
 ]

airtrain/integrations/openai/models_config.py

@@ -11,6 +11,20 @@ class OpenAIModelConfig(NamedTuple):
 
 
 OPENAI_MODELS: Dict[str, OpenAIModelConfig] = {
+    "gpt-4.5-preview": OpenAIModelConfig(
+        display_name="GPT-4.5 Preview",
+        base_model="gpt-4.5-preview",
+        input_price=Decimal("75.00"),
+        cached_input_price=Decimal("37.50"),
+        output_price=Decimal("150.00"),
+    ),
+    "gpt-4.5-preview-2025-02-27": OpenAIModelConfig(
+        display_name="GPT-4.5 Preview (2025-02-27)",
+        base_model="gpt-4.5-preview",
+        input_price=Decimal("75.00"),
+        cached_input_price=Decimal("37.50"),
+        output_price=Decimal("150.00"),
+    ),
     "gpt-4o": OpenAIModelConfig(
         display_name="GPT-4 Optimized",
         base_model="gpt-4o",
@@ -25,69 +39,160 @@ OPENAI_MODELS: Dict[str, OpenAIModelConfig] = {
         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"),
+    "gpt-4o-audio-preview": 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("15.00"),
+        output_price=Decimal("10.00"),
     ),
     "gpt-4o-audio-preview-2024-12-17": OpenAIModelConfig(
-        display_name="GPT-4 Optimized Audio Preview",
+        display_name="GPT-4 Optimized Audio Preview (2024-12-17)",
         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(
+    "gpt-4o-realtime-preview": 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(
+    "gpt-4o-realtime-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Realtime Preview (2024-12-17)",
+        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": 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(
+    "gpt-4o-mini-2024-07-18": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini (2024-07-18)",
+        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": 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(
+    "gpt-4o-mini-audio-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Audio Preview (2024-12-17)",
+        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": 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(
+    "gpt-4o-mini-realtime-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Realtime Preview (2024-12-17)",
+        base_model="gpt-4o-mini-realtime-preview",
+        input_price=Decimal("0.60"),
+        cached_input_price=Decimal("0.30"),
+        output_price=Decimal("2.40"),
+    ),
+    "o1": 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(
+    "o1-2024-12-17": OpenAIModelConfig(
+        display_name="O1 (2024-12-17)",
+        base_model="o1",
+        input_price=Decimal("15.00"),
+        cached_input_price=Decimal("7.50"),
+        output_price=Decimal("60.00"),
+    ),
+    "o3-mini": 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(
+    "o3-mini-2025-01-31": OpenAIModelConfig(
+        display_name="O3 Mini (2025-01-31)",
+        base_model="o3-mini",
+        input_price=Decimal("1.10"),
+        cached_input_price=Decimal("0.55"),
+        output_price=Decimal("4.40"),
+    ),
+    "o1-mini": 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"),
     ),
+    "o1-mini-2024-09-12": OpenAIModelConfig(
+        display_name="O1 Mini (2024-09-12)",
+        base_model="o1-mini",
+        input_price=Decimal("1.10"),
+        cached_input_price=Decimal("0.55"),
+        output_price=Decimal("4.40"),
+    ),
+    "gpt-4o-mini-search-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Search Preview",
+        base_model="gpt-4o-mini-search-preview",
+        input_price=Decimal("0.15"),
+        cached_input_price=None,
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-mini-search-preview-2025-03-11": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Search Preview (2025-03-11)",
+        base_model="gpt-4o-mini-search-preview",
+        input_price=Decimal("0.15"),
+        cached_input_price=None,
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-search-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Search Preview",
+        base_model="gpt-4o-search-preview",
+        input_price=Decimal("2.50"),
+        cached_input_price=None,
+        output_price=Decimal("10.00"),
+    ),
+    "gpt-4o-search-preview-2025-03-11": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Search Preview (2025-03-11)",
+        base_model="gpt-4o-search-preview",
+        input_price=Decimal("2.50"),
+        cached_input_price=None,
+        output_price=Decimal("10.00"),
+    ),
+    "computer-use-preview": OpenAIModelConfig(
+        display_name="Computer Use Preview",
+        base_model="computer-use-preview",
+        input_price=Decimal("3.00"),
+        cached_input_price=None,
+        output_price=Decimal("12.00"),
+    ),
+    "computer-use-preview-2025-03-11": OpenAIModelConfig(
+        display_name="Computer Use Preview (2025-03-11)",
+        base_model="computer-use-preview",
+        input_price=Decimal("3.00"),
+        cached_input_price=None,
+        output_price=Decimal("12.00"),
+    ),
 }
 
 

airtrain/integrations/openai/skills.py

@@ -1,7 +1,8 @@
-from typing import AsyncGenerator, List, Optional, Dict, TypeVar, Type, Generator
+from typing import AsyncGenerator, List, Optional, Dict, TypeVar, Type, Generator, Union
 from pydantic import Field, BaseModel
 from openai import OpenAI, AsyncOpenAI
 from openai.types.chat import ChatCompletionChunk
+import numpy as np
 
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
@@ -232,3 +233,110 @@ class OpenAIParserSkill(Skill[OpenAIParserInput, OpenAIParserOutput]):
 
         except Exception as e:
             raise ProcessingError(f"OpenAI parsing failed: {str(e)}")
+
+
+class OpenAIEmbeddingsInput(InputSchema):
+    """Schema for OpenAI embeddings input"""
+
+    texts: Union[str, List[str]] = Field(
+        ..., description="Text or list of texts to generate embeddings for"
+    )
+    model: str = Field(
+        default="text-embedding-3-large", description="OpenAI embeddings model to use"
+    )
+    encoding_format: str = Field(
+        default="float", description="The format of the embeddings: 'float' or 'base64'"
+    )
+    dimensions: Optional[int] = Field(
+        default=None, description="Optional number of dimensions for the embeddings"
+    )
+
+
+class OpenAIEmbeddingsOutput(OutputSchema):
+    """Schema for OpenAI embeddings output"""
+
+    embeddings: List[List[float]] = Field(..., description="List of embeddings vectors")
+    used_model: str = Field(..., description="Model used for generating embeddings")
+    tokens_used: int = Field(..., description="Number of tokens used")
+
+
+class OpenAIEmbeddingsSkill(Skill[OpenAIEmbeddingsInput, OpenAIEmbeddingsOutput]):
+    """Skill for generating embeddings using OpenAI models"""
+
+    input_schema = OpenAIEmbeddingsInput
+    output_schema = OpenAIEmbeddingsOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+        self.async_client = AsyncOpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def process(self, input_data: OpenAIEmbeddingsInput) -> OpenAIEmbeddingsOutput:
+        """Generate embeddings for the input text(s)"""
+        try:
+            # Handle single text input
+            texts = (
+                [input_data.texts]
+                if isinstance(input_data.texts, str)
+                else input_data.texts
+            )
+
+            # Create embeddings
+            response = self.client.embeddings.create(
+                model=input_data.model,
+                input=texts,
+                encoding_format=input_data.encoding_format,
+                dimensions=input_data.dimensions,
+            )
+
+            # Extract embeddings
+            embeddings = [data.embedding for data in response.data]
+
+            return OpenAIEmbeddingsOutput(
+                embeddings=embeddings,
+                used_model=response.model,
+                tokens_used=response.usage.total_tokens,
+            )
+        except Exception as e:
+            raise ProcessingError(f"OpenAI embeddings generation failed: {str(e)}")
+
+    async def process_async(
+        self, input_data: OpenAIEmbeddingsInput
+    ) -> OpenAIEmbeddingsOutput:
+        """Async version of the embeddings generation"""
+        try:
+            # Handle single text input
+            texts = (
+                [input_data.texts]
+                if isinstance(input_data.texts, str)
+                else input_data.texts
+            )
+
+            # Create embeddings
+            response = await self.async_client.embeddings.create(
+                model=input_data.model,
+                input=texts,
+                encoding_format=input_data.encoding_format,
+                dimensions=input_data.dimensions,
+            )
+
+            # Extract embeddings
+            embeddings = [data.embedding for data in response.data]
+
+            return OpenAIEmbeddingsOutput(
+                embeddings=embeddings,
+                used_model=response.model,
+                tokens_used=response.usage.total_tokens,
+            )
+        except Exception as e:
+            raise ProcessingError(
+                f"OpenAI async embeddings generation failed: {str(e)}"
+            )

{airtrain-0.1.39.dist-info → airtrain-0.1.40.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.39
+Version: 0.1.40
 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.39.dist-info → airtrain-0.1.40.dist-info}/RECORD

@@ -1,4 +1,4 @@
-airtrain/__init__.py,sha256=n0L7jCq_aew4QXp1TUM8TzyXHSQwd60t1WdnXkbTgrk,2099
+airtrain/__init__.py,sha256=fRBP9B8-SN5didcE8h1LLAPTuQTxiODu_-bbKGWLDR4,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,9 +13,10 @@ 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=-3Vz2bqAUNvVHEZxFGUv5BfzJZsO_7MRyLifbHsweEE,1488
-airtrain/integrations/anthropic/__init__.py,sha256=qwlWLDh1rEVizYFbW8430z-f1SxHio7_Gaw5cCTUtoo,274
+airtrain/integrations/__init__.py,sha256=rk9QFl0Dd7Qp4rULhi_u4smwsJwk69Kg_-fv0GQ43iw,1782
+airtrain/integrations/anthropic/__init__.py,sha256=F4kB5fuj7nYgTVcgzeHGc91LT96FZfsCJVBVCnTRh-k,541
 airtrain/integrations/anthropic/credentials.py,sha256=hlTSw9HX66kYNaeQUtn0JjdZQBMNkzzFOJOoLOOzvcY,1246
+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/aws/__init__.py,sha256=3x7v2NxpAfI-U-YgwQeH5PtsmUrNLPMfLyUGFLiBjbs,155
 airtrain/integrations/aws/credentials.py,sha256=nN-daKAl7qOb_VdRpsThG8gN5GeSUkx-ji5E_gF_vYw,1444
@@ -23,15 +24,16 @@ airtrain/integrations/aws/skills.py,sha256=TQiMXeXRRcJ14fe8Xi7Uk20iS6_INbcznuLGt
 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/fireworks/__init__.py,sha256=9pJvP0u1FJbNtB0oHa09mHVJLctELf_c27LOYyDk2ZI,271
+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/credentials.py,sha256=UpcwR9V5Hbk5sJbjFDJDbHMRqc90IQSqAvrtJCOvwEo,524
-airtrain/integrations/fireworks/models.py,sha256=F-MddbLCLAsTjwRr1l6IpJxOegyY4pD7jN9ySPiypSo,593
+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=OB4epD4CSTxExUCC1oMJ_8rHLOoftlxf0FUoIVrd4mA,5163
-airtrain/integrations/fireworks/structured_completion_skills.py,sha256=IXG4gsZDSfuscrmKIHfnyHkBaCV7zlPInaWXb95iC5k,6428
-airtrain/integrations/fireworks/structured_requests_skills.py,sha256=oRpbKMOcKgY2js16uNkIx866UEwEYSNgFPNYn9cLO3U,8409
+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/google/__init__.py,sha256=ElwgcXfbg_gGMm6zbkMXCQPFKZUb-yTJk986o19A7Cs,214
 airtrain/integrations/google/credentials.py,sha256=KSvWNqW8Mjr4MkysRvUqlrOSGdShNIe5u2OPO6vRrWY,2047
@@ -42,11 +44,11 @@ airtrain/integrations/groq/skills.py,sha256=qFyxC_2xZYnByAPo5p2aHbrqhdHYCoIdvDRA
 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=K-NY2_T1T6SEOgkpbUA55cWvK2nr2NOJgLCqmmtaCno,371
+airtrain/integrations/openai/__init__.py,sha256=w5V7lxvrKtrrjyqGoppEKg9ORKKQ2cxaLOpgCZdm_H8,541
 airtrain/integrations/openai/chinese_assistant.py,sha256=MMhv4NBOoEQ0O22ZZtP255rd5ajHC9l6FPWIjpqxBOA,1581
 airtrain/integrations/openai/credentials.py,sha256=NfRyp1QgEtgm8cxt2-BOLq-6d0X-Pcm80NnfHM8p0FY,1470
-airtrain/integrations/openai/models_config.py,sha256=bzosqqpDy2AJxu2vGdk2H4voqEGlv7LORR6fpJLhNic,3962
-airtrain/integrations/openai/skills.py,sha256=gikb9RBH1ggSSUwDDE7t6cg3LZrCAPooXg04MgqAJ-0,8862
+airtrain/integrations/openai/models_config.py,sha256=W9mu_z9tCC4ZUKHSJ6Hk4X09TRZLqEhT7TtRY5JEk5g,8007
+airtrain/integrations/openai/skills.py,sha256=1dvRJYrnU2hOmGRlkHBtyR6P8D7aIwHZfUKmjlReWrQ,12821
 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
@@ -63,8 +65,8 @@ airtrain/integrations/together/rerank_skill.py,sha256=gjH24hLWCweWKPyyfKZMG3K_g9
 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.39.dist-info/METADATA,sha256=b0g4vq6XCWe55sSnRlt6bcsK_ynvyj7SkDYTxEvrQU8,5375
-airtrain-0.1.39.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-airtrain-0.1.39.dist-info/entry_points.txt,sha256=rrJ36IUsyq6n1dSfTWXqVAgpQLPRWDfCqwd6_3B-G0U,52
-airtrain-0.1.39.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
-airtrain-0.1.39.dist-info/RECORD,,
+airtrain-0.1.40.dist-info/METADATA,sha256=3-KnsKsMdriztxTjdEWoJF599N3oicxpw2x17-mrQaw,5375
+airtrain-0.1.40.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
+airtrain-0.1.40.dist-info/entry_points.txt,sha256=rrJ36IUsyq6n1dSfTWXqVAgpQLPRWDfCqwd6_3B-G0U,52
+airtrain-0.1.40.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.40.dist-info/RECORD,,

{airtrain-0.1.39.dist-info → airtrain-0.1.40.dist-info}/WHEEL

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