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

airtrain/core/skills.py

@@ -1,8 +1,17 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Optional, Type, Generic, TypeVar
 from uuid import UUID, uuid4
+import time
+import functools
 from .schemas import InputSchema, OutputSchema
 
+# Import telemetry
+from airtrain.telemetry import (
+    telemetry,
+    SkillInitTelemetryEvent,
+    SkillProcessTelemetryEvent,
+)
+
 # Generic type variables for input and output schemas
 InputT = TypeVar("InputT", bound=InputSchema)
 OutputT = TypeVar("OutputT", bound=OutputSchema)
@@ -17,6 +26,92 @@ class Skill(ABC, Generic[InputT, OutputT]):
     input_schema: Type[InputT]
     output_schema: Type[OutputT]
     _skill_id: Optional[UUID] = None
+    _original_process = None
+
+    def __init__(self):
+        """Initialize the skill and capture telemetry."""
+        # Initialize skill_id if not already set
+        if not self._skill_id:
+            self._skill_id = uuid4()
+        
+        # Monkey patch the process method if it hasn't been patched yet
+        # This allows us to add telemetry without changing the API
+        if not hasattr(self.__class__, '_patched_process'):
+            # Store the original process method implementation from this instance
+            # This is crucial for proper behavior with inheritance
+            self.__class__._original_process = self.__class__.process
+            
+            # Create a wrapper function that will capture telemetry
+            def _create_wrapper(original_method):
+                @functools.wraps(original_method)
+                def wrapped_process(instance, input_data):
+                    start_time = time.time()
+                    error = None
+                    
+                    try:
+                        # Call the original process method
+                        result = original_method(instance, input_data)
+                        return result
+                    except Exception as e:
+                        error = str(e)
+                        raise
+                    finally:
+                        duration = time.time() - start_time
+                        
+                        try:
+                            # Serialize input data for telemetry
+                            serialized_input = None
+                            try:
+                                # Convert input_data to dict if it's a Pydantic model
+                                if hasattr(input_data, "dict"):
+                                    serialized_input = input_data.dict()
+                                # If it's a dataclass
+                                elif hasattr(input_data, "__dataclass_fields__"):
+                                    from dataclasses import asdict
+                                    serialized_input = asdict(input_data)
+                                # Fallback
+                                else:
+                                    serialized_input = {
+                                        "__str__": str(input_data)
+                                    }
+                            except Exception:
+                                # If serialization fails, provide simple info
+                                serialized_input = {"error": "Failed to serialize input data"}
+                            
+                            telemetry.capture(
+                                SkillProcessTelemetryEvent(
+                                    skill_id=str(instance.skill_id),
+                                    skill_class=instance.__class__.__name__,
+                                    input_schema=instance.input_schema.__name__,
+                                    output_schema=instance.output_schema.__name__,
+                                    input_data=serialized_input,
+                                    duration_seconds=duration,
+                                    error=error,
+                                )
+                            )
+                        except Exception:
+                            # Silently continue if telemetry fails
+                            pass
+                
+                return wrapped_process
+            
+            # Replace the process method with our wrapped version at the class level
+            self.__class__.process = _create_wrapper(self.__class__._original_process)
+            
+            # Mark this class as patched to prevent double-patching
+            self.__class__._patched_process = True
+        
+        # Capture telemetry for initialization
+        try:
+            telemetry.capture(
+                SkillInitTelemetryEvent(
+                    skill_id=str(self.skill_id),
+                    skill_class=self.__class__.__name__,
+                )
+            )
+        except Exception:
+            # Silently continue if telemetry fails
+            pass
 
     @abstractmethod
     def process(self, input_data: InputT) -> OutputT:
@@ -34,6 +129,13 @@ class Skill(ABC, Generic[InputT, OutputT]):
         """
         pass
 
+    def __call__(self, input_data: InputT) -> OutputT:
+        """Make the skill callable, with input/output validation."""
+        self.validate_input(input_data)
+        result = self.process(input_data)
+        self.validate_output(result)
+        return result
+
     def validate_input(self, input_data: Any) -> None:
         """
         Validate input data before processing.

airtrain/integrations/__init__.py

