airtrain 0.1.6__tar.gz → 0.1.8__tar.gz

{airtrain-0.1.6/airtrain.egg-info → airtrain-0.1.8}/PKG-INFO

@@ -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.6 → airtrain-0.1.8}/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-0.1.6 → airtrain-0.1.8}/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-0.1.8/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-0.1.8/airtrain/integrations/aws/__init__.py

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

airtrain-0.1.8/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-0.1.8/airtrain/integrations/cerebras/__init__.py

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

airtrain-0.1.8/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-0.1.8/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-0.1.8/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-0.1.8/airtrain/integrations/groq/__init__.py

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

airtrain-0.1.8/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-0.1.8/airtrain/integrations/ollama/__init__.py

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

airtrain-0.1.8/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-0.1.8/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-0.1.8/airtrain/integrations/sambanova/__init__.py

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

airtrain-0.1.8/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-0.1.8/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-0.1.8/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 → airtrain-0.1.8/airtrain.egg-info}/PKG-INFO

@@ -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.6 → airtrain-0.1.8}/airtrain.egg-info/SOURCES.txt

@@ -25,18 +25,34 @@ airtrain/core/__pycache__/credentials.cpython-310.pyc
 airtrain/core/__pycache__/schemas.cpython-310.pyc
 airtrain/core/__pycache__/skills.cpython-310.pyc
 airtrain/integrations/__init__.py
+airtrain/integrations/anthropic/__init__.py
 airtrain/integrations/anthropic/credentials.py
 airtrain/integrations/anthropic/skills.py
+airtrain/integrations/aws/__init__.py
 airtrain/integrations/aws/credentials.py
+airtrain/integrations/aws/skills.py
+airtrain/integrations/cerebras/__init__.py
 airtrain/integrations/cerebras/credentials.py
+airtrain/integrations/cerebras/skills.py
+airtrain/integrations/google/__init__.py
 airtrain/integrations/google/credentials.py
+airtrain/integrations/google/skills.py
+airtrain/integrations/groq/__init__.py
 airtrain/integrations/groq/credentials.py
+airtrain/integrations/groq/skills.py
+airtrain/integrations/ollama/__init__.py
 airtrain/integrations/ollama/credentials.py
+airtrain/integrations/ollama/skills.py
+airtrain/integrations/openai/__init__.py
 airtrain/integrations/openai/chinese_assistant.py
 airtrain/integrations/openai/credentials.py
 airtrain/integrations/openai/skills.py
+airtrain/integrations/sambanova/__init__.py
 airtrain/integrations/sambanova/credentials.py
+airtrain/integrations/sambanova/skills.py
+airtrain/integrations/together/__init__.py
 airtrain/integrations/together/credentials.py
+airtrain/integrations/together/skills.py
 examples/credentials_usage.py
 examples/schema_usage.py
 examples/skill_usage.py

{airtrain-0.1.6 → airtrain-0.1.8}/examples/creating-skills/chinese_anthropic_assistant.py

@@ -36,6 +36,18 @@ class ChineseAnthropicSkill(AnthropicChatSkill):
     output_schema = AnthropicOutput
 
     def process(self, input_data: T) -> AnthropicOutput:
+        """
+        Process the input and ensure Chinese language response.
+
+        Args:
+            input_data: The input data containing user's query and settings
+
+        Returns:
+            AnthropicOutput: The model's response in Chinese
+
+        Raises:
+            ProcessingError: If processing fails
+        """
         if "你是" not in input_data.system_prompt:
             input_data.system_prompt = (
                 "你是一个中文助手。" + input_data.system_prompt + "请用中文回答。"