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

airtrain/integrations/google/skills.py

@@ -0,0 +1,122 @@
+from typing import List, Optional, Dict, Any
+from pydantic import Field
+import google.generativeai as genai
+from loguru import logger
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import GeminiCredentials
+
+
+class GoogleGenerationConfig(InputSchema):
+    """Schema for Google generation config"""
+
+    temperature: float = Field(
+        default=1.0, description="Temperature for response generation", ge=0, le=1
+    )
+    top_p: float = Field(
+        default=0.95, description="Top p sampling parameter", ge=0, le=1
+    )
+    top_k: int = Field(default=40, description="Top k sampling parameter")
+    max_output_tokens: int = Field(
+        default=8192, description="Maximum tokens in response"
+    )
+    response_mime_type: str = Field(
+        default="text/plain", description="Response MIME type"
+    )
+
+
+class GoogleInput(InputSchema):
+    """Schema for Google 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 | List[Dict[str, str]]]] = Field(
+        default_factory=list,
+        description="List of conversation messages in Google's format",
+    )
+    model: str = Field(default="gemini-1.5-flash", description="Google model to use")
+    generation_config: GoogleGenerationConfig = Field(
+        default_factory=GoogleGenerationConfig, description="Generation configuration"
+    )
+
+
+class GoogleOutput(OutputSchema):
+    """Schema for Google 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 GoogleChatSkill(Skill[GoogleInput, GoogleOutput]):
+    """Skill for Google chat"""
+
+    input_schema = GoogleInput
+    output_schema = GoogleOutput
+
+    def __init__(self, credentials: Optional[GeminiCredentials] = None):
+        super().__init__()
+        self.credentials = credentials or GeminiCredentials.from_env()
+        genai.configure(api_key=self.credentials.gemini_api_key.get_secret_value())
+
+    def _convert_history_format(
+        self, history: List[Dict[str, str]]
+    ) -> List[Dict[str, List[Dict[str, str]]]]:
+        """Convert standard history format to Google's format"""
+        google_history = []
+        for msg in history:
+            google_msg = {
+                "role": "user" if msg["role"] == "user" else "model",
+                "parts": [{"text": msg["content"]}],
+            }
+            google_history.append(google_msg)
+        return google_history
+
+    def process(self, input_data: GoogleInput) -> GoogleOutput:
+        try:
+            # Create generation config
+            generation_config = {
+                "temperature": input_data.generation_config.temperature,
+                "top_p": input_data.generation_config.top_p,
+                "top_k": input_data.generation_config.top_k,
+                "max_output_tokens": input_data.generation_config.max_output_tokens,
+                "response_mime_type": input_data.generation_config.response_mime_type,
+            }
+
+            # Initialize model
+            model = genai.GenerativeModel(
+                model_name=input_data.model,
+                generation_config=generation_config,
+                system_instruction=input_data.system_prompt,
+            )
+
+            # Convert history format if needed
+            history = (
+                input_data.conversation_history
+                if input_data.conversation_history
+                else self._convert_history_format([])
+            )
+
+            # Start chat session
+            chat = model.start_chat(history=history)
+
+            # Send message and get response
+            response = chat.send_message(input_data.user_input)
+
+            return GoogleOutput(
+                response=response.text,
+                used_model=input_data.model,
+                usage={
+                    "prompt_tokens": 0,
+                    "completion_tokens": 0,
+                    "total_tokens": 0,
+                },  # Google API doesn't provide usage stats
+            )
+
+        except Exception as e:
+            logger.exception(f"Google processing failed: {str(e)}")
+            raise ProcessingError(f"Google processing failed: {str(e)}")

airtrain/integrations/groq/__init__.py