@@ -0,0 +1,74 @@
+"""Airtrain integrations package"""
+
+# Credentials imports
+from .openai.credentials import OpenAICredentials
+from .aws.credentials import AWSCredentials
+from .google.credentials import GoogleCloudCredentials
+from .anthropic.credentials import AnthropicCredentials
+from .groq.credentials import GroqCredentials
+from .together.credentials import TogetherAICredentials
+from .ollama.credentials import OllamaCredentials
+from .sambanova.credentials import SambanovaCredentials
+from .cerebras.credentials import CerebrasCredentials
+from .perplexity.credentials import PerplexityCredentials
+
+# Skills imports
+from .openai.skills import OpenAIChatSkill, OpenAIParserSkill
+from .anthropic.skills import AnthropicChatSkill
+from .aws.skills import AWSBedrockSkill
+from .google.skills import GoogleChatSkill
+from .groq.skills import GroqChatSkill
+from .together.skills import TogetherAIChatSkill
+from .ollama.skills import OllamaChatSkill
+from .sambanova.skills import SambanovaChatSkill
+from .cerebras.skills import CerebrasChatSkill
+from .perplexity.skills import PerplexityChatSkill, PerplexityStreamingChatSkill
+
+# Model configurations
+from .openai.models_config import OPENAI_MODELS, OpenAIModelConfig
+from .anthropic.models_config import ANTHROPIC_MODELS, AnthropicModelConfig
+from .perplexity.models_config import PERPLEXITY_MODELS_CONFIG
+
+# Combined modules
+from .combined.list_models_factory import (
+    ListModelsSkillFactory,
+    GenericListModelsInput,
+    GenericListModelsOutput,
+)
+
+__all__ = [
+    # Credentials
+    "OpenAICredentials",
+    "AWSCredentials",
+    "GoogleCloudCredentials",
+    "AnthropicCredentials",
+    "GroqCredentials",
+    "TogetherAICredentials",
+    "OllamaCredentials",
+    "SambanovaCredentials",
+    "CerebrasCredentials",
+    "PerplexityCredentials",
+    # Skills
+    "OpenAIChatSkill",
+    "OpenAIParserSkill",
+    "AnthropicChatSkill",
+    "AWSBedrockSkill",
+    "GoogleChatSkill",
+    "GroqChatSkill",
+    "TogetherAIChatSkill",
+    "OllamaChatSkill",
+    "SambanovaChatSkill",
+    "CerebrasChatSkill",
+    "PerplexityChatSkill",
+    "PerplexityStreamingChatSkill",
+    # Model configurations
+    "OPENAI_MODELS",
+    "OpenAIModelConfig",
+    "ANTHROPIC_MODELS",
+    "AnthropicModelConfig",
+    "PERPLEXITY_MODELS_CONFIG",
+    # Combined modules
+    "ListModelsSkillFactory",
+    "GenericListModelsInput",
+    "GenericListModelsOutput",
+]

airtrain/integrations/anthropic/__init__.py

@@ -0,0 +1,33 @@
+"""Anthropic integration for Airtrain"""
+
+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,
+)
+from .list_models import (
+    AnthropicListModelsSkill,
+    AnthropicListModelsInput,
+    AnthropicListModelsOutput,
+    AnthropicModel,
+)
+
+__all__ = [
+    "AnthropicCredentials",
+    "AnthropicChatSkill",
+    "AnthropicInput",
+    "AnthropicOutput",
+    "ANTHROPIC_MODELS",
+    "AnthropicModelConfig",
+    "get_model_config",
+    "get_default_model",
+    "calculate_cost",
+    "AnthropicListModelsSkill",
+    "AnthropicListModelsInput",
+    "AnthropicListModelsOutput",
+    "AnthropicModel",
+]

airtrain/integrations/anthropic/credentials.py

@@ -0,0 +1,32 @@
+from pydantic import Field, SecretStr, validator
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+from anthropic import Anthropic
+
+
+class AnthropicCredentials(BaseCredentials):
+    """Anthropic API credentials"""
+
+    anthropic_api_key: SecretStr = Field(..., description="Anthropic API key")
+    version: str = Field(default="2023-06-01", description="API Version")
+
+    _required_credentials = {"anthropic_api_key"}
+
+    @validator("anthropic_api_key")
+    def validate_api_key_format(cls, v: SecretStr) -> SecretStr:
+        key = v.get_secret_value()
+        if not key.startswith("sk-ant-"):
+            raise ValueError("Anthropic API key must start with 'sk-ant-'")
+        return v
+
+    async def validate_credentials(self) -> bool:
+        """Validate Anthropic credentials"""
+        try:
+            client = Anthropic(api_key=self.anthropic_api_key.get_secret_value())
+            client.messages.create(
+                model="claude-3-opus-20240229",
+                max_tokens=1,
+                messages=[{"role": "user", "content": "Hi"}],
+            )
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid Anthropic credentials: {str(e)}")

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)}")