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

airtrain/core/schemas.py

@@ -0,0 +1,237 @@
+from typing import Any, Dict, Optional, Type, Union, cast, get_args, get_origin
+from pydantic import BaseModel, ValidationError, create_model
+import json
+from uuid import UUID, uuid4
+
+
+class AirtrainSchema(BaseModel):
+    """Base schema class for all Airtrain schemas"""
+
+    _schema_id: Optional[UUID] = None
+    _schema_version: str = "1.0.0"
+
+    @classmethod
+    def _extract_field_type(cls, field_props: Dict) -> Type:
+        """
+        Extract Python type from field properties
+
+        Args:
+            field_props: Field properties from JSON schema
+
+        Returns:
+            Python type for the field
+        """
+        # Handle direct type specification
+        if "type" in field_props:
+            return cls._map_json_type_to_python(field_props["type"])
+
+        # Handle anyOf/oneOf cases
+        for union_key in ["anyOf", "oneOf"]:
+            if union_key in field_props:
+                types = []
+                for type_option in field_props[union_key]:
+                    if "type" in type_option:
+                        if type_option["type"] == "null":
+                            types.append(type(None))
+                        else:
+                            types.append(
+                                cls._map_json_type_to_python(type_option["type"])
+                            )
+
+                # If we have types, create a Union
+                if types:
+                    return Union[tuple(types)] if len(types) > 1 else types[0]
+
+        # Default to Any if type cannot be determined
+        return Any
+
+    @classmethod
+    def _get_field_config(cls, field_props: Dict) -> tuple:
+        """
+        Get field type and default value configuration
+
+        Args:
+            field_props: Field properties from JSON schema
+
+        Returns:
+            Tuple of (field_type, field_default)
+        """
+        field_type = cls._extract_field_type(field_props)
+
+        # Handle default values
+        if "default" in field_props:
+            return (field_type, field_props["default"])
+
+        # Handle Optional/Union types
+        if get_origin(field_type) is Union and type(None) in get_args(field_type):
+            return (field_type, None)
+
+        # No default value
+        return (field_type, ...)
+
+    @classmethod
+    def from_json_schema(cls, json_schema: str | Dict) -> "AirtrainSchema":
+        """
+        Create an AirtrainSchema from a JSON schema
+
+        Args:
+            json_schema: JSON schema string or dictionary
+
+        Returns:
+            AirtrainSchema instance
+
+        Raises:
+            ValidationError: If schema is invalid
+        """
+        if isinstance(json_schema, str):
+            json_schema = json.loads(json_schema)
+
+        # Convert JSON schema to Pydantic model
+        assert isinstance(json_schema, dict)
+        model_fields = {}
+        required_fields = json_schema.get("required", [])
+
+        for field_name, field_props in json_schema["properties"].items():
+            field_type, field_default = cls._get_field_config(field_props)
+
+            # Override default for required fields
+            if field_name in required_fields:
+                field_default = ...
+
+            model_fields[field_name] = (field_type, field_default)
+
+        # Create dynamic model using create_model
+        DynamicSchema = create_model("DynamicSchema", __base__=cls, **model_fields)
+
+        return cast(AirtrainSchema, DynamicSchema)
+
+    @classmethod
+    def from_pydantic_schema(cls, pydantic_schema: Type[BaseModel]) -> "AirtrainSchema":
+        """
+        Create an AirtrainSchema from a Pydantic model
+
+        Args:
+            pydantic_schema: Pydantic model class
+
+        Returns:
+            AirtrainSchema instance
+        """
+        # Get JSON schema from pydantic model
+        schema = pydantic_schema.model_json_schema()
+
+        # Create new schema using from_json_schema
+        return cls.from_json_schema(schema)
+
+    @staticmethod
+    def _map_json_type_to_python(json_type: str) -> Type:
+        """Map JSON schema types to Python types"""
+        type_mapping = {
+            "string": str,
+            "integer": int,
+            "number": float,
+            "boolean": bool,
+            "array": list,
+            "object": dict,
+        }
+        assert json_type in type_mapping, f"Unsupported JSON type: {json_type}"
+        return type_mapping[json_type]
+
+    def validate_custom(self) -> None:
+        """
+        Perform custom validation beyond Pydantic's built-in validation
+        To be implemented by subclasses
+
+        Raises:
+            ValidationError: If custom validation fails
+        """
+        pass
+
+    def validate_all(self) -> None:
+        """
+        Perform all validations including Pydantic and custom
+
+        Raises:
+            ValidationError: If any validation fails
+        """
+        # Pydantic validation happens automatically
+        try:
+            self.validate_custom()
+        except Exception as e:
+            raise ValidationError(f"Custom validation failed: {str(e)}")
+
+    def publish(self) -> UUID:
+        """
+        Publish schema to make it available for use
+
+        Returns:
+            UUID: Unique identifier for the published schema
+        """
+        if not self._schema_id:
+            self._schema_id = uuid4()
+        # TODO: Implement actual publishing logic
+        return self._schema_id
+
+    @classmethod
+    def get_by_id(cls, schema_id: UUID) -> "AirtrainSchema":
+        """
+        Retrieve a published schema by ID
+
+        Args:
+            schema_id: UUID of the published schema
+
+        Returns:
+            AirtrainSchema instance
+
+        Raises:
+            ValueError: If schema not found
+        """
+        # TODO: Implement schema retrieval logic
+        raise NotImplementedError("Schema retrieval not implemented yet")
+
+
+class InputSchema(AirtrainSchema):
+    """Schema for task/skill inputs"""
+
+    def validate_input_specific(self) -> None:
+        """
+        Perform input-specific validations
+        To be implemented by subclasses
+
+        Raises:
+            ValidationError: If validation fails
+        """
+        pass
+
+    def validate_custom(self) -> None:
+        """
+        Override custom validation to include input-specific validation
+
+        Raises:
+            ValidationError: If validation fails
+        """
+        super().validate_custom()
+        self.validate_input_specific()
+
+
+class OutputSchema(AirtrainSchema):
+    """Schema for task/skill outputs"""
+
+    def validate_output_specific(self) -> None:
+        """
+        Perform output-specific validations
+        To be implemented by subclasses
+
+        Raises:
+            ValidationError: If validation fails
+        """
+        pass
+
+    def validate_custom(self) -> None:
+        """
+        Override custom validation to include output-specific validation
+
+        Raises:
+            ValidationError: If validation fails
+        """
+        super().validate_custom()
+        self.validate_output_specific()

