airtrain 0.1.3__tar.gz → 0.1.5__tar.gz

airtrain-0.1.5/.gitignore

@@ -0,0 +1,183 @@
+package
+.env
+.mypy_cache
+firebaseadmin.json
+**pyc=
+50mb_test3.bin
+*bin
+**bin
+token.pickle
+temp_workspace
+temp*
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# PyPI configuration file
+.pypirc

airtrain-0.1.5/EXPERIMENTS/integrations_examples/anthropic_with_image.py

@@ -0,0 +1,43 @@
+import anthropic
+
+client = anthropic.Anthropic(
+    # defaults to os.environ.get("ANTHROPIC_API_KEY")
+    api_key="my_api_key",
+)
+
+# Replace placeholders like {{PR_DESCRIPTION}} with real values,
+# because the SDK does not support variables.
+message = client.messages.create(
+    model="claude-3-5-sonnet-20241022",
+    max_tokens=8192,
+    temperature=0,
+    system="You are an experienced software engineer tasked with reviewing a GitHub Pull Request (PR). Your goal is to analyze the code quality and suggest improvements. Follow these steps carefully:\n\n1. Review the PR description:\n<PR_DESCRIPTION>\n{{PR_DESCRIPTION}}\n</PR_DESCRIPTION>\n\n2. Examine the code changes:\n<CODE_CHANGES>\n{{CODE_CHANGES}}\n</CODE_CHANGES>\n\n3. Consider any existing comments:\n<EXISTING_COMMENTS>\n{{EXISTING_COMMENTS}}\n</EXISTING_COMMENTS>\n\n4. Analyze the code quality:\n   a. Check for adherence to coding standards and best practices\n   b. Evaluate code readability and maintainability\n   c. Assess performance implications\n   d. Look for potential bugs or edge cases\n   e. Consider security implications\n\n5. Suggest improvements:\n   a. Identify areas where the code can be optimized or simplified\n   b. Propose alternative approaches if applicable\n   c. Recommend additional tests or error handling if needed\n\n6. Format your response as follows:\n   <code_review>\n   <quality_analysis>\n   Provide a detailed analysis of the code quality, addressing points 4a-4e.\n   </quality_analysis>\n\n   <improvement_suggestions>\n   List your suggestions for improvement, addressing points 5a-5c. Number each suggestion.\n   </improvement_suggestions>\n\n   <summary>\n   Provide a brief summary of your overall assessment and key recommendations.\n   </summary>\n   </code_review>\n\nRemember to be constructive and specific in your feedback. Use code snippets or pseudocode to illustrate your suggestions when appropriate. If you need clarification on any part of the code or PR description, state your assumptions clearly.\n\nDo not comment on aspects unrelated to code quality or potential improvements. Focus solely on the technical aspects of the code changes presented.",
+    messages=[
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": "\nAnalyze the above examples and give me some updates. Analyze this image as well.\n\nOne more image is this. Can you image this as well.",
+                },
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": "image/jpeg",
+                        "data": "<base64_encoded_image>",
+                    },
+                },
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": "image/jpeg",
+                        "data": "<base64_encoded_image>",
+                    },
+                },
+            ],
+        }
+    ],
+)
+print(message.content)

{airtrain-0.1.3 → airtrain-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.3
+Version: 0.1.5
 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.3 → airtrain-0.1.5}/airtrain/__init__.py

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

{airtrain-0.1.3 → airtrain-0.1.5}/airtrain/core/credentials.py

@@ -5,7 +5,7 @@ from pathlib import Path
 from abc import ABC, abstractmethod
 import dotenv
 from pydantic import BaseModel, Field, SecretStr
-import yaml  # type: ignore
+import yaml
 
 
 class CredentialError(Exception):
@@ -100,7 +100,7 @@ class BaseCredentials(BaseModel):
         else:
             raise ValueError(f"Unsupported file format: {file_path.suffix}")
 
-    def validate_credentials(self) -> None:
+    async def validate_credentials(self) -> bool:
         """Validate that all required credentials are present"""
         missing = []
         for field_name in self._required_credentials:
@@ -114,6 +114,7 @@ class BaseCredentials(BaseModel):
             raise CredentialValidationError(
                 f"Missing required credentials: {', '.join(missing)}"
             )
+        return True
 
     def clear_from_env(self) -> None:
         """Remove credentials from environment variables"""
@@ -122,32 +123,3 @@ class BaseCredentials(BaseModel):
             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-0.1.5/airtrain/integrations/__init__.py

@@ -0,0 +1,26 @@
+"""Airtrain integrations package"""
+
+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 .anthropic.skills import AnthropicChatSkill
+
+__all__ = [
+    "OpenAICredentials",
+    "AWSCredentials",
+    "GoogleCloudCredentials",
+    "AnthropicCredentials",
+    "AnthropicChatSkill",
+    "GroqCredentials",
+    "TogetherAICredentials",
+    "OllamaCredentials",
+    "SambanovaCredentials",
+    "CerebrasCredentials",
+]

airtrain-0.1.5/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-0.1.5/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-0.1.5/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-0.1.5/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-0.1.5/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-0.1.5/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-0.1.5/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-0.1.5/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)