airtrain 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

airtrain/__init__.py

@@ -0,0 +1,9 @@
+"""Airtrain - A platform for building and deploying AI agents with structured skills"""
+
+__version__ = "0.1.1"
+
+from .core.skills import Skill
+from .core.schemas import InputSchema, OutputSchema
+from .core.credentials import BaseCredentials
+
+__all__ = ["Skill", "InputSchema", "OutputSchema", "BaseCredentials"]

airtrain/core/__init__.py

@@ -0,0 +1,7 @@
+"""Core modules for Airtrain"""
+
+from .skills import Skill, ProcessingError
+from .schemas import InputSchema, OutputSchema
+from .credentials import BaseCredentials
+
+__all__ = ["Skill", "ProcessingError", "InputSchema", "OutputSchema", "BaseCredentials"]

airtrain/core/credentials.py

@@ -0,0 +1,153 @@
+from typing import Dict, List, Optional, Set
+import os
+import json
+from pathlib import Path
+from abc import ABC, abstractmethod
+import dotenv
+from pydantic import BaseModel, Field, SecretStr
+import yaml  # type: ignore
+
+
+class CredentialError(Exception):
+    """Base exception for credential-related errors"""
+
+    pass
+
+
+class CredentialNotFoundError(CredentialError):
+    """Raised when a required credential is not found"""
+
+    pass
+
+
+class CredentialValidationError(CredentialError):
+    """Raised when credentials fail validation"""
+
+    pass
+
+
+class BaseCredentials(BaseModel):
+    """Base class for all credential configurations"""
+
+    _loaded: bool = False
+    _required_credentials: Set[str] = set()
+
+    def load_to_env(self) -> None:
+        """Load credentials into environment variables"""
+        for field_name, field_value in self:
+            if isinstance(field_value, SecretStr):
+                value = field_value.get_secret_value()
+            else:
+                value = str(field_value)
+            os.environ[field_name.upper()] = value
+        self._loaded = True
+
+    @classmethod
+    def from_env(cls) -> "BaseCredentials":
+        """Create credentials instance from environment variables"""
+        field_values = {}
+        for field_name in cls.model_fields:
+            env_key = field_name.upper()
+            if env_value := os.getenv(env_key):
+                field_values[field_name] = env_value
+        return cls(**field_values)
+
+    @classmethod
+    def from_file(cls, file_path: Path) -> "BaseCredentials":
+        """Load credentials from a file (supports .env, .json, .yaml)"""
+        if not file_path.exists():
+            raise FileNotFoundError(f"Credentials file not found: {file_path}")
+
+        if file_path.suffix == ".env":
+            dotenv.load_dotenv(file_path)
+            return cls.from_env()
+
+        elif file_path.suffix == ".json":
+            with open(file_path) as f:
+                data = json.load(f)
+            return cls(**data)
+
+        elif file_path.suffix in {".yaml", ".yml"}:
+            with open(file_path) as f:
+                data = yaml.safe_load(f)
+            return cls(**data)
+
+        else:
+            raise ValueError(f"Unsupported file format: {file_path.suffix}")
+
+    def save_to_file(self, file_path: Path) -> None:
+        """Save credentials to a file"""
+        data = self.model_dump(exclude={"_loaded"})
+
+        # Convert SecretStr to plain strings for saving
+        for key, value in data.items():
+            if isinstance(value, SecretStr):
+                data[key] = value.get_secret_value()
+
+        if file_path.suffix == ".env":
+            with open(file_path, "w") as f:
+                for key, value in data.items():
+                    f.write(f"{key.upper()}={value}\n")
+
+        elif file_path.suffix == ".json":
+            with open(file_path, "w") as f:
+                json.dump(data, f, indent=2)
+
+        elif file_path.suffix in {".yaml", ".yml"}:
+            with open(file_path, "w") as f:
+                yaml.dump(data, f)
+
+        else:
+            raise ValueError(f"Unsupported file format: {file_path.suffix}")
+
+    def validate_credentials(self) -> None:
+        """Validate that all required credentials are present"""
+        missing = []
+        for field_name in self._required_credentials:
+            value = getattr(self, field_name, None)
+            if value is None or (
+                isinstance(value, SecretStr) and not value.get_secret_value()
+            ):
+                missing.append(field_name)
+
+        if missing:
+            raise CredentialValidationError(
+                f"Missing required credentials: {', '.join(missing)}"
+            )
+
+    def clear_from_env(self) -> None:
+        """Remove credentials from environment variables"""
+        for field_name in self.model_fields:
+            env_key = field_name.upper()
+            if env_key in os.environ:
+                del os.environ[env_key]
+        self._loaded = False
+
+
+class OpenAICredentials(BaseCredentials):
+    """OpenAI API credentials"""
+
+    api_key: SecretStr = Field(..., description="OpenAI API key")
+    organization_id: Optional[str] = Field(None, description="OpenAI organization ID")
+
+    _required_credentials = {"api_key"}
+
+
+class AWSCredentials(BaseCredentials):
+    """AWS credentials"""
+
+    aws_access_key_id: SecretStr
+    aws_secret_access_key: SecretStr
+    aws_region: str = "us-east-1"
+    aws_session_token: Optional[SecretStr] = None
+
+    _required_credentials = {"aws_access_key_id", "aws_secret_access_key"}
+
+
+class GoogleCloudCredentials(BaseCredentials):
+    """Google Cloud credentials"""
+
+    project_id: str
+    service_account_key: SecretStr
+
+    _required_credentials = {"project_id", "service_account_key"}

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,167 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, Type, Generic, TypeVar
+from uuid import UUID, uuid4
+from .schemas import InputSchema, OutputSchema
+
+# 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
+
+    @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 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-0.1.0.dist-info → airtrain-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.0
+Version: 0.1.1
 Summary: A platform for building and deploying AI agents with structured skills
 Home-page: https://github.com/rosaboyle/airtrain.dev
 Author: Dheeraj Pai