airtrain/core/skills.py

@@ -0,0 +1,269 @@
+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)
+
+
+class Skill(ABC, Generic[InputT, OutputT]):
+    """
+    Abstract base class for all skills in Airtrain.
+    Each skill must define input/output schemas and implement core processing logic.
+    """
+
+    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:
+        """
+        Process the input and generate output according to defined schemas.
+
+        Args:
+            input_data: Validated input conforming to input_schema
+
+        Returns:
+            Output conforming to output_schema
+
+        Raises:
+            ProcessingError: If processing fails
+        """
+        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.
+
+        Args:
+            input_data: Raw input data
+
+        Raises:
+            InputValidationError: If validation fails
+        """
+        if not isinstance(input_data, self.input_schema):
+            raise InputValidationError(
+                f"Input must be an instance of {self.input_schema.__name__}"
+            )
+        input_data.validate_all()
+
+    def validate_output(self, output_data: Any) -> None:
+        """
+        Validate output data after processing.
+
+        Args:
+            output_data: Processed output data
+
+        Raises:
+            OutputValidationError: If validation fails
+        """
+        if not isinstance(output_data, self.output_schema):
+            raise OutputValidationError(
+                f"Output must be an instance of {self.output_schema.__name__}"
+            )
+        output_data.validate_all()
+
+    def evaluate(self, test_dataset: Optional["Dataset"] = None) -> "EvaluationResult":
+        """
+        Evaluate skill performance.
+
+        Args:
+            test_dataset: Optional dataset for evaluation
+
+        Returns:
+            EvaluationResult containing metrics
+        """
+        if not test_dataset:
+            test_dataset = self.get_default_test_dataset()
+
+        results = []
+        for test_case in test_dataset:
+            try:
+                output = self.process(test_case.input)
+                results.append(self.compare_output(output, test_case.expected))
+            except Exception as e:
+                results.append(EvaluationError(str(e)))
+
+        return EvaluationResult(results)
+
+    def get_default_test_dataset(self) -> "Dataset":
+        """Get default test dataset for evaluation"""
+        raise NotImplementedError("No default test dataset provided")
+
+    def compare_output(self, actual: OutputT, expected: OutputT) -> Dict:
+        """
+        Compare actual output with expected output
+
+        Args:
+            actual: Actual output from processing
+            expected: Expected output from test case
+
+        Returns:
+            Dictionary containing comparison metrics
+        """
+        raise NotImplementedError("Output comparison not implemented")
+
+    @property
+    def skill_id(self) -> UUID:
+        """Unique identifier for the skill"""
+        if not self._skill_id:
+            self._skill_id = uuid4()
+        return self._skill_id
+
+
+class ProcessingError(Exception):
+    """Raised when skill processing fails"""
+
+    pass
+
+
+class InputValidationError(Exception):
+    """Raised when input validation fails"""
+
+    pass
+
+
+class OutputValidationError(Exception):
+    """Raised when output validation fails"""
+
+    pass
+
+
+class EvaluationError:
+    """Represents an error during evaluation"""
+
+    def __init__(self, message: str):
+        self.message = message
+
+
+class EvaluationResult:
+    """Contains results from skill evaluation"""
+
+    def __init__(self, results: list):
+        self.results = results
+
+    def get_metrics(self) -> Dict:
+        """Calculate evaluation metrics"""
+        return {
+            "total_cases": len(self.results),
+            "successful": len(
+                [r for r in self.results if not isinstance(r, EvaluationError)]
+            ),
+            "failed": len([r for r in self.results if isinstance(r, EvaluationError)]),
+            "results": self.results,
+        }
+
+
+class Dataset:
+    """Represents a test dataset for skill evaluation"""
+
+    def __init__(self, test_cases: list):
+        self.test_cases = test_cases
+
+    def __iter__(self):
+        return iter(self.test_cases)

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