airtrain 0.1.53__py3-none-any.whl → 0.1.58__py3-none-any.whl

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

@@ -1,8 +1,9 @@
-from typing import Generator, Optional, Dict, Any, List
-from pydantic import Field
+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
 
 
@@ -12,22 +13,59 @@ class GroqInput(InputSchema):
     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",
+        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",
+        description=(
+            "List of previous conversation messages in "
+            "[{'role': 'user|assistant', 'content': 'message'}] format"
+        ),
     )
     model: str = Field(
-        default="deepseek-r1-distill-llama-70b-specdec", description="Groq model to use"
+        default="llama-3.3-70b-versatile", 
+        description="Groq model to use"
+    )
+    max_tokens: int = Field(
+        default=4096, 
+        description="Maximum tokens in response"
     )
-    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
+        default=0.7, 
+        description="Temperature for response generation", 
+        ge=0, 
+        le=1
     )
     stream: bool = Field(
-        default=False, description="Whether to stream the response progressively"
+        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):
@@ -38,6 +76,9 @@ class GroqOutput(OutputSchema):
     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]):
@@ -101,24 +142,59 @@ class GroqChatSkill(Skill[GroqInput, GroqOutput]):
                     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)
-                completion = self.client.chat.completions.create(
-                    model=input_data.model,
-                    messages=messages,
-                    temperature=input_data.temperature,
-                    max_tokens=input_data.max_tokens,
-                    stream=False,
-                )
-                response = completion.choices[0].message.content
+                
+                # 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
+                response=response, 
+                used_model=input_data.model, 
+                usage=usage,
+                tool_calls=tool_calls
             )
 
         except Exception as e:

airtrain/integrations/together/__init__.py

@@ -1,7 +1,14 @@
 """Together AI integration module"""
 
 from .credentials import TogetherAICredentials
-from .skills import TogetherAIChatSkill
+from .skills import TogetherAIChatSkill, TogetherAIInput, TogetherAIOutput
+from .models_config import (
+    get_model_config_with_capabilities,
+    get_max_completion_tokens,
+    supports_tool_use,
+    supports_json_mode,
+    TOGETHER_MODELS_CONFIG,
+)
 from .list_models import (
     TogetherListModelsSkill,
     TogetherListModelsInput,
@@ -12,8 +19,15 @@ from .models import TogetherModel
 __all__ = [
     "TogetherAICredentials", 
     "TogetherAIChatSkill",
+    "TogetherAIInput",
+    "TogetherAIOutput",
     "TogetherListModelsSkill",
     "TogetherListModelsInput",
     "TogetherListModelsOutput",
     "TogetherModel",
+    "get_model_config_with_capabilities",
+    "get_max_completion_tokens",
+    "supports_tool_use",
+    "supports_json_mode",
+    "TOGETHER_MODELS_CONFIG",
 ]

airtrain/integrations/together/models_config.py

@@ -1,4 +1,4 @@
-from typing import Dict, NamedTuple
+from typing import Dict, NamedTuple, Any
 
 
 class ModelConfig(NamedTuple):
@@ -272,6 +272,128 @@ def get_default_model() -> str:
     return "meta-llama/Llama-3.3-70B-Instruct-Turbo"
 
 
+# Model configuration with capabilities for each model
+TOGETHER_MODELS_CONFIG = {
+    "meta-llama/Llama-3.1-8B-Instruct": {
+        "name": "Llama 3.1 8B Instruct",
+        "context_window": 128000,
+        "max_completion_tokens": 8192,
+        "tool_use": True,
+        "json_mode": True,
+    },
+    "meta-llama/Llama-3.1-70B-Instruct": {
+        "name": "Llama 3.1 70B Instruct",
+        "context_window": 128000,
+        "max_completion_tokens": 32768,
+        "tool_use": True,
+        "json_mode": True,
+    },
+    "mistralai/Mixtral-8x7B-Instruct-v0.1": {
+        "name": "Mixtral 8x7B Instruct v0.1",
+        "context_window": 32768,
+        "max_completion_tokens": 8192,
+        "tool_use": True,
+        "json_mode": True,
+    },
+    "meta-llama/Meta-Llama-3-8B-Instruct": {
+        "name": "Meta Llama 3 8B Instruct",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "tool_use": True,
+        "json_mode": True,
+    },
+    "meta-llama/Meta-Llama-3-70B-Instruct": {
+        "name": "Meta Llama 3 70B Instruct",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "tool_use": True,
+        "json_mode": True,
+    },
+    "deepseek-ai/DeepSeek-Coder-V2": {
+        "name": "DeepSeek Coder V2",
+        "context_window": 128000,
+        "max_completion_tokens": 16384,
+        "tool_use": True,
+        "json_mode": True,
+    },
+    "deepseek-ai/DeepSeek-V2": {
+        "name": "DeepSeek V2",
+        "context_window": 128000,
+        "max_completion_tokens": 16384,
+        "tool_use": True,
+        "json_mode": True,
+    },
+    "deepseek-ai/DeepSeek-R1": {
+        "name": "DeepSeek R1",
+        "context_window": 32768,
+        "max_completion_tokens": 8192,
+        "tool_use": False,
+        "json_mode": False,
+    },
+    # Qwen models
+    "Qwen/Qwen2.5-72B-Instruct-Turbo": {
+        "context_window": 128000,
+        "max_completion_tokens": 4096,
+        "tool_use": True,
+        "json_mode": True,
+    },
+    "Qwen/Qwen2.5-7B-Instruct": {
+        "context_window": 32768,
+        "max_completion_tokens": 4096,
+        "tool_use": True,
+        "json_mode": True,
+    },
+}
+
+
+def get_model_config_with_capabilities(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 TOGETHER_MODELS_CONFIG:
+        return TOGETHER_MODELS_CONFIG[model_id]
+    
+    # Try to find a match with different format or case
+    normalized_id = model_id.lower().replace("-", "").replace("_", "").replace("/", "")
+    for config_id, config in TOGETHER_MODELS_CONFIG.items():
+        norm_config_id = config_id.lower().replace("-", "").replace("_", "").replace("/", "")
+        if normalized_id == norm_config_id:
+            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,
+        "json_mode": False,
+    }
+
+
+def supports_tool_use(model_id: str) -> bool:
+    """Check if a model supports tool use."""
+    return get_model_config_with_capabilities(model_id).get("tool_use", False)
+
+
+def supports_json_mode(model_id: str) -> bool:
+    """Check if a model supports JSON mode."""
+    return get_model_config_with_capabilities(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_with_capabilities(model_id).get("max_completion_tokens", 1024)
+
+
 if __name__ == "__main__":
     print(len(TOGETHER_MODELS))
     print(get_model_config("meta-llama/Llama-3.3-70B-Instruct-Turbo"))

airtrain/integrations/together/skills.py

@@ -1,9 +1,10 @@
-from typing import Optional, Dict, Any, List, Generator
-from pydantic import Field
+from typing import Optional, Dict, Any, List, Generator, Union
+from pydantic import Field, validator
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
 from .credentials import TogetherAICredentials
 from .models import TogetherAIImageInput, TogetherAIImageOutput, GeneratedImage
+from .models_config import get_max_completion_tokens
 from pathlib import Path
 import base64
 import time
@@ -20,16 +21,53 @@ class TogetherAIInput(InputSchema):
     )
     conversation_history: List[Dict[str, str]] = Field(
         default_factory=list,
-        description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
+        description=(
+            "List of previous conversation messages in "
+            "[{'role': 'user|assistant', 'content': 'message'}] format"
+        ),
     )
     model: str = Field(
-        default="deepseek-ai/DeepSeek-R1", description="Together AI model to use"
+        default="meta-llama/Llama-3.1-8B-Instruct", 
+        description="Together AI model to use"
+    )
+    max_tokens: int = Field(
+        default=4096, 
+        description="Maximum tokens in response"
     )
-    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
+        default=0.7, 
+        description="Temperature for response generation", 
+        ge=0, 
+        le=1
+    )
+    stream: bool = Field(
+        default=False, 
+        description="Whether to stream the response"
+    )
+    tools: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description=(
+            "A list of tools the model may use. "
+            "Currently only functions supported."
+        ),
     )
