airtrain 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

airtrain/__init__.py

@@ -1,6 +1,6 @@
 """Airtrain - A platform for building and deploying AI agents with structured skills"""
 
-__version__ = "0.1.6"
+__version__ = "0.1.8"
 
 from .core.skills import Skill
 from .core.schemas import InputSchema, OutputSchema

airtrain/integrations/__init__.py

@@ -1,5 +1,6 @@
 """Airtrain integrations package"""
 
+# Credentials imports
 from .openai.credentials import OpenAICredentials
 from .aws.credentials import AWSCredentials
 from .google.credentials import GoogleCloudCredentials
@@ -10,17 +11,37 @@ from .ollama.credentials import OllamaCredentials
 from .sambanova.credentials import SambanovaCredentials
 from .cerebras.credentials import CerebrasCredentials
 
+# Skills imports
+from .openai.skills import OpenAIChatSkill, OpenAIParserSkill
 from .anthropic.skills import AnthropicChatSkill
+from .aws.skills import AWSBedrockSkill
+from .google.skills import VertexAISkill
+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
 
 __all__ = [
+    # Credentials
     "OpenAICredentials",
     "AWSCredentials",
     "GoogleCloudCredentials",
     "AnthropicCredentials",
-    "AnthropicChatSkill",
     "GroqCredentials",
     "TogetherAICredentials",
     "OllamaCredentials",
     "SambanovaCredentials",
     "CerebrasCredentials",
+    # Skills
+    "OpenAIChatSkill",
+    "OpenAIParserSkill",
+    "AnthropicChatSkill",
+    "AWSBedrockSkill",
+    "VertexAISkill",
+    "GroqChatSkill",
+    "TogetherAIChatSkill",
+    "OllamaChatSkill",
+    "SambanovaChatSkill",
+    "CerebrasChatSkill",
 ]

airtrain/integrations/anthropic/__init__.py

@@ -0,0 +1,11 @@
+"""Anthropic integration for Airtrain"""
+
+from .credentials import AnthropicCredentials
+from .skills import AnthropicChatSkill, AnthropicInput, AnthropicOutput
+
+__all__ = [
+    "AnthropicCredentials",
+    "AnthropicChatSkill",
+    "AnthropicInput",
+    "AnthropicOutput",
+]

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