@@ -95,7 +95,7 @@ creds.load_to_env()
 
 ## Documentation
 
-For detailed documentation, visit [our documentation site](https://airtrain.readthedocs.io/).
+For detailed documentation, visit [our documentation site](https://docs.airtrain.dev/).
 
 ## Contributing
 

airtrain-0.1.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+airtrain/__init__.py,sha256=dqQKBcKKk6Xis8BNi-BygiK1W51cppG3Sh5rdefqBys,312
+airtrain/core/__init__.py,sha256=9h7iKwTzZocCPc9bU6j8bA02BokteWIOcO1uaqGMcrk,254
+airtrain/core/credentials.py,sha256=CzUZkAFxrSMC0nq70zybkkJmeIZDYiNBuzfivOTEgH0,4773
+airtrain/core/schemas.py,sha256=MMXrDviC4gRea_QaPpbjgO--B_UKxnD7YrxqZOLJZZU,7003
+airtrain/core/skills.py,sha256=LljalzeSHK5eQPTAOEAYc5D8Qn1kVSfiz9WgziTD5UM,4688
+airtrain-0.1.1.dist-info/METADATA,sha256=3qlD2n866n3emRI9yEnfYC4dzLhwXGyeb6lqqFnwCyM,2786
+airtrain-0.1.1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+airtrain-0.1.1.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.1.dist-info/RECORD,,

airtrain-0.1.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+airtrain

airtrain-0.1.0.dist-info/RECORD

@@ -1,4 +0,0 @@
-airtrain-0.1.0.dist-info/METADATA,sha256=2wLV8r1XTMiq6YUJpSCANE5A2KbXlfDUjSAxHOwq8iw,2792
-airtrain-0.1.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-airtrain-0.1.0.dist-info/top_level.txt,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-airtrain-0.1.0.dist-info/RECORD,,

airtrain-0.1.0.dist-info/top_level.txt

@@ -1 +0,0 @@
-