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

airtrain/__init__.py

@@ -1,43 +1,65 @@
-"""Airtrain - A platform for building and deploying AI agents with structured skills"""
+"""
+Airtrain: AI Agent Framework
 
-__version__ = "0.1.58"
+This library provides a flexible framework for building AI agents
+that can complete complex tasks using AI models, skills, and tools.
+"""
+
+__version__ = "0.1.61"
 
 import sys
 
 # Core imports
-from .core.skills import Skill, ProcessingError
-from .core.schemas import InputSchema, OutputSchema
-from .core.credentials import BaseCredentials
+from .core import Skill, ProcessingError, InputSchema, OutputSchema, BaseCredentials
 
 # Integration imports - Credentials
-from .integrations.openai.credentials import OpenAICredentials
-from .integrations.aws.credentials import AWSCredentials
-from .integrations.google.credentials import GoogleCloudCredentials
-from .integrations.anthropic.credentials import AnthropicCredentials
-from .integrations.groq.credentials import GroqCredentials
-from .integrations.together.credentials import TogetherAICredentials
-from .integrations.ollama.credentials import OllamaCredentials
-from .integrations.sambanova.credentials import SambanovaCredentials
-from .integrations.cerebras.credentials import CerebrasCredentials
+from .integrations import (
+    # OpenAI
+    OpenAICredentials,
+    OpenAIChatSkill,
+    OpenAICompletionSkill,
+    # Anthropic
+    AnthropicCredentials,
+    AnthropicChatSkill,
+    # Together.ai
+    TogetherAICredentials,
+    TogetherChatSkill,
+    # Fireworks
+    FireworksCredentials,
+    FireworksChatSkill,
+    # Google
+    GeminiCredentials,
+    GeminiChatSkill,
+    # Search
+    ExaCredentials,
+    ExaSearchSkill,
+    ExaSearchInputSchema,
+    ExaSearchOutputSchema,
+)
 
 # Integration imports - Skills
-from .integrations.openai.skills import OpenAIChatSkill, OpenAIParserSkill
-from .integrations.anthropic.skills import AnthropicChatSkill
 from .integrations.aws.skills import AWSBedrockSkill
 from .integrations.google.skills import GoogleChatSkill
 from .integrations.groq.skills import GroqChatSkill
-from .integrations.together.skills import TogetherAIChatSkill
 from .integrations.ollama.skills import OllamaChatSkill
 from .integrations.sambanova.skills import SambanovaChatSkill
 from .integrations.cerebras.skills import CerebrasChatSkill
 
 # Tool imports
 from .tools import (
-    StatefulTool,
-    StatelessTool,
-    register_tool,
     ToolFactory,
-    execute_tool_call
+    register_tool,
+    StatelessTool,
+    StatefulTool,
+    BaseTool,
+    ListDirectoryTool,
+    DirectoryTreeTool,
+    ApiCallTool,
+    ExecuteCommandTool,
+    FindFilesTool,
+    TerminalNavigationTool,
+    SearchTermTool,
+    RunPytestTool,
 )
 
 # Agent imports
@@ -48,7 +70,7 @@ from .agents import (
     BaseMemory,
     ShortTermMemory,
     LongTermMemory,
-    SharedMemory
+    SharedMemory,
 )
 
 # Telemetry import - must be imported after version is defined