-    stream: bool = Field(default=False, description="Whether to stream the response")
+    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 TogetherAIOutput(OutputSchema):
@@ -38,6 +76,9 @@ class TogetherAIOutput(OutputSchema):
     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")
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        default=None, description="Tool calls generated by the model"
+    )
 
 
 class TogetherAIChatSkill(Skill[TogetherAIInput, TogetherAIOutput]):
@@ -102,27 +143,83 @@ class TogetherAIChatSkill(Skill[TogetherAIInput, TogetherAIOutput]):
                 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)
-                completion = self.client.chat.completions.create(
-                    model=input_data.model,
-                    messages=messages,
-                    temperature=input_data.temperature,
-                    max_tokens=input_data.max_tokens,
-                    stream=False,
-                )
-                response = completion.choices[0].message.content
+                
+                # 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
+                
+                try:
+                    completion = self.client.chat.completions.create(**api_params)
+                    response = completion.choices[0].message.content or ""
+                    
+                    # Extract usage information
+                    usage = (
+                        completion.usage.model_dump()
+                        if hasattr(completion, "usage")
+                        else {}
+                    )
+                    
+                    # 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
+                        ]
+                except Exception as api_error:
+                    # Provide more specific error messages for common tool call issues
+                    error_message = str(api_error)
+                    if "tool_call_id" in error_message:
+                        raise ProcessingError(
+                            "Tool call error: Missing tool_call_id in conversation history. "
+                            "When adding tool responses to conversation history, "
+                            "make sure each message with role='tool' includes a 'tool_call_id' field."
+                        )
+                    elif "messages" in error_message and "tool" in error_message:
+                        raise ProcessingError(
+                            "Tool call error: Invalid message format in conversation history. "
+                            f"Original error: {error_message}"
+                        )
+                    else:
+                        # Re-raise with original error message
+                        raise ProcessingError(f"Together AI API error: {error_message}")
 
             return TogetherAIOutput(
                 response=response,
                 used_model=input_data.model,
-                usage=(
-                    completion.usage.model_dump()
-                    if hasattr(completion, "usage")
-                    else {}
-                ),
+                usage=usage,
+                tool_calls=tool_calls
             )
 
+        except ProcessingError:
+            # Re-raise ProcessingError without modification
+            raise
         except Exception as e:
             raise ProcessingError(f"Together AI processing failed: {str(e)}")