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

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=131072, 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,19 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+
+
+class CerebrasCredentials(BaseCredentials):
+    """Cerebras credentials"""
+
+    cerebras_api_key: SecretStr = Field(..., description="Cerebras API key")
+
+    _required_credentials = {"cerebras_api_key"}
+
+    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,127 @@
+from typing import List, Optional, Dict, Any, Generator
+from pydantic import Field
+from cerebras.cloud.sdk import Cerebras
+from loguru import logger
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import CerebrasCredentials
+
+
+class CerebrasInput(InputSchema):
+    """Schema for Cerebras 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",
+    )
+    conversation_history: List[Dict[str, str]] = Field(
+        default_factory=list,
+        description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
+    )
+    model: str = Field(default="llama3.1-8b", description="Cerebras model to use")
+    max_tokens: Optional[int] = Field(
+        default=131072, description="Maximum tokens in response"
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+    stream: bool = Field(
+        default=False, description="Whether to stream the response progressively"
+    )
+
+
+class CerebrasOutput(OutputSchema):
+    """Schema for Cerebras 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")
+
+
+class CerebrasChatSkill(Skill[CerebrasInput, CerebrasOutput]):
+    """Skill for Cerebras chat"""
+
+    input_schema = CerebrasInput
+    output_schema = CerebrasOutput
+
+    def __init__(self, credentials: Optional[CerebrasCredentials] = None):
+        super().__init__()
+        self.credentials = credentials or CerebrasCredentials.from_env()
+        self.client = Cerebras(
+            api_key=self.credentials.cerebras_api_key.get_secret_value()
+        )
+
+    def _build_messages(self, input_data: CerebrasInput) -> List[Dict[str, str]]:
+        """
+        Build messages list from input data including conversation history.
+
+        Args:
+            input_data: The input data containing system prompt, conversation history, and user input
+
+        Returns:
+            List[Dict[str, str]]: List of messages in the format required by Cerebras
+        """
+        messages = [{"role": "system", "content": input_data.system_prompt}]
+
+        # Add conversation history if present
+        if input_data.conversation_history:
+            messages.extend(input_data.conversation_history)
+
+        # Add current user input
+        messages.append({"role": "user", "content": input_data.user_input})
+
+        return messages
+
+    def process_stream(self, input_data: CerebrasInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
+        try:
+            messages = self._build_messages(input_data)
+
+            stream = self.client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+                stream=True,
+            )
+
+            for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
+
+        except Exception as e:
+            logger.exception(f"Cerebras streaming failed: {str(e)}")
+            raise ProcessingError(f"Cerebras streaming failed: {str(e)}")
+
+    def process(self, input_data: CerebrasInput) -> CerebrasOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+                usage = {}  # Usage stats not available in streaming
+            else:
+                messages = self._build_messages(input_data)
+                response = self.client.chat.completions.create(
+                    model=input_data.model,
+                    messages=messages,
+                    temperature=input_data.temperature,
+                    max_tokens=input_data.max_tokens,
+                )
+                usage = (
+                    response.usage.model_dump() if hasattr(response, "usage") else {}
+                )
+
+            return CerebrasOutput(
+                response=response.choices[0].message.content,
+                used_model=input_data.model,
+                usage=usage,
+            )
+
+        except Exception as e:
+            logger.exception(f"Cerebras processing failed: {str(e)}")
+            raise ProcessingError(f"Cerebras processing failed: {str(e)}")

airtrain/integrations/combined/__init__.py

@@ -0,0 +1,21 @@
+"""Combined integration modules for Airtrain"""
+
+from .groq_fireworks_skills import (
+    GroqFireworksSkill,
+    GroqFireworksInput,
+    GroqFireworksOutput
+)
+from .list_models_factory import (
+    ListModelsSkillFactory, 
+    GenericListModelsInput, 
+    GenericListModelsOutput
+)
+
+__all__ = [
+    "GroqFireworksSkill",
+    "GroqFireworksInput",
+    "GroqFireworksOutput",
+    "ListModelsSkillFactory",
+    "GenericListModelsInput",
+    "GenericListModelsOutput"
+] 

airtrain/integrations/combined/groq_fireworks_skills.py

@@ -0,0 +1,126 @@
+from typing import Optional, Dict, Any, List
+from pydantic import Field
+import requests
+from groq import Groq
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from airtrain.integrations.fireworks.completion_skills import (
+    FireworksCompletionSkill,
+    FireworksCompletionInput,
+)
+
+
+class GroqFireworksInput(InputSchema):
+    """Schema for combined Groq and Fireworks input"""
+
+    user_input: str = Field(..., description="User's input text")
+    groq_model: str = Field(
+        default="mixtral-8x7b-32768", description="Groq model to use"
+    )
+    fireworks_model: str = Field(
+        default="accounts/fireworks/models/deepseek-r1",
+        description="Fireworks model to use",
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation"
+    )
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
+
+
+class GroqFireworksOutput(OutputSchema):
+    """Schema for combined Groq and Fireworks output"""
+
+    combined_response: str
+    groq_response: str
+    fireworks_response: str
+    used_models: Dict[str, str]
+    usage: Dict[str, Dict[str, int]]
+
+
+class GroqFireworksSkill(Skill[GroqFireworksInput, GroqFireworksOutput]):
+    """Skill combining Groq and Fireworks responses"""
+
+    input_schema = GroqFireworksInput
+    output_schema = GroqFireworksOutput
+
+    def __init__(
+        self,
+        groq_api_key: Optional[str] = None,
+        fireworks_skill: Optional[FireworksCompletionSkill] = None,
+    ):
+        """Initialize the skill with optional API keys"""
+        super().__init__()
+        self.groq_client = Groq(api_key=groq_api_key)
+        self.fireworks_skill = fireworks_skill or FireworksCompletionSkill()
+
+    def _get_groq_response(self, input_data: GroqFireworksInput) -> Dict[str, Any]:
+        """Get response from Groq"""
+        try:
+            completion = self.groq_client.chat.completions.create(
+                model=input_data.groq_model,
+                messages=[{"role": "user", "content": input_data.user_input}],
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+            return {
+                "response": completion.choices[0].message.content,
+                "usage": completion.usage.model_dump(),
+            }
+        except Exception as e:
+            raise ProcessingError(f"Groq request failed: {str(e)}")
+
+    def _get_fireworks_response(
+        self, groq_response: str, input_data: GroqFireworksInput
+    ) -> Dict[str, Any]:
+        """Get response from Fireworks"""
+        try:
+            formatted_prompt = (
+                f"<USER>{input_data.user_input}</USER>\n<ASSISTANT>{groq_response}"
+            )
+
+            fireworks_input = FireworksCompletionInput(
+                prompt=formatted_prompt,
+                model=input_data.fireworks_model,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+
+            result = self.fireworks_skill.process(fireworks_input)
+            return {"response": result.response, "usage": result.usage}
+        except Exception as e:
+            raise ProcessingError(f"Fireworks request failed: {str(e)}")
+
+    def process(self, input_data: GroqFireworksInput) -> GroqFireworksOutput:
+        """Process the input using both Groq and Fireworks"""
+        try:
+            # Get Groq response
+            groq_result = self._get_groq_response(input_data)
+
+            # Get Fireworks response
+            fireworks_result = self._get_fireworks_response(
+                groq_result["response"], input_data
+            )
+
+            # Combine responses in the required format
+            combined_response = (
+                f"<USER>{input_data.user_input}</USER>\n"
+                f"<ASSISTANT>{groq_result['response']} {fireworks_result['response']}"
+            )
+
+            return GroqFireworksOutput(
+                combined_response=combined_response,
+                groq_response=groq_result["response"],
+                fireworks_response=fireworks_result["response"],
+                used_models={
+                    "groq": input_data.groq_model,
+                    "fireworks": input_data.fireworks_model,
+                },
+                usage={
+                    "groq": groq_result["usage"],
+                    "fireworks": fireworks_result["usage"],
+                },
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Combined processing failed: {str(e)}")

airtrain/integrations/combined/list_models_factory.py

@@ -0,0 +1,210 @@
+from typing import Optional, Dict, Any, List
+from pydantic import Field
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from airtrain.core.credentials import BaseCredentials
+
+# Import existing list models skills
+from airtrain.integrations.openai.list_models import OpenAIListModelsSkill
+from airtrain.integrations.anthropic.list_models import AnthropicListModelsSkill
+from airtrain.integrations.together.list_models import TogetherListModelsSkill
+from airtrain.integrations.fireworks.list_models import FireworksListModelsSkill
+
+# Import credentials
+from airtrain.integrations.groq.credentials import GroqCredentials
+from airtrain.integrations.cerebras.credentials import CerebrasCredentials
+from airtrain.integrations.sambanova.credentials import SambanovaCredentials
+from airtrain.integrations.perplexity.credentials import PerplexityCredentials
+
+# Import Perplexity list models
+from airtrain.integrations.perplexity.list_models import PerplexityListModelsSkill
+
+
+# Generic list models input schema
+class GenericListModelsInput(InputSchema):
+    """Generic schema for listing models from any provider"""
+
+    api_models_only: bool = Field(
+        default=False,
+        description=(
+            "If True, fetch models from the API only. If False, use local config."
+        ),
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+        extra = "allow"
+
+
+# Generic list models output schema
+class GenericListModelsOutput(OutputSchema):
+    """Generic schema for list models output from any provider"""
+
+    models: List[Dict[str, Any]] = Field(
+        default_factory=list, description="List of models"
+    )
+    provider: str = Field(..., description="Provider name")
+
+
+# Base class for stub implementations
+class BaseListModelsSkill(Skill[GenericListModelsInput, GenericListModelsOutput]):
+    """Base skill for listing models"""
+
+    input_schema = GenericListModelsInput
+    output_schema = GenericListModelsOutput
+
+    def __init__(self, provider: str, credentials: Optional[BaseCredentials] = None):
+        """Initialize the skill with provider name and optional credentials"""
+        super().__init__()
+        self.provider = provider
+        self.credentials = credentials
+
+    def get_models(self) -> List[Dict[str, Any]]:
+        """Return list of models. To be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement get_models()")
+
+    def process(self, input_data: GenericListModelsInput) -> GenericListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            models = self.get_models()
+            return GenericListModelsOutput(models=models, provider=self.provider)
+        except Exception as e:
+            raise ProcessingError(f"Failed to list {self.provider} models: {str(e)}")
+
+
+# Groq implementation
+class GroqListModelsSkill(BaseListModelsSkill):
+    """Skill for listing Groq models"""
+
+    def __init__(self, credentials: Optional[GroqCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__(provider="groq", credentials=credentials)
+
+    def get_models(self) -> List[Dict[str, Any]]:
+        """Return list of Groq models."""
+        # Default Groq models from trmx_agent config
+        models = [
+            {
+                "id": "llama-3.3-70b-versatile",
+                "display_name": "Llama 3.3 70B Versatile (Tool Use)",
+            },
+            {
+                "id": "llama-3.1-8b-instant",
+                "display_name": "Llama 3.1 8B Instant (Tool Use)",
+            },
+            {
+                "id": "mixtral-8x7b-32768",
+                "display_name": "Mixtral 8x7B (32K) (Tool Use)",
+            },
+            {"id": "gemma2-9b-it", "display_name": "Gemma 2 9B IT (Tool Use)"},
+            {"id": "qwen-qwq-32b", "display_name": "Qwen QWQ 32B (Tool Use)"},
+            {
+                "id": "qwen-2.5-coder-32b",
+                "display_name": "Qwen 2.5 Coder 32B (Tool Use)",
+            },
+            {"id": "qwen-2.5-32b", "display_name": "Qwen 2.5 32B (Tool Use)"},
+            {
+                "id": "deepseek-r1-distill-qwen-32b",
+                "display_name": "DeepSeek R1 Distill Qwen 32B (Tool Use)",
+            },
+            {
+                "id": "deepseek-r1-distill-llama-70b",
+                "display_name": "DeepSeek R1 Distill Llama 70B (Tool Use)",
+            },
+        ]
+        return models
+
+
+# Cerebras implementation
+class CerebrasListModelsSkill(BaseListModelsSkill):
+    """Skill for listing Cerebras models"""
+
+    def __init__(self, credentials: Optional[CerebrasCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__(provider="cerebras", credentials=credentials)
+
+    def get_models(self) -> List[Dict[str, Any]]:
+        """Return list of Cerebras models."""
+        # Default Cerebras models from trmx_agent config
+        models = [
+            {
+                "id": "cerebras/Cerebras-GPT-13B-v0.1",
+                "display_name": "Cerebras GPT 13B v0.1",
+            },
+            {
+                "id": "cerebras/Cerebras-GPT-111M-v0.9",
+                "display_name": "Cerebras GPT 111M v0.9",
+            },
+            {
+                "id": "cerebras/Cerebras-GPT-590M-v0.7",
+                "display_name": "Cerebras GPT 590M v0.7",
+            },
+        ]
+        return models
+
+
+# Sambanova implementation
+class SambanovaListModelsSkill(BaseListModelsSkill):
+    """Skill for listing Sambanova models"""
+
+    def __init__(self, credentials: Optional[SambanovaCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__(provider="sambanova", credentials=credentials)
+
+    def get_models(self) -> List[Dict[str, Any]]:
+        """Return list of Sambanova models."""
+        # Limited Sambanova model information
+        models = [
+            {"id": "sambanova/samba-1", "display_name": "Samba-1"},
+            {"id": "sambanova/samba-2", "display_name": "Samba-2"},
+        ]
+        return models
+
+
+# Factory class
+class ListModelsSkillFactory:
+    """Factory for creating list models skills for different providers"""
+
+    # Map provider names to their corresponding list models skills
+    _PROVIDER_MAP = {
+        "openai": OpenAIListModelsSkill,
+        "anthropic": AnthropicListModelsSkill,
+        "together": TogetherListModelsSkill,
+        "fireworks": FireworksListModelsSkill,
+        "groq": GroqListModelsSkill,
+        "cerebras": CerebrasListModelsSkill,
+        "sambanova": SambanovaListModelsSkill,
+        "perplexity": PerplexityListModelsSkill,
+    }
+
+    @classmethod
+    def get_skill(cls, provider: str, credentials=None):
+        """Return a list models skill for the specified provider
+
+        Args:
+            provider (str): The provider name (case-insensitive)
+            credentials: Optional credentials for the provider
+
+        Returns:
+            A ListModelsSkill instance for the specified provider
+
+        Raises:
+            ValueError: If the provider is not supported
+        """
+        provider = provider.lower()
+
+        if provider not in cls._PROVIDER_MAP:
+            supported = ", ".join(cls.get_supported_providers())
+            raise ValueError(
+                f"Unsupported provider: {provider}. "
+                f"Supported providers are: {supported}"
+            )
+
+        skill_class = cls._PROVIDER_MAP[provider]
+        return skill_class(credentials=credentials)
+
+    @classmethod
+    def get_supported_providers(cls):
+        """Return a list of supported provider names"""
+        return list(cls._PROVIDER_MAP.keys())

airtrain/integrations/fireworks/__init__.py

@@ -0,0 +1,21 @@
+"""Fireworks AI integration module"""
+
+from .credentials import FireworksCredentials
+from .skills import FireworksChatSkill, FireworksInput, FireworksOutput
+from .list_models import (
+    FireworksListModelsSkill,
+    FireworksListModelsInput,
+    FireworksListModelsOutput,
+)
+from .models import FireworksModel
+
+__all__ = [
+    "FireworksCredentials",
+    "FireworksChatSkill",
+    "FireworksInput",
+    "FireworksOutput",
+    "FireworksListModelsSkill",
+    "FireworksListModelsInput",
+    "FireworksListModelsOutput",
+    "FireworksModel",
+]