@@ -63,33 +85,41 @@ __all__ = [
     "InputSchema",
     "OutputSchema",
     "BaseCredentials",
-    # Credentials   
+    # OpenAI Integration
     "OpenAICredentials",
-    "AWSCredentials",
-    "GoogleCloudCredentials",
-    "AnthropicCredentials",
-    "GroqCredentials",
-    "TogetherAICredentials",
-    "OllamaCredentials",
-    "SambanovaCredentials",
-    "CerebrasCredentials",
-    # Skills
     "OpenAIChatSkill",
-    "OpenAIParserSkill",
+    "OpenAICompletionSkill",
+    # Anthropic Integration
+    "AnthropicCredentials",
     "AnthropicChatSkill",
-    "AWSBedrockSkill",
-    "GoogleChatSkill",
-    "GroqChatSkill",
-    "TogetherAIChatSkill",
-    "OllamaChatSkill",
-    "SambanovaChatSkill",
-    "CerebrasChatSkill",
+    # Together Integration
+    "TogetherAICredentials",
+    "TogetherChatSkill",
+    # Fireworks Integration
+    "FireworksCredentials",
+    "FireworksChatSkill",
+    # Google Integration
+    "GeminiCredentials",
+    "GeminiChatSkill",
+    # Search Integration
+    "ExaCredentials",
+    "ExaSearchSkill",
+    "ExaSearchInputSchema",
+    "ExaSearchOutputSchema",
     # Tools
-    "StatefulTool",
-    "StatelessTool",
-    "register_tool",
     "ToolFactory",
-    "execute_tool_call",
+    "register_tool",
+    "StatelessTool",
+    "StatefulTool",
+    "BaseTool",
+    "ListDirectoryTool",
+    "DirectoryTreeTool",
+    "ApiCallTool",
+    "ExecuteCommandTool",
+    "FindFilesTool",
+    "TerminalNavigationTool",
+    "SearchTermTool",
+    "RunPytestTool",
     # Agents
     "BaseAgent",
     "AgentFactory",
@@ -111,7 +141,7 @@ try:
                 f"{sys.version_info.major}."
                 f"{sys.version_info.minor}."
                 f"{sys.version_info.micro}"
-            )
+            ),
         )
     )
 except Exception:

airtrain/core/credentials.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, Union
 import os
 import json
 from pathlib import Path
@@ -53,30 +53,70 @@ class BaseCredentials(BaseModel):
         return cls(**field_values)
 
     @classmethod
-    def from_file(cls, file_path: Path) -> "BaseCredentials":
-        """Load credentials from a file (supports .env, .json, .yaml)"""
+    def from_file(cls, file_path: str | Path) -> "BaseCredentials":
+        """Load credentials from a file (supports .env, .json, .yaml).
+
+        Args:
+            file_path: Path to load credentials from. Can be a string or Path object.
+                Supported formats: .env, .json, .yaml/.yml
+
+        Returns:
+            BaseCredentials: Initialized credentials object
+
+        Raises:
+            FileNotFoundError: If the credentials file does not exist
+            ValueError: If the file format is not supported
+        """
+        # Convert to Path object if string
+        if isinstance(file_path, str):
+            file_path = Path(file_path)
+
         if not file_path.exists():
             raise FileNotFoundError(f"Credentials file not found: {file_path}")
 
-        if file_path.suffix == ".env":
+        # Get file extension, default to .env if none provided
+        suffix = file_path.suffix
+        if not suffix:
+            # Try to find a file with the same name but different extension
+            for ext in [".env", ".json", ".yaml", ".yml"]:
+                potential_path = file_path.with_suffix(ext)
+                if potential_path.exists():
+                    file_path = potential_path
+                    suffix = ext
+                    break
+            # If no file was found, default to .env
+            if not suffix:
+                file_path = file_path.with_suffix(".env")
+                suffix = ".env"
+
+        if suffix == ".env":
             dotenv.load_dotenv(file_path)
             return cls.from_env()
 
-        elif file_path.suffix == ".json":
+        elif suffix == ".json":
             with open(file_path) as f:
                 data = json.load(f)
             return cls(**data)
 
-        elif file_path.suffix in {".yaml", ".yml"}:
+        elif suffix in {".yaml", ".yml"}:
             with open(file_path) as f:
                 data = yaml.safe_load(f)
             return cls(**data)
 
         else:
-            raise ValueError(f"Unsupported file format: {file_path.suffix}")
+            raise ValueError(f"Unsupported file format: {suffix}")
+
+    def save_to_file(self, file_path: str | Path) -> None:
+        """Save credentials to a file.
+
+        Args:
+            file_path: Path to save credentials to. Can be a string or Path object.
+                Supported formats: .env, .json, .yaml/.yml
+        """
+        # Convert to Path object if string
+        if isinstance(file_path, str):
+            file_path = Path(file_path)
 
-    def save_to_file(self, file_path: Path) -> None:
-        """Save credentials to a file"""
         data = self.model_dump(exclude={"_loaded"})
 
         # Convert SecretStr to plain strings for saving
@@ -84,21 +124,27 @@ class BaseCredentials(BaseModel):
             if isinstance(value, SecretStr):
                 data[key] = value.get_secret_value()
 
-        if file_path.suffix == ".env":
+        # Get file extension, default to .env if none provided
+        suffix = file_path.suffix
+        if not suffix:
+            file_path = file_path.with_suffix(".env")
+            suffix = ".env"
+
+        if suffix == ".env":
             with open(file_path, "w") as f:
                 for key, value in data.items():
                     f.write(f"{key.upper()}={value}\n")
 
-        elif file_path.suffix == ".json":
+        elif suffix == ".json":
             with open(file_path, "w") as f:
                 json.dump(data, f, indent=2)
 
-        elif file_path.suffix in {".yaml", ".yml"}:
+        elif suffix in {".yaml", ".yml"}:
             with open(file_path, "w") as f:
                 yaml.dump(data, f)
 
         else:
-            raise ValueError(f"Unsupported file format: {file_path.suffix}")
+            raise ValueError(f"Unsupported file format: {suffix}")
 
     async def validate_credentials(self) -> bool:
         """Validate that all required credentials are present"""

airtrain/integrations/__init__.py

@@ -10,6 +10,7 @@ from .together.credentials import TogetherAICredentials
 from .ollama.credentials import OllamaCredentials
 from .sambanova.credentials import SambanovaCredentials
 from .cerebras.credentials import CerebrasCredentials
+from .perplexity.credentials import PerplexityCredentials
 
 # Skills imports
 from .openai.skills import OpenAIChatSkill, OpenAIParserSkill
@@ -21,16 +22,18 @@ from .together.skills import TogetherAIChatSkill
 from .ollama.skills import OllamaChatSkill
 from .sambanova.skills import SambanovaChatSkill
 from .cerebras.skills import CerebrasChatSkill
+from .perplexity.skills import PerplexityChatSkill, PerplexityStreamingChatSkill
 
 # Model configurations
 from .openai.models_config import OPENAI_MODELS, OpenAIModelConfig
 from .anthropic.models_config import ANTHROPIC_MODELS, AnthropicModelConfig
+from .perplexity.models_config import PERPLEXITY_MODELS_CONFIG
 
 # Combined modules
 from .combined.list_models_factory import (
     ListModelsSkillFactory,
     GenericListModelsInput,
-    GenericListModelsOutput
+    GenericListModelsOutput,
 )
 
 __all__ = [
@@ -44,6 +47,7 @@ __all__ = [
     "OllamaCredentials",
     "SambanovaCredentials",
     "CerebrasCredentials",
+    "PerplexityCredentials",
     # Skills
     "OpenAIChatSkill",
     "OpenAIParserSkill",
@@ -55,11 +59,14 @@ __all__ = [
     "OllamaChatSkill",
     "SambanovaChatSkill",
     "CerebrasChatSkill",
+    "PerplexityChatSkill",
+    "PerplexityStreamingChatSkill",
     # Model configurations
     "OPENAI_MODELS",
     "OpenAIModelConfig",
     "ANTHROPIC_MODELS",
     "AnthropicModelConfig",
+    "PERPLEXITY_MODELS_CONFIG",
     # Combined modules
     "ListModelsSkillFactory",
     "GenericListModelsInput",

airtrain/integrations/combined/list_models_factory.py

@@ -15,16 +15,21 @@ from airtrain.integrations.fireworks.list_models import FireworksListModelsSkill
 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:
@@ -32,38 +37,33 @@ class GenericListModelsInput(InputSchema):
         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"
+        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:
@@ -76,24 +76,42 @@ class BaseListModelsSkill(Skill[GenericListModelsInput, GenericListModelsOutput]
 # 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": "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-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)"},
+            {
+                "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
 
@@ -101,18 +119,27 @@ class GroqListModelsSkill(BaseListModelsSkill):
 # 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"}
+            {
+                "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
 
@@ -120,17 +147,17 @@ class CerebrasListModelsSkill(BaseListModelsSkill):
 # 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"}
+            {"id": "sambanova/samba-2", "display_name": "Samba-2"},
         ]
         return models
 
@@ -138,7 +165,7 @@ class SambanovaListModelsSkill(BaseListModelsSkill):
 # 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,
@@ -147,36 +174,37 @@ class ListModelsSkillFactory:
         "fireworks": FireworksListModelsSkill,
         "groq": GroqListModelsSkill,
         "cerebras": CerebrasListModelsSkill,
-        "sambanova": SambanovaListModelsSkill
+        "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()) 
+        return list(cls._PROVIDER_MAP.keys())

airtrain/integrations/perplexity/__init__.py

@@ -0,0 +1,49 @@
+"""Perplexity AI integration module"""
+
+from .credentials import PerplexityCredentials
+from .skills import (
+    PerplexityInput,
+    PerplexityOutput,
+    PerplexityChatSkill,
+    PerplexityCitation,
+    PerplexityStreamingChatSkill,
+    PerplexityStreamOutput,
+)
+from .list_models import (
+    PerplexityListModelsSkill,
+    StandalonePerplexityListModelsSkill,
+    PerplexityListModelsInput,
+    PerplexityListModelsOutput,
+)
+from .models_config import (
+    get_model_config,
+    get_default_model,
+    supports_citations,
+    supports_search,
+    get_models_by_category,
+    PERPLEXITY_MODELS_CONFIG,
+)
+
+__all__ = [
+    # Credentials
+    "PerplexityCredentials",
+    # Skills
+    "PerplexityInput",
+    "PerplexityOutput",
+    "PerplexityChatSkill",
+    "PerplexityCitation",
+    "PerplexityStreamingChatSkill",
+    "PerplexityStreamOutput",
+    # List Models
+    "PerplexityListModelsSkill",
+    "StandalonePerplexityListModelsSkill",
+    "PerplexityListModelsInput",
+    "PerplexityListModelsOutput",
+    # Model Config
+    "get_model_config",
+    "get_default_model",
+    "supports_citations",
+    "supports_search",
+    "get_models_by_category",
+    "PERPLEXITY_MODELS_CONFIG",
+]

airtrain/integrations/perplexity/credentials.py

@@ -0,0 +1,43 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+import requests
+
+
+class PerplexityCredentials(BaseCredentials):
+    """Perplexity AI API credentials"""
+
+    perplexity_api_key: SecretStr = Field(..., description="Perplexity AI API key")
+
+    _required_credentials = {"perplexity_api_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Perplexity AI credentials by making a test API call"""
+        try:
+            headers = {
+                "Authorization": f"Bearer {self.perplexity_api_key.get_secret_value()}",
+                "Content-Type": "application/json",
+            }
+
+            # Small API call to check if credentials are valid
+            data = {
+                "model": "sonar-pro",
+                "messages": [{"role": "user", "content": "Test"}],
+                "max_tokens": 1,
+            }
+
+            # Make a synchronous request for validation
+            response = requests.post(
+                "https://api.perplexity.ai/chat/completions", headers=headers, json=data
+            )
+
+            if response.status_code == 200:
+                return True
+            else:
+                raise CredentialValidationError(
+                    f"Invalid Perplexity AI credentials: {response.status_code} - {response.text}"
+                )
+
+        except Exception as e:
+            raise CredentialValidationError(
+                f"Invalid Perplexity AI credentials: {str(e)}"
+            )