airtrain 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

airtrain/integrations/anthropic/list_models.py

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

airtrain/integrations/anthropic/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/anthropic/skills.py

@@ -0,0 +1,155 @@
+from typing import List, Optional, Dict, Any, Generator
+from pydantic import Field
+from anthropic import Anthropic
+import base64
+from pathlib import Path
+from loguru import logger
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import AnthropicCredentials
+
+
+class AnthropicInput(InputSchema):
+    """Schema for Anthropic chat input"""
+
+    user_input: str = Field(..., description="User's input text")
+    system_prompt: str = Field(
+        default="You are a helpful assistant.",
+        description="System prompt to guide the model's behavior",
+    )
+    conversation_history: List[Dict[str, str]] = Field(
+        default_factory=list,
+        description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
+    )
+    model: str = Field(
+        default="claude-3-opus-20240229", description="Anthropic model to use"
+    )
+    max_tokens: Optional[int] = Field(
+        default=131072, description="Maximum tokens in response"
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+    images: List[Path] = Field(
+        default_factory=list,
+        description="List of image paths to include in the message",
+    )
+    stream: bool = Field(
+        default=False, description="Whether to stream the response progressively"
+    )
+
+
+class AnthropicOutput(OutputSchema):
+    """Schema for Anthropic chat output"""
+
+    response: str = Field(..., description="Model's response text")
+    used_model: str = Field(..., description="Model used for generation")
+    usage: Dict[str, Any] = Field(
+        default_factory=dict, description="Usage statistics from the API"
+    )
+
+
+class AnthropicChatSkill(Skill[AnthropicInput, AnthropicOutput]):
+    """Skill for Anthropic chat"""
+
+    input_schema = AnthropicInput
+    output_schema = AnthropicOutput
+
+    def __init__(self, credentials: Optional[AnthropicCredentials] = None):
+        super().__init__()
+        self.credentials = credentials or AnthropicCredentials.from_env()
+        self.client = Anthropic(
+            api_key=self.credentials.anthropic_api_key.get_secret_value()
+        )
+
+    def _build_messages(self, input_data: AnthropicInput) -> List[Dict[str, Any]]:
+        """
+        Build messages list from input data including conversation history.
+
+        Args:
+            input_data: The input data containing system prompt, conversation history, and user input
+
+        Returns:
+            List[Dict[str, Any]]: List of messages in the format required by Anthropic
+        """
+        messages = []
+
+        # Add conversation history if present
+        if input_data.conversation_history:
+            messages.extend(input_data.conversation_history)
+
+        # Prepare user message content
+        user_message = {"type": "text", "text": input_data.user_input}
+
+        # Add images if present
+        if input_data.images:
+            content = []
+            for image_path in input_data.images:
+                with open(image_path, "rb") as img_file:
+                    base64_image = base64.b64encode(img_file.read()).decode("utf-8")
+                content.append(
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": "image/jpeg",
+                            "data": base64_image,
+                        },
+                    }
+                )
+            content.append(user_message)
+            messages.append({"role": "user", "content": content})
+        else:
+            messages.append({"role": "user", "content": [user_message]})
+
+        return messages
+
+    def process_stream(self, input_data: AnthropicInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
+        try:
+            messages = self._build_messages(input_data)
+
+            with self.client.beta.messages.stream(
+                model=input_data.model,
+                system=input_data.system_prompt,
+                messages=messages,
+                max_tokens=input_data.max_tokens,
+                temperature=input_data.temperature,
+            ) as stream:
+                for chunk in stream.text_stream:
+                    yield chunk
+
+        except Exception as e:
+            logger.exception(f"Anthropic streaming failed: {str(e)}")
+            raise ProcessingError(f"Anthropic streaming failed: {str(e)}")
+
+    def process(self, input_data: AnthropicInput) -> AnthropicOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+                usage = {}  # Usage stats not available in streaming
+            else:
+                messages = self._build_messages(input_data)
+                response = self.client.messages.create(
+                    model=input_data.model,
+                    system=input_data.system_prompt,
+                    messages=messages,
+                    max_tokens=input_data.max_tokens,
+                    temperature=input_data.temperature,
+                )
+                usage = response.usage.model_dump() if response.usage else {}
+
+            return AnthropicOutput(
+                response=response.content[0].text,
+                used_model=input_data.model,
+                usage=usage,
+            )
+
+        except Exception as e:
+            logger.exception(f"Anthropic processing failed: {str(e)}")
+            raise ProcessingError(f"Anthropic processing failed: {str(e)}")

airtrain/integrations/aws/__init__.py

@@ -0,0 +1,6 @@
+"""AWS integration module"""
+
+from .credentials import AWSCredentials
+from .skills import AWSBedrockSkill
+
+__all__ = ["AWSCredentials", "AWSBedrockSkill"]

airtrain/integrations/aws/credentials.py

@@ -0,0 +1,36 @@
+from typing import Optional
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+import boto3
+
+
+class AWSCredentials(BaseCredentials):
+    """AWS credentials"""
+
+    aws_access_key_id: SecretStr = Field(..., description="AWS Access Key ID")
+    aws_secret_access_key: SecretStr = Field(..., description="AWS Secret Access Key")
+    aws_region: str = Field(default="us-east-1", description="AWS Region")
+    aws_session_token: Optional[SecretStr] = Field(
+        None, description="AWS Session Token"
+    )
+
+    _required_credentials = {"aws_access_key_id", "aws_secret_access_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate AWS credentials by making a test API call"""
+        try:
+            session = boto3.Session(
+                aws_access_key_id=self.aws_access_key_id.get_secret_value(),
+                aws_secret_access_key=self.aws_secret_access_key.get_secret_value(),
+                aws_session_token=(
+                    self.aws_session_token.get_secret_value()
+                    if self.aws_session_token
+                    else None
+                ),
+                region_name=self.aws_region,
+            )
+            sts = session.client("sts")
+            sts.get_caller_identity()
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid AWS credentials: {str(e)}")

airtrain/integrations/aws/skills.py

@@ -0,0 +1,98 @@
+from typing import List, Optional, Dict, Any
+from pydantic import Field
+import boto3
+from pathlib import Path
+from loguru import logger
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import AWSCredentials
+
+
+class AWSBedrockInput(InputSchema):
+    """Schema for AWS Bedrock chat input"""
+
+    user_input: str = Field(..., description="User's input text")
+    system_prompt: str = Field(
+        default="You are a helpful assistant.",
+        description="System prompt to guide the model's behavior",
+    )
+    model: str = Field(
+        default="anthropic.claude-3-sonnet-20240229-v1:0",
+        description="AWS Bedrock model to use",
+    )
+    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
+    )
+    images: Optional[List[Path]] = Field(
+        default=None,
+        description="Optional list of image paths to include in the message",
+    )
+
+
+class AWSBedrockOutput(OutputSchema):
+    """Schema for AWS Bedrock chat output"""
+
+    response: str = Field(..., description="Model's response text")
+    used_model: str = Field(..., description="Model used for generation")
+    usage: Dict[str, Any] = Field(
+        default_factory=dict, description="Usage statistics from the API"
+    )
+
+
+class AWSBedrockSkill(Skill[AWSBedrockInput, AWSBedrockOutput]):
+    """Skill for interacting with AWS Bedrock models"""
+
+    input_schema = AWSBedrockInput
+    output_schema = AWSBedrockOutput
+
+    def __init__(self, credentials: Optional[AWSCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or AWSCredentials.from_env()
+        self.client = boto3.client(
+            "bedrock-runtime",
+            aws_access_key_id=self.credentials.aws_access_key_id.get_secret_value(),
+            aws_secret_access_key=self.credentials.aws_secret_access_key.get_secret_value(),
+            region_name=self.credentials.aws_region,
+        )
+
+    def process(self, input_data: AWSBedrockInput) -> AWSBedrockOutput:
+        """Process the input using AWS Bedrock API"""
+        try:
+            logger.info(f"Processing request with model {input_data.model}")
+
+            # Prepare request body based on model provider
+            if "anthropic" in input_data.model:
+                request_body = {
+                    "anthropic_version": "bedrock-2023-05-31",
+                    "max_tokens": input_data.max_tokens,
+                    "temperature": input_data.temperature,
+                    "system": input_data.system_prompt,
+                    "messages": [{"role": "user", "content": input_data.user_input}],
+                }
+            else:
+                raise ProcessingError(f"Unsupported model: {input_data.model}")
+
+            response = self.client.invoke_model(
+                modelId=input_data.model, body=request_body
+            )
+
+            # Parse response based on model provider
+            if "anthropic" in input_data.model:
+                response_data = response["body"]["completion"]
+                usage = {
+                    "input_tokens": response["body"]["usage"]["input_tokens"],
+                    "output_tokens": response["body"]["usage"]["output_tokens"],
+                }
+            else:
+                raise ProcessingError(f"Unsupported model response: {input_data.model}")
+
+            return AWSBedrockOutput(
+                response=response_data, used_model=input_data.model, usage=usage
+            )
+
+        except Exception as e:
+            logger.exception(f"AWS Bedrock processing failed: {str(e)}")
+            raise ProcessingError(f"AWS Bedrock processing failed: {str(e)}")

airtrain/integrations/cerebras/__init__.py

@@ -0,0 +1,6 @@
+"""Cerebras integration module"""
+
+from .credentials import CerebrasCredentials
+from .skills import CerebrasChatSkill
+
+__all__ = ["CerebrasCredentials", "CerebrasChatSkill"]

airtrain/integrations/cerebras/credentials.py

@@ -0,0 +1,19 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+
+
+class CerebrasCredentials(BaseCredentials):
+    """Cerebras credentials"""
+
+    cerebras_api_key: SecretStr = Field(..., description="Cerebras API key")
+
+    _required_credentials = {"cerebras_api_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Cerebras credentials"""
+        try:
+            # Implement Cerebras-specific validation
+            # This would depend on their API client implementation
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid Cerebras credentials: {str(e)}")

airtrain/integrations/cerebras/skills.py

@@ -0,0 +1,127 @@
+from typing import List, Optional, Dict, Any, Generator
+from pydantic import Field
+from cerebras.cloud.sdk import Cerebras
+from loguru import logger
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import CerebrasCredentials
+
+
+class CerebrasInput(InputSchema):
+    """Schema for Cerebras chat input"""
+
+    user_input: str = Field(..., description="User's input text")
+    system_prompt: str = Field(
+        default="You are a helpful assistant.",
+        description="System prompt to guide the model's behavior",
+    )
+    conversation_history: List[Dict[str, str]] = Field(
+        default_factory=list,
+        description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
+    )
+    model: str = Field(default="llama3.1-8b", description="Cerebras model to use")
+    max_tokens: Optional[int] = Field(
+        default=131072, description="Maximum tokens in response"
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+    stream: bool = Field(
+        default=False, description="Whether to stream the response progressively"
+    )
+
+
+class CerebrasOutput(OutputSchema):
+    """Schema for Cerebras chat output"""
+
+    response: str = Field(..., description="Model's response text")
+    used_model: str = Field(..., description="Model used for generation")
+    usage: Dict[str, Any] = Field(default_factory=dict, description="Usage statistics")
+
+
+class CerebrasChatSkill(Skill[CerebrasInput, CerebrasOutput]):
+    """Skill for Cerebras chat"""
+
+    input_schema = CerebrasInput
+    output_schema = CerebrasOutput
+
+    def __init__(self, credentials: Optional[CerebrasCredentials] = None):
+        super().__init__()
+        self.credentials = credentials or CerebrasCredentials.from_env()
+        self.client = Cerebras(
+            api_key=self.credentials.cerebras_api_key.get_secret_value()
+        )
+
+    def _build_messages(self, input_data: CerebrasInput) -> List[Dict[str, str]]:
+        """
+        Build messages list from input data including conversation history.
+
+        Args:
+            input_data: The input data containing system prompt, conversation history, and user input
+
+        Returns:
+            List[Dict[str, str]]: List of messages in the format required by Cerebras
+        """
+        messages = [{"role": "system", "content": input_data.system_prompt}]
+
+        # Add conversation history if present
+        if input_data.conversation_history:
+            messages.extend(input_data.conversation_history)
+
+        # Add current user input
+        messages.append({"role": "user", "content": input_data.user_input})
+
+        return messages
+
+    def process_stream(self, input_data: CerebrasInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
+        try:
+            messages = self._build_messages(input_data)
+
+            stream = self.client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+                stream=True,
+            )
+
+            for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
+
+        except Exception as e:
+            logger.exception(f"Cerebras streaming failed: {str(e)}")
+            raise ProcessingError(f"Cerebras streaming failed: {str(e)}")
+
+    def process(self, input_data: CerebrasInput) -> CerebrasOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+                usage = {}  # Usage stats not available in streaming
+            else:
+                messages = self._build_messages(input_data)
+                response = self.client.chat.completions.create(
+                    model=input_data.model,
+                    messages=messages,
+                    temperature=input_data.temperature,
+                    max_tokens=input_data.max_tokens,
+                )
+                usage = (
+                    response.usage.model_dump() if hasattr(response, "usage") else {}
+                )
+
+            return CerebrasOutput(
+                response=response.choices[0].message.content,
+                used_model=input_data.model,
+                usage=usage,
+            )
+
+        except Exception as e:
+            logger.exception(f"Cerebras processing failed: {str(e)}")
+            raise ProcessingError(f"Cerebras processing failed: {str(e)}")

airtrain/integrations/combined/__init__.py

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