@@ -0,0 +1,135 @@
+from typing import List, Optional, Dict, Any
+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",
+    )
+    model: str = Field(
+        default="claude-3-opus-20240229", description="Anthropic model to use"
+    )
+    max_tokens: int = Field(default=1024, 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 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 interacting with Anthropic's Claude models"""
+
+    input_schema = AnthropicInput
+    output_schema = AnthropicOutput
+
+    def __init__(self, credentials: Optional[AnthropicCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or AnthropicCredentials.from_env()
+        self.client = Anthropic(
+            api_key=self.credentials.anthropic_api_key.get_secret_value()
+        )
+
+    def _encode_image(self, image_path: Path) -> Dict[str, Any]:
+        """Convert image to base64 for API consumption"""
+        try:
+            if not image_path.exists():
+                raise FileNotFoundError(f"Image file not found: {image_path}")
+
+            with open(image_path, "rb") as img_file:
+                encoded = base64.b64encode(img_file.read()).decode()
+                return {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": f"image/{image_path.suffix[1:]}",
+                        "data": encoded,
+                    },
+                }
+        except Exception as e:
+            logger.error(f"Failed to encode image {image_path}: {str(e)}")
+            raise ProcessingError(f"Image encoding failed: {str(e)}")
+
+    def process(self, input_data: AnthropicInput) -> AnthropicOutput:
+        """Process the input using Anthropic's API"""
+        try:
+            logger.info(f"Processing request with model {input_data.model}")
+
+            # Prepare message content
+            content = []
+
+            # Add text content
+            content.append({"type": "text", "text": input_data.user_input})
+
+            # Add images if provided
+            if input_data.images:
+                logger.debug(f"Processing {len(input_data.images)} images")
+                for image_path in input_data.images:
+                    content.append(self._encode_image(image_path))
+
+            # Create message
+            response = self.client.messages.create(
+                model=input_data.model,
+                max_tokens=input_data.max_tokens,
+                temperature=input_data.temperature,
+                system=input_data.system_prompt,
+                messages=[{"role": "user", "content": content}],
+            )
+
+            # Validate response content
+            if not response.content:
+                logger.error("Empty response received from Anthropic API")
+                raise ProcessingError("Empty response received from Anthropic API")
+
+            if not isinstance(response.content, list) or not response.content:
+                logger.error("Invalid response format from Anthropic API")
+                raise ProcessingError("Invalid response format from Anthropic API")
+
+            first_content = response.content[0]
+            if not hasattr(first_content, "text"):
+                logger.error("Response content does not contain text")
+                raise ProcessingError("Response content does not contain text")
+
+            logger.success("Successfully processed Anthropic request")
+
+            # Create output
+            return AnthropicOutput(
+                response=first_content.text,
+                used_model=response.model,
+                usage={
+                    "input_tokens": response.usage.input_tokens,
+                    "output_tokens": response.usage.output_tokens,
+                },
+            )
+
+        except ProcessingError:
+            # Re-raise ProcessingError without modification
+            raise
+        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=1024, 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,22 @@
+from pydantic import Field, SecretStr, HttpUrl
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+from typing import Optional
+
+
+class CerebrasCredentials(BaseCredentials):
+    """Cerebras credentials"""
+
+    api_key: SecretStr = Field(..., description="Cerebras API key")
+    endpoint_url: HttpUrl = Field(..., description="Cerebras API endpoint")
+    project_id: Optional[str] = Field(None, description="Cerebras Project ID")
+
+    _required_credentials = {"api_key", "endpoint_url"}
+
+    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,41 @@
+from typing import Optional, Dict, Any
+from pydantic import Field
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import CerebrasCredentials
+
+
+class CerebrasInput(InputSchema):
+    """Schema for Cerebras 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="cerebras-gpt", description="Cerebras model to use")
+    max_tokens: int = Field(default=1024, description="Maximum tokens in response")
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+
+
+class CerebrasOutput(OutputSchema):
+    """Schema for Cerebras 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 - Not Implemented"""
+
+    input_schema = CerebrasInput
+    output_schema = CerebrasOutput
+
+    def __init__(self, credentials: Optional[CerebrasCredentials] = None):
+        raise NotImplementedError("CerebrasChatSkill is not implemented yet")
+
+    def process(self, input_data: CerebrasInput) -> CerebrasOutput:
+        raise NotImplementedError("CerebrasChatSkill is not implemented yet")

airtrain/integrations/google/__init__.py

@@ -0,0 +1,6 @@
+"""Google Cloud integration module"""
+
+from .credentials import GoogleCloudCredentials
+from .skills import VertexAISkill
+
+__all__ = ["GoogleCloudCredentials", "VertexAISkill"]

airtrain/integrations/google/credentials.py

@@ -0,0 +1,27 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+from google.cloud import storage
+
+
+class GoogleCloudCredentials(BaseCredentials):
+    """Google Cloud credentials"""
+
+    project_id: str = Field(..., description="Google Cloud Project ID")
+    service_account_key: SecretStr = Field(..., description="Service Account Key JSON")
+
+    _required_credentials = {"project_id", "service_account_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Google Cloud credentials"""
+        try:
+            # Initialize with service account key
+            storage_client = storage.Client.from_service_account_info(
+                self.service_account_key.get_secret_value()
+            )
+            # Test API call
+            storage_client.list_buckets(max_results=1)
+            return True
+        except Exception as e:
+            raise CredentialValidationError(
+                f"Invalid Google Cloud credentials: {str(e)}"
+            )

airtrain/integrations/google/skills.py

@@ -0,0 +1,41 @@
+from typing import Optional, Dict, Any
+from pydantic import Field
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import GoogleCloudCredentials
+
+
+class VertexAIInput(InputSchema):
+    """Schema for Google Vertex AI 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="text-bison", description="Vertex AI model to use")
+    max_tokens: int = Field(default=1024, description="Maximum tokens in response")
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+
+
+class VertexAIOutput(OutputSchema):
+    """Schema for Vertex AI 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 VertexAISkill(Skill[VertexAIInput, VertexAIOutput]):
+    """Skill for Google Vertex AI - Not Implemented"""
+
+    input_schema = VertexAIInput
+    output_schema = VertexAIOutput
+
+    def __init__(self, credentials: Optional[GoogleCloudCredentials] = None):
+        raise NotImplementedError("VertexAISkill is not implemented yet")
+
+    def process(self, input_data: VertexAIInput) -> VertexAIOutput:
+        raise NotImplementedError("VertexAISkill is not implemented yet")

airtrain/integrations/groq/__init__.py

@@ -0,0 +1,6 @@
+"""Groq integration module"""
+
+from .credentials import GroqCredentials
+from .skills import GroqChatSkill
+
+__all__ = ["GroqCredentials", "GroqChatSkill"]

airtrain/integrations/groq/credentials.py

@@ -0,0 +1,24 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+from groq import Groq
+
+
+class GroqCredentials(BaseCredentials):
+    """Groq API credentials"""
+
+    api_key: SecretStr = Field(..., description="Groq API key")
+
+    _required_credentials = {"api_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Groq credentials"""
+        try:
+            client = Groq(api_key=self.api_key.get_secret_value())
+            await client.chat.completions.create(
+                messages=[{"role": "user", "content": "Hi"}],
+                model="mixtral-8x7b-32768",
+                max_tokens=1,
+            )
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid Groq credentials: {str(e)}")

airtrain/integrations/groq/skills.py

@@ -0,0 +1,41 @@
+from typing import Optional, Dict, Any
+from pydantic import Field
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import GroqCredentials
+
+
+class GroqInput(InputSchema):
+    """Schema for Groq 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="mixtral-8x7b", description="Groq model to use")
+    max_tokens: int = Field(default=1024, description="Maximum tokens in response")
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+
+
+class GroqOutput(OutputSchema):
+    """Schema for Groq 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 GroqChatSkill(Skill[GroqInput, GroqOutput]):
+    """Skill for Groq - Not Implemented"""
+
+    input_schema = GroqInput
+    output_schema = GroqOutput
+
+    def __init__(self, credentials: Optional[GroqCredentials] = None):
+        raise NotImplementedError("GroqChatSkill is not implemented yet")
+
+    def process(self, input_data: GroqInput) -> GroqOutput:
+        raise NotImplementedError("GroqChatSkill is not implemented yet")

airtrain/integrations/ollama/__init__.py

@@ -0,0 +1,6 @@
+"""Ollama integration module"""
+
+from .credentials import OllamaCredentials
+from .skills import OllamaChatSkill
+
+__all__ = ["OllamaCredentials", "OllamaChatSkill"]

airtrain/integrations/ollama/credentials.py

@@ -0,0 +1,26 @@
+from pydantic import Field
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+from importlib.util import find_spec
+
+
+class OllamaCredentials(BaseCredentials):
+    """Ollama credentials"""
+
+    host: str = Field(default="http://localhost:11434", description="Ollama host URL")
+    timeout: int = Field(default=30, description="Request timeout in seconds")
+
+    async def validate_credentials(self) -> bool:
+        """Validate Ollama credentials"""
+        if find_spec("ollama") is None:
+            raise CredentialValidationError(
+                "Ollama package is not installed. Please install it using: pip install ollama"
+            )
+
+        try:
+            from ollama import Client
+
+            client = Client(host=self.host)
+            await client.list()
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid Ollama connection: {str(e)}")

airtrain/integrations/ollama/skills.py

@@ -0,0 +1,41 @@
+from typing import Optional, Dict, Any
+from pydantic import Field
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import OllamaCredentials
+
+
+class OllamaInput(InputSchema):
+    """Schema for Ollama 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="llama2", description="Ollama model to use")
+    max_tokens: int = Field(default=1024, description="Maximum tokens in response")
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+
+
+class OllamaOutput(OutputSchema):
+    """Schema for Ollama 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 OllamaChatSkill(Skill[OllamaInput, OllamaOutput]):
+    """Skill for Ollama - Not Implemented"""
+
+    input_schema = OllamaInput
+    output_schema = OllamaOutput
+
+    def __init__(self, credentials: Optional[OllamaCredentials] = None):
+        raise NotImplementedError("OllamaChatSkill is not implemented yet")
+
+    def process(self, input_data: OllamaInput) -> OllamaOutput:
+        raise NotImplementedError("OllamaChatSkill is not implemented yet")

airtrain/integrations/openai/__init__.py

@@ -0,0 +1,19 @@
+from .skills import (
+    OpenAIChatSkill,
+    OpenAIInput,
+    OpenAIParserSkill,
+    OpenAIOutput,
+    OpenAIParserInput,
+    OpenAIParserOutput,
+)
+from .credentials import OpenAICredentials
+
+__all__ = [
+    "OpenAIChatSkill",
+    "OpenAIInput",
+    "OpenAIParserSkill",
+    "OpenAIParserInput",
+    "OpenAIParserOutput",
+    "OpenAICredentials",
+    "OpenAIOutput",
+]

airtrain/integrations/openai/chinese_assistant.py

@@ -0,0 +1,42 @@
+from typing import Optional, TypeVar
+from pydantic import Field
+from .skills import OpenAIChatSkill, OpenAIInput, OpenAIOutput
+from .credentials import OpenAICredentials
+
+T = TypeVar("T", bound=OpenAIInput)
+
+
+class ChineseAssistantInput(OpenAIInput):
+    """Schema for Chinese Assistant input"""
+
+    user_input: str = Field(
+        ..., description="User's input text (can be in any language)"
+    )
+    system_prompt: str = Field(
+        default="你是一个有帮助的助手。请用中文回答所有问题，即使问题是用其他语言问的。回答要准确、礼貌、专业。",
+        description="System prompt in Chinese",
+    )
+    model: str = Field(default="gpt-4o", description="OpenAI model to use")
+    max_tokens: int = Field(default=8096, description="Maximum tokens in response")
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+
+
+class ChineseAssistantSkill(OpenAIChatSkill):
+    """Skill for Chinese language assistance"""
+
+    input_schema = ChineseAssistantInput
+    output_schema = OpenAIOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        super().__init__(credentials)
+
+    def process(self, input_data: T) -> OpenAIOutput:
+        # Add language check to ensure response is in Chinese
+        if "你是" not in input_data.system_prompt:
+            input_data.system_prompt = (
+                "你是一个中文助手。" + input_data.system_prompt + "请用中文回答。"
+            )
+
+        return super().process(input_data)

airtrain/integrations/openai/credentials.py

@@ -0,0 +1,39 @@
+from datetime import datetime, timedelta
+from typing import Optional
+from pydantic import Field, SecretStr, validator
+from openai import OpenAI
+
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+
+
+class OpenAICredentials(BaseCredentials):
+    """OpenAI API credentials with enhanced validation"""
+
+    openai_api_key: SecretStr = Field(..., description="OpenAI API key")
+    openai_organization_id: Optional[str] = Field(
+        None, description="OpenAI organization ID", pattern="^org-[A-Za-z0-9]{24}$"
+    )
+
+    _required_credentials = {"openai_api_key"}
+
+    @validator("openai_api_key")
+    def validate_api_key_format(cls, v: SecretStr) -> SecretStr:
+        key = v.get_secret_value()
+        if not key.startswith("sk-"):
+            raise ValueError("OpenAI API key must start with 'sk-'")
+        if len(key) < 40:
+            raise ValueError("OpenAI API key appears to be too short")
+        return v
+
+    async def validate_credentials(self) -> bool:
+        """Validate credentials by making a test API call"""
+        try:
+            client = OpenAI(
+                api_key=self.openai_api_key.get_secret_value(),
+                organization=self.openai_organization_id,
+            )
+            # Make minimal API call to validate
+            await client.models.list(limit=1)
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid OpenAI credentials: {str(e)}")

airtrain/integrations/openai/skills.py

@@ -0,0 +1,208 @@
+from typing import List, Optional, Dict, Any, TypeVar, Type
+from pydantic import Field, BaseModel
+from openai import OpenAI
+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 OpenAICredentials
+
+
+class OpenAIInput(InputSchema):
+    """Schema for OpenAI 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="gpt-4o", description="OpenAI model to use")
+    max_tokens: int = Field(default=8192, description="Maximum tokens in response")
+    temperature: float = Field(
+        default=0.2, 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",
+    )
+    functions: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description="Optional function definitions for function calling",
+    )
+    function_call: Optional[str] = Field(
+        default=None,
+        description="Controls function calling behavior",
+    )
+
+
+class OpenAIOutput(OutputSchema):
+    """Schema for OpenAI 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"
+    )
+    function_call: Optional[Dict[str, Any]] = Field(
+        default=None, description="Function call information if applicable"
+    )
+
+
+class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
+    """Skill for interacting with OpenAI's models"""
+
+    input_schema = OpenAIInput
+    output_schema = OpenAIOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def _encode_image(self, image_path: Path) -> Dict[str, Any]:
+        """Convert image to base64 for API consumption"""
+        try:
+            if not image_path.exists():
+                raise FileNotFoundError(f"Image file not found: {image_path}")
+
+            with open(image_path, "rb") as img_file:
+                encoded = base64.b64encode(img_file.read()).decode()
+                return {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/jpeg;base64,{encoded}"},
+                }
+        except Exception as e:
+            logger.error(f"Failed to encode image {image_path}: {str(e)}")
+            raise ProcessingError(f"Image encoding failed: {str(e)}")
+
+    def process(self, input_data: OpenAIInput) -> OpenAIOutput:
+        """Process the input using OpenAI's API"""
+        try:
+            logger.info(f"Processing request with model {input_data.model}")
+
+            # Prepare message content
+            content = []
+
+            # Add text content
+            content.append({"type": "text", "text": input_data.user_input})
+
+            # Add images if provided
+            if input_data.images:
+                logger.debug(f"Processing {len(input_data.images)} images")
+                for image_path in input_data.images:
+                    content.append(self._encode_image(image_path))
+
+            # Prepare messages
+            messages = [
+                {"role": "system", "content": input_data.system_prompt},
+                {"role": "user", "content": content},
+            ]
+
+            # Create completion parameters
+            params = {
+                "model": input_data.model,
+                "messages": messages,
+                "temperature": input_data.temperature,
+                "max_tokens": input_data.max_tokens,
+            }
+
+            # Add function calling if provided
+            if input_data.functions:
+                params["functions"] = input_data.functions
+                params["function_call"] = input_data.function_call
+
+            # Create chat completion
+            response = self.client.chat.completions.create(**params)
+
+            # Extract function call if present
+            function_call = None
+            if response.choices[0].message.function_call:
+                function_call = response.choices[0].message.function_call.model_dump()
+
+            logger.success("Successfully processed OpenAI request")
+
+            return OpenAIOutput(
+                response=response.choices[0].message.content or "",
+                used_model=response.model,
+                usage={
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "completion_tokens": response.usage.completion_tokens,
+                    "total_tokens": response.usage.total_tokens,
+                },
+                function_call=function_call,
+            )
+
+        except Exception as e:
+            logger.exception(f"OpenAI processing failed: {str(e)}")
+            raise ProcessingError(f"OpenAI processing failed: {str(e)}")
+
+
+ResponseT = TypeVar("ResponseT", bound=BaseModel)
+
+
+class OpenAIParserInput(InputSchema):
+    """Schema for OpenAI structured output input"""
+
+    user_input: str
+    system_prompt: str = "You are a helpful assistant that provides structured data."
+    model: str = "gpt-4o"
+    temperature: float = 0.7
+    max_tokens: Optional[int] = None
+    response_model: Type[ResponseT]
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class OpenAIParserOutput(OutputSchema):
+    """Schema for OpenAI structured output"""
+
+    parsed_response: BaseModel
+    used_model: str
+    tokens_used: int
+
+
+class OpenAIParserSkill(Skill[OpenAIParserInput, OpenAIParserOutput]):
+    """Skill for getting structured responses from OpenAI"""
+
+    input_schema = OpenAIParserInput
+    output_schema = OpenAIParserOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def process(self, input_data: OpenAIParserInput) -> OpenAIParserOutput:
+        try:
+            # Use parse method instead of create
+            completion = self.client.beta.chat.completions.parse(
+                model=input_data.model,
+                messages=[
+                    {"role": "system", "content": input_data.system_prompt},
+                    {"role": "user", "content": input_data.user_input},
+                ],
+                response_format=input_data.response_model,
+            )
+
+            if completion.choices[0].message.parsed is None:
+                raise ProcessingError("Failed to parse response")
+
+            return OpenAIParserOutput(
+                parsed_response=completion.choices[0].message.parsed,
+                used_model=completion.model,
+                tokens_used=completion.usage.total_tokens,
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"OpenAI parsing failed: {str(e)}")

airtrain/integrations/sambanova/__init__.py

@@ -0,0 +1,6 @@
+"""Sambanova integration module"""
+
+from .credentials import SambanovaCredentials
+from .skills import SambanovaChatSkill
+
+__all__ = ["SambanovaCredentials", "SambanovaChatSkill"]

airtrain/integrations/sambanova/credentials.py

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

airtrain/integrations/sambanova/skills.py

@@ -0,0 +1,41 @@
+from typing import Optional, Dict, Any
+from pydantic import Field
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import SambanovaCredentials
+
+
+class SambanovaInput(InputSchema):
+    """Schema for Sambanova 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="sambanova-llm", description="Sambanova model to use")
+    max_tokens: int = Field(default=1024, description="Maximum tokens in response")
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+
+
+class SambanovaOutput(OutputSchema):
+    """Schema for Sambanova 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 SambanovaChatSkill(Skill[SambanovaInput, SambanovaOutput]):
+    """Skill for Sambanova - Not Implemented"""
+
+    input_schema = SambanovaInput
+    output_schema = SambanovaOutput
+
+    def __init__(self, credentials: Optional[SambanovaCredentials] = None):
+        raise NotImplementedError("SambanovaChatSkill is not implemented yet")
+
+    def process(self, input_data: SambanovaInput) -> SambanovaOutput:
+        raise NotImplementedError("SambanovaChatSkill is not implemented yet")

airtrain/integrations/together/__init__.py

@@ -0,0 +1,6 @@
+"""Together AI integration module"""
+
+from .credentials import TogetherAICredentials
+from .skills import TogetherAIChatSkill
+
+__all__ = ["TogetherAICredentials", "TogetherAIChatSkill"]

airtrain/integrations/together/credentials.py

@@ -0,0 +1,22 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+import together
+
+
+class TogetherAICredentials(BaseCredentials):
+    """Together AI credentials"""
+
+    api_key: SecretStr = Field(..., description="Together AI API key")
+
+    _required_credentials = {"api_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Together AI credentials"""
+        try:
+            together.api_key = self.api_key.get_secret_value()
+            await together.Models.list()
+            return True
+        except Exception as e:
+            raise CredentialValidationError(
+                f"Invalid Together AI credentials: {str(e)}"
+            )

airtrain/integrations/together/skills.py

@@ -0,0 +1,43 @@
+from typing import Optional, Dict, Any
+from pydantic import Field
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import TogetherAICredentials
+
+
+class TogetherAIInput(InputSchema):
+    """Schema for Together AI 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="togethercomputer/llama-2-70b", description="Together AI model to use"
+    )
+    max_tokens: int = Field(default=1024, description="Maximum tokens in response")
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+
+
+class TogetherAIOutput(OutputSchema):
+    """Schema for Together AI 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 TogetherAIChatSkill(Skill[TogetherAIInput, TogetherAIOutput]):
+    """Skill for Together AI - Not Implemented"""
+
+    input_schema = TogetherAIInput
+    output_schema = TogetherAIOutput
+
+    def __init__(self, credentials: Optional[TogetherAICredentials] = None):
+        raise NotImplementedError("TogetherAIChatSkill is not implemented yet")
+
+    def process(self, input_data: TogetherAIInput) -> TogetherAIOutput:
+        raise NotImplementedError("TogetherAIChatSkill is not implemented yet")

{airtrain-0.1.6.dist-info → airtrain-0.1.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.6
+Version: 0.1.8
 Summary: A platform for building and deploying AI agents with structured skills
 Home-page: https://github.com/rosaboyle/airtrain.dev
 Author: Dheeraj Pai

airtrain-0.1.8.dist-info/RECORD

@@ -0,0 +1,38 @@
+airtrain/__init__.py,sha256=VKqqTKi_5O7itXPoDiYPbAAgnxMd5opQeJHsRPyUNTM,312
+airtrain/core/__init__.py,sha256=9h7iKwTzZocCPc9bU6j8bA02BokteWIOcO1uaqGMcrk,254
+airtrain/core/credentials.py,sha256=PgQotrQc46J5djidKnkK1znUv3fyNkUFDO-m2Kn_Gzo,4006
+airtrain/core/schemas.py,sha256=MMXrDviC4gRea_QaPpbjgO--B_UKxnD7YrxqZOLJZZU,7003
+airtrain/core/skills.py,sha256=LljalzeSHK5eQPTAOEAYc5D8Qn1kVSfiz9WgziTD5UM,4688
+airtrain/integrations/__init__.py,sha256=Y-nxbWlPGy0Txt6qkRzLKQGapfAfL_q9W6U3ISTSg_I,1486
+airtrain/integrations/anthropic/__init__.py,sha256=qwlWLDh1rEVizYFbW8430z-f1SxHio7_Gaw5cCTUtoo,274
+airtrain/integrations/anthropic/credentials.py,sha256=hlTSw9HX66kYNaeQUtn0JjdZQBMNkzzFOJOoLOOzvcY,1246
+airtrain/integrations/anthropic/skills.py,sha256=sT7dBYPVCsICYjgBjUlyyP84A8h9OkbgkslvKHk3Tjs,5273
+airtrain/integrations/aws/__init__.py,sha256=3x7v2NxpAfI-U-YgwQeH5PtsmUrNLPMfLyUGFLiBjbs,155
+airtrain/integrations/aws/credentials.py,sha256=nN-daKAl7qOb_VdRpsThG8gN5GeSUkx-ji5E_gF_vYw,1444
+airtrain/integrations/aws/skills.py,sha256=TQiMXeXRRcJ14fe8Xi7Uk20iS6_INbcznuLGtMorcKY,3870
+airtrain/integrations/cerebras/__init__.py,sha256=zAD-qV38OzHhMCz1z-NvjjqcYEhURbm8RWTOKHNqbew,174
+airtrain/integrations/cerebras/credentials.py,sha256=IFkn8LxMAaOpvEWXDpb94VQGtqcDxQ7rZHKH-tX4Nuw,884
+airtrain/integrations/cerebras/skills.py,sha256=O9vwFzvv_tUOwFOVE8CszAQEac711eVYVUj_8dVMTpc,1596
+airtrain/integrations/google/__init__.py,sha256=INZFNOcNebz3m-Ggk07ZjmX0kNHIbTe_St9gBlZBki8,176
+airtrain/integrations/google/credentials.py,sha256=yyl-MWl06wr4SWvcvJGSpJ3hGTz21ByrRSr_3np5cbU,1030
+airtrain/integrations/google/skills.py,sha256=uwmgetl5Ien7fLOA5HIZdqoL6AZnexFDyzfsrGuJ1RU,1606
+airtrain/integrations/groq/__init__.py,sha256=B_X2fXbsJfFD6GquKeVCsEJjwd9Ygbq1uEHlV4Jy7YE,154
+airtrain/integrations/groq/credentials.py,sha256=A8-VIyoZTkHFQb-O-lmu-UrgaLZ3hfWfzzigkYteESk,829
+airtrain/integrations/groq/skills.py,sha256=Qy6SBAb19SzOFuqgcLyzdyRBp4D7jKqsEeJ6UTDaqMM,1528
+airtrain/integrations/ollama/__init__.py,sha256=zMHBsGzViVrvxAeJmfq6r-ZfSE6Dy5QcKLhe4d5fEcM,164
+airtrain/integrations/ollama/credentials.py,sha256=D7O4kUAb_VHs5s1ncUN9Ezhu5PvLfgj3RifAkB9sEZk,940
+airtrain/integrations/ollama/skills.py,sha256=M_Un8D5VJ5XtPEq9IClzqV3jCPBoFTSm2ve6EO8W2JU,1556
+airtrain/integrations/openai/__init__.py,sha256=K-NY2_T1T6SEOgkpbUA55cWvK2nr2NOJgLCqmmtaCno,371
+airtrain/integrations/openai/chinese_assistant.py,sha256=MMhv4NBOoEQ0O22ZZtP255rd5ajHC9l6FPWIjpqxBOA,1581
+airtrain/integrations/openai/credentials.py,sha256=NfRyp1QgEtgm8cxt2-BOLq-6d0X-Pcm80NnfHM8p0FY,1470
+airtrain/integrations/openai/skills.py,sha256=Olg9-6f_p2XgkVwwcB9tvjAMApmM2EK81i8LP4qVVvs,7676
+airtrain/integrations/sambanova/__init__.py,sha256=dp_263iOckM_J9pOEvyqpf3FrejD6-_x33r0edMCTe0,179
+airtrain/integrations/sambanova/credentials.py,sha256=U36RAEIPNuwo-vTrt3U9kkkj2GfdqSclA1ttOYHxS-w,784
+airtrain/integrations/sambanova/skills.py,sha256=Po1ur_QFwzVIugbkk2mt73WdXDz_Gr9ASlUc9Y12Kok,1614
+airtrain/integrations/together/__init__.py,sha256=we4KXn_pUs6Dxo3QcB-t40BSRraQFdKg2nXw7yi2FjM,185
+airtrain/integrations/together/credentials.py,sha256=y5M6ZQrfYJLJbClxEasq4HaVyZM0l5lFshwVP6jq2E4,720
+airtrain/integrations/together/skills.py,sha256=YMOULyk2TX32rCjhxK29e4ehn8iIzMXpg3xmdYtuyQQ,1664
+airtrain-0.1.8.dist-info/METADATA,sha256=o8NbBrkTDmyZBRZJQ2fxti7ffsVjD8U1F2R-w0aSbpw,4380
+airtrain-0.1.8.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+airtrain-0.1.8.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.8.dist-info/RECORD,,

airtrain-0.1.6.dist-info/RECORD

@@ -1,10 +0,0 @@
-airtrain/__init__.py,sha256=t0n2IItXdHzaLZf_1zSUmSf3buF-8Y4FJol1LAzSclk,312
-airtrain/core/__init__.py,sha256=9h7iKwTzZocCPc9bU6j8bA02BokteWIOcO1uaqGMcrk,254
-airtrain/core/credentials.py,sha256=PgQotrQc46J5djidKnkK1znUv3fyNkUFDO-m2Kn_Gzo,4006
-airtrain/core/schemas.py,sha256=MMXrDviC4gRea_QaPpbjgO--B_UKxnD7YrxqZOLJZZU,7003
-airtrain/core/skills.py,sha256=LljalzeSHK5eQPTAOEAYc5D8Qn1kVSfiz9WgziTD5UM,4688
-airtrain/integrations/__init__.py,sha256=PRKI_A-KE307C4lpXgFAsZA2oFtTl1kt_4CrRUF2rpU,832
-airtrain-0.1.6.dist-info/METADATA,sha256=C0hcwg_Am0cUqrv8vKfNJ_AeMdBfGhxrSEmzDbjJA7o,4380
-airtrain-0.1.6.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-airtrain-0.1.6.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
-airtrain-0.1.6.dist-info/RECORD,,