@@ -0,0 +1,23 @@
+"""Groq integration module"""
+
+from .credentials import GroqCredentials
+from .skills import GroqChatSkill
+from .models_config import (
+    get_model_config,
+    get_default_model,
+    supports_tool_use,
+    supports_parallel_tool_use,
+    supports_json_mode,
+    GROQ_MODELS_CONFIG,
+)
+
+__all__ = [
+    "GroqCredentials", 
+    "GroqChatSkill",
+    "get_model_config",
+    "get_default_model",
+    "supports_tool_use",
+    "supports_parallel_tool_use",
+    "supports_json_mode",
+    "GROQ_MODELS_CONFIG",
+]

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"""
+
+    groq_api_key: SecretStr = Field(..., description="Groq API key")
+
+    _required_credentials = {"groq_api_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Groq credentials"""
+        try:
+            client = Groq(api_key=self.groq_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/models_config.py

@@ -0,0 +1,162 @@
+"""Configuration of Groq model capabilities."""
+
+from typing import Dict, Any
+
+
+# Model configuration with capabilities for each model
+GROQ_MODELS_CONFIG = {
+    "llama-3.3-70b-versatile": {
+        "name": "Llama 3.3 70B Versatile",
+        "context_window": 128000,
+        "max_completion_tokens": 32768,
+        "tool_use": True,
+        "parallel_tool_use": True,
+        "json_mode": True,
+    },
+    "llama-3.1-8b-instant": {
+        "name": "Llama 3.1 8B Instant",
+        "context_window": 128000,
+        "max_completion_tokens": 8192,
+        "tool_use": True,
+        "parallel_tool_use": True,
+        "json_mode": True,
+    },
+    "mixtral-8x7b-32768": {
+        "name": "Mixtral 8x7B (32K)",
+        "context_window": 32768,
+        "max_completion_tokens": 8192,
+        "tool_use": True,
+        "parallel_tool_use": False,
+        "json_mode": True,
+    },
+    "gemma2-9b-it": {
+        "name": "Gemma 2 9B IT",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "tool_use": True,
+        "parallel_tool_use": False,
+        "json_mode": True,
+    },
+    "qwen-qwq-32b": {
+        "name": "Qwen QWQ 32B",
+        "context_window": 128000,
+        "max_completion_tokens": 16384,
+        "tool_use": True,
+        "parallel_tool_use": True,
+        "json_mode": True,
+    },
+    "qwen-2.5-coder-32b": {
+        "name": "Qwen 2.5 Coder 32B",
+        "context_window": 128000,
+        "max_completion_tokens": 16384,
+        "tool_use": True,
+        "parallel_tool_use": True,
+        "json_mode": True,
+    },
+    "qwen-2.5-32b": {
+        "name": "Qwen 2.5 32B",
+        "context_window": 128000,
+        "max_completion_tokens": 16384,
+        "tool_use": True,
+        "parallel_tool_use": True,
+        "json_mode": True,
+    },
+    "deepseek-r1-distill-qwen-32b": {
+        "name": "DeepSeek R1 Distill Qwen 32B",
+        "context_window": 128000,
+        "max_completion_tokens": 16384,
+        "tool_use": True,
+        "parallel_tool_use": True,
+        "json_mode": True,
+    },
+    "deepseek-r1-distill-llama-70b": {
+        "name": "DeepSeek R1 Distill Llama 70B",
+        "context_window": 128000,
+        "max_completion_tokens": 16384,
+        "tool_use": True,
+        "parallel_tool_use": True,
+        "json_mode": True,
+    },
+    "deepseek-r1-distill-llama-70b-specdec": {
+        "name": "DeepSeek R1 Distill Llama 70B SpecDec",
+        "context_window": 128000,
+        "max_completion_tokens": 16384,
+        "tool_use": False,
+        "parallel_tool_use": False,
+        "json_mode": False,
+    },
+    "llama3-70b-8192": {
+        "name": "Llama 3 70B (8K)",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "tool_use": False,
+        "parallel_tool_use": False,
+        "json_mode": False,
+    },
+    "llama3-8b-8192": {
+        "name": "Llama 3 8B (8K)",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "tool_use": False,
+        "parallel_tool_use": False,
+        "json_mode": False,
+    },
+}
+
+
+def get_model_config(model_id: str) -> Dict[str, Any]:
+    """
+    Get the configuration for a specific model.
+    
+    Args:
+        model_id: The model ID to get configuration for
+        
+    Returns:
+        Dict with model configuration
+        
+    Raises:
+        ValueError: If model_id is not found in configuration
+    """
+    if model_id in GROQ_MODELS_CONFIG:
+        return GROQ_MODELS_CONFIG[model_id]
+    
+    # Try to find a match with different format or case
+    normalized_id = model_id.lower().replace("-", "").replace("_", "")
+    for config_id, config in GROQ_MODELS_CONFIG.items():
+        if normalized_id == config_id.lower().replace("-", "").replace("_", ""):
+            return config
+    
+    # Default configuration for unknown models
+    return {
+        "name": model_id,
+        "context_window": 4096,  # Conservative default
+        "max_completion_tokens": 1024,  # Conservative default
+        "tool_use": False,
+        "parallel_tool_use": False,
+        "json_mode": False,
+    }
+
+
+def get_default_model() -> str:
+    """Get the default model ID for Groq."""
+    return "llama-3.3-70b-versatile"
+
+
+def supports_tool_use(model_id: str) -> bool:
+    """Check if a model supports tool use."""
+    return get_model_config(model_id).get("tool_use", False)
+
+
+def supports_parallel_tool_use(model_id: str) -> bool:
+    """Check if a model supports parallel tool use."""
+    return get_model_config(model_id).get("parallel_tool_use", False)
+
+
+def supports_json_mode(model_id: str) -> bool:
+    """Check if a model supports JSON mode."""
+    return get_model_config(model_id).get("json_mode", False)
+
+
+def get_max_completion_tokens(model_id: str) -> int:
+    """Get the maximum number of completion tokens for a model."""
+    return get_model_config(model_id).get("max_completion_tokens", 1024) 

airtrain/integrations/groq/skills.py

@@ -0,0 +1,201 @@
+from typing import Generator, Optional, Dict, Any, List, Union
+from pydantic import Field, validator
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import GroqCredentials
+from .models_config import get_max_completion_tokens, get_model_config
+from groq import Groq
+
+
+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"
+        ),
+    )
+    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="llama-3.3-70b-versatile", 
+        description="Groq model to use"
+    )
+    max_tokens: int = Field(
+        default=4096, 
+        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"
+    )
+    tools: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description=(
+            "A list of tools the model may use. "
+            "Currently only functions supported."
+        ),
+    )
+    tool_choice: Optional[Union[str, Dict[str, Any]]] = Field(
+        default=None,
+        description=(
+            "Controls which tool is called by the model. "
+            "'none', 'auto', or specific tool."
+        ),
+    )
+    
+    @validator('max_tokens')
+    def validate_max_tokens(cls, v, values):
+        """Validate that max_tokens doesn't exceed the model's limit."""
+        if 'model' in values:
+            model_id = values['model']
+            max_limit = get_max_completion_tokens(model_id)
+            if v > max_limit:
+                return max_limit
+        return v
+
+
+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 from the API"
+    )
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        default=None, description="Tool calls generated by the model"
+    )
+
+
+class GroqChatSkill(Skill[GroqInput, GroqOutput]):
+    """Skill for Groq chat"""
+
+    input_schema = GroqInput
+    output_schema = GroqOutput
+
+    def __init__(self, credentials: Optional[GroqCredentials] = None):
+        super().__init__()
+        self.credentials = credentials or GroqCredentials.from_env()
+        self.client = Groq(api_key=self.credentials.groq_api_key.get_secret_value())
+
+    def _build_messages(self, input_data: GroqInput) -> 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 Groq
+        """
+        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: GroqInput) -> 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:
+            raise ProcessingError(f"Groq streaming failed: {str(e)}")
+
+    def process(self, input_data: GroqInput) -> GroqOutput:
+        """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
+                tool_calls = None  # Tool calls not available in streaming
+            else:
+                messages = self._build_messages(input_data)
+                
+                # Prepare API call parameters
+                api_params = {
+                    "model": input_data.model,
+                    "messages": messages,
+                    "temperature": input_data.temperature,
+                    "max_tokens": input_data.max_tokens,
+                    "stream": False,
+                }
+                
+                # Add tools and tool_choice if provided
+                if input_data.tools:
+                    api_params["tools"] = input_data.tools
+                
+                if input_data.tool_choice:
+                    api_params["tool_choice"] = input_data.tool_choice
+                
+                completion = self.client.chat.completions.create(**api_params)
+                response = completion.choices[0].message.content or ""
+                
+                # Extract usage information
+                usage = {
+                    "total_tokens": completion.usage.total_tokens,
+                    "prompt_tokens": completion.usage.prompt_tokens,
+                    "completion_tokens": completion.usage.completion_tokens,
+                }
+                
+                # Check for tool calls in the response
+                tool_calls = None
+                if (
+                    hasattr(completion.choices[0].message, "tool_calls") 
+                    and completion.choices[0].message.tool_calls
+                ):
+                    tool_calls = [
+                        {
+                            "id": tool_call.id,
+                            "type": tool_call.type,
+                            "function": {
+                                "name": tool_call.function.name,
+                                "arguments": tool_call.function.arguments
+                            }
+                        }
+                        for tool_call in completion.choices[0].message.tool_calls
+                    ]
+
+            return GroqOutput(
+                response=response, 
+                used_model=input_data.model, 
+                usage=usage,
+                tool_calls=tool_calls
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Groq processing failed: {str(e)}")

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=131072, 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,37 @@
+"""OpenAI API integration."""
+
+from .skills import (
+    OpenAIChatSkill,
+    OpenAIInput,
+    OpenAIParserSkill,
+    OpenAIOutput,
+    OpenAIParserInput,
+    OpenAIParserOutput,
+    OpenAIEmbeddingsSkill,
+    OpenAIEmbeddingsInput,
+    OpenAIEmbeddingsOutput,
+)
+from .credentials import OpenAICredentials
+from .list_models import (
+    OpenAIListModelsSkill,
+    OpenAIListModelsInput,
+    OpenAIListModelsOutput,
+    OpenAIModel,
+)
+
+__all__ = [
+    "OpenAIChatSkill",
+    "OpenAIInput",
+    "OpenAIParserSkill",
+    "OpenAIParserInput",
+    "OpenAIParserOutput",
+    "OpenAICredentials",
+    "OpenAIOutput",
+    "OpenAIEmbeddingsSkill",
+    "OpenAIEmbeddingsInput",
+    "OpenAIEmbeddingsOutput",
+    "OpenAIListModelsSkill",
+    "OpenAIListModelsInput",
+    "OpenAIListModelsOutput",
+    "OpenAIModel",
+]

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=131072, 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)