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

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)}"
+            )

airtrain/integrations/perplexity/list_models.py

@@ -0,0 +1,112 @@
+from typing import Dict, Any, List, Optional
+import requests
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.integrations.combined.list_models_factory import (
+    BaseListModelsSkill,
+    GenericListModelsInput,
+    GenericListModelsOutput,
+)
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import PerplexityCredentials
+from .models_config import PERPLEXITY_MODELS_CONFIG
+
+
+class PerplexityListModelsInput(InputSchema):
+    """Schema for listing Perplexity AI models"""
+
+    api_models_only: bool = False
+
+
+class PerplexityListModelsOutput(OutputSchema):
+    """Schema for Perplexity AI models listing output"""
+
+    models: List[Dict[str, Any]]
+    provider: str = "perplexity"
+
+
+class PerplexityListModelsSkill(BaseListModelsSkill):
+    """Skill for listing Perplexity AI models"""
+
+    def __init__(self, credentials: Optional[PerplexityCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__(provider="perplexity", credentials=credentials)
+        self.credentials = credentials
+
+    def get_models(self) -> List[Dict[str, Any]]:
+        """Return list of Perplexity AI models."""
+        models = []
+
+        # Add models from the configuration
+        for model_id, config in PERPLEXITY_MODELS_CONFIG.items():
+            models.append(
+                {
+                    "id": model_id,
+                    "display_name": config["name"],
+                    "description": config.get("description", ""),
+                    "category": config.get("category", "unknown"),
+                    "capabilities": {
+                        "citations": config.get("citations", False),
+                        "search": config.get("search", False),
+                        "context_window": config.get("context_window", 8192),
+                        "max_completion_tokens": config.get(
+                            "max_completion_tokens", 4096
+                        ),
+                    },
+                }
+            )
+
+        return 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="perplexity")
+        except Exception as e:
+            raise ProcessingError(f"Failed to list Perplexity AI models: {str(e)}")
+
+
+# Standalone version directly using the Perplexity-specific schemas
+class StandalonePerplexityListModelsSkill(
+    Skill[PerplexityListModelsInput, PerplexityListModelsOutput]
+):
+    """Standalone skill for listing Perplexity AI models"""
+
+    input_schema = PerplexityListModelsInput
+    output_schema = PerplexityListModelsOutput
+
+    def __init__(self, credentials: Optional[PerplexityCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials
+
+    def process(
+        self, input_data: PerplexityListModelsInput
+    ) -> PerplexityListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            models = []
+
+            # Add models from the configuration
+            for model_id, config in PERPLEXITY_MODELS_CONFIG.items():
+                models.append(
+                    {
+                        "id": model_id,
+                        "display_name": config["name"],
+                        "description": config.get("description", ""),
+                        "category": config.get("category", "unknown"),
+                        "capabilities": {
+                            "citations": config.get("citations", False),
+                            "search": config.get("search", False),
+                            "context_window": config.get("context_window", 8192),
+                            "max_completion_tokens": config.get(
+                                "max_completion_tokens", 4096
+                            ),
+                        },
+                    }
+                )
+
+            return PerplexityListModelsOutput(models=models, provider="perplexity")
+        except Exception as e:
+            raise ProcessingError(f"Failed to list Perplexity AI models: {str(e)}")

airtrain/integrations/perplexity/models_config.py

@@ -0,0 +1,128 @@
+"""Configuration of Perplexity AI model capabilities."""
+
+from typing import Dict, Any
+
+
+# Model configuration with capabilities for each Perplexity AI model
+PERPLEXITY_MODELS_CONFIG = {
+    # Search Models
+    "sonar-pro": {
+        "name": "Sonar Pro",
+        "description": "Advanced search offering with grounding, supporting complex queries and follow-ups.",
+        "category": "search",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "citations": True,
+        "search": True,
+    },
+    "sonar": {
+        "name": "Sonar",
+        "description": "Lightweight, cost-effective search model with grounding.",
+        "category": "search",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "citations": True,
+        "search": True,
+    },
+    # Research Models
+    "sonar-deep-research": {
+        "name": "Sonar Deep Research",
+        "description": "Expert-level research model conducting exhaustive searches and generating comprehensive reports.",
+        "category": "research",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "citations": True,
+        "search": True,
+    },
+    # Reasoning Models
+    "sonar-reasoning-pro": {
+        "name": "Sonar Reasoning Pro",
+        "description": "Premier reasoning offering powered by DeepSeek R1 with Chain of Thought (CoT).",
+        "category": "reasoning",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "citations": True,
+        "search": True,
+        "chain_of_thought": True,
+    },
+    "sonar-reasoning": {
+        "name": "Sonar Reasoning",
+        "description": "Fast, real-time reasoning model designed for quick problem-solving with search.",
+        "category": "reasoning",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "citations": True,
+        "search": True,
+        "chain_of_thought": True,
+    },
+    # Offline Models
+    "r1-1776": {
+        "name": "R1-1776",
+        "description": "A version of DeepSeek R1 post-trained for uncensored, unbiased, and factual information.",
+        "category": "offline",
+        "context_window": 8192,
+        "max_completion_tokens": 4096,
+        "citations": False,
+        "search": False,
+    },
+}
+
+
+def get_model_config(model_id: str) -> Dict[str, Any]:
+    """
+    Get the configuration for a specific Perplexity AI 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 PERPLEXITY_MODELS_CONFIG:
+        return PERPLEXITY_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 PERPLEXITY_MODELS_CONFIG.items():
+        if normalized_id == config_id.lower().replace("-", "").replace("_", ""):
+            return config
+
+    # If model not found, raise an error
+    raise ValueError(
+        f"Model '{model_id}' not found in Perplexity AI models configuration"
+    )
+
+
+def get_default_model() -> str:
+    """Get the default model ID for Perplexity AI."""
+    return "sonar-pro"
+
+
+def supports_citations(model_id: str) -> bool:
+    """Check if a model supports citations."""
+    return get_model_config(model_id).get("citations", False)
+
+
+def supports_search(model_id: str) -> bool:
+    """Check if a model uses search capabilities."""
+    return get_model_config(model_id).get("search", False)
+
+
+def get_models_by_category(category: str) -> Dict[str, Dict[str, Any]]:
+    """
+    Get all models belonging to a specific category.
+
+    Args:
+        category: Category to filter by ('search', 'research', 'reasoning', 'offline')
+
+    Returns:
+        Dict of model IDs and their configurations that match the category
+    """
+    return {
+        model_id: config
+        for model_id, config in PERPLEXITY_MODELS_CONFIG.items()
+        if config.get("category") == category
+    }

airtrain/integrations/perplexity/skills.py

@@ -0,0 +1,279 @@
+from typing import Dict, Any, List, Optional, Generator, Union
+from pydantic import Field, validator
+import requests
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import PerplexityCredentials
+from .models_config import get_model_config, get_default_model
+
+
+class PerplexityInput(InputSchema):
+    """Schema for Perplexity AI chat input"""
+
+    user_input: str = Field(..., description="User's input text")
+    system_prompt: Optional[str] = Field(
+        default=None,
+        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="sonar-pro",
+        description="Perplexity AI model to use",
+    )
+    temperature: Optional[float] = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+    max_tokens: Optional[int] = Field(
+        default=500, description="Maximum tokens in response"
+    )
+    top_p: Optional[float] = Field(
+        default=1.0, description="Top-p (nucleus) sampling parameter", ge=0, le=1
+    )
+    top_k: Optional[int] = Field(
+        default=None,
+        description="Top-k sampling parameter",
+    )
+    presence_penalty: Optional[float] = Field(
+        default=None,
+        description="Presence penalty parameter",
+    )
+    frequency_penalty: Optional[float] = Field(
+        default=None,
+        description="Frequency penalty parameter",
+    )
+
+    @validator("model")
+    def validate_model(cls, v):
+        """Validate that the model is supported by Perplexity AI."""
+        try:
+            get_model_config(v)
+            return v
+        except ValueError as e:
+            raise ValueError(f"Invalid Perplexity AI model: {v}. {str(e)}")
+
+
+class PerplexityCitation(OutputSchema):
+    """Schema for Perplexity AI citation information"""
+
+    url: str = Field(..., description="URL of the citation source")
+    title: Optional[str] = Field(None, description="Title of the cited source")
+    snippet: Optional[str] = Field(None, description="Text snippet from the citation")
+
+
+class PerplexityOutput(OutputSchema):
+    """Schema for Perplexity AI chat output"""
+
+    response: str = Field(..., description="Model's response text")
+    used_model: str = Field(..., description="Model used for generation")
+    usage: Dict[str, int] = Field(..., description="Usage statistics from the API")
+    citations: Optional[List[PerplexityCitation]] = Field(
+        default=None, description="Citations used in the response, if available"
+    )
+    search_queries: Optional[List[str]] = Field(
+        default=None, description="Search queries used, if available"
+    )
+
+
+class PerplexityChatSkill(Skill[PerplexityInput, PerplexityOutput]):
+    """Skill for interacting with Perplexity AI models"""
+
+    input_schema = PerplexityInput
+    output_schema = PerplexityOutput
+
+    def __init__(self, credentials: Optional[PerplexityCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or PerplexityCredentials.from_env()
+        self.api_url = "https://api.perplexity.ai/chat/completions"
+
+    def _build_messages(self, input_data: PerplexityInput) -> List[Dict[str, str]]:
+        """Build messages list from input data including conversation history."""
+        messages = []
+
+        # Add system prompt if provided
+        if input_data.system_prompt:
+            messages.append({"role": "system", "content": input_data.system_prompt})
+
+        # Add conversation history
+        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 _prepare_api_parameters(self, input_data: PerplexityInput) -> Dict[str, Any]:
+        """Prepare parameters for the API request."""
+        parameters = {
+            "model": input_data.model,
+            "messages": self._build_messages(input_data),
+            "max_tokens": input_data.max_tokens,
+        }
+
+        # Add optional parameters if provided
+        if input_data.temperature is not None:
+            parameters["temperature"] = input_data.temperature
+
+        if input_data.top_p is not None:
+            parameters["top_p"] = input_data.top_p
+
+        if input_data.top_k is not None:
+            parameters["top_k"] = input_data.top_k
+
+        if input_data.presence_penalty is not None:
+            parameters["presence_penalty"] = input_data.presence_penalty
+
+        if input_data.frequency_penalty is not None:
+            parameters["frequency_penalty"] = input_data.frequency_penalty
+
+        return parameters
+
+    def process(self, input_data: PerplexityInput) -> PerplexityOutput:
+        """Process the input and return the complete response."""
+        try:
+            # Prepare headers with API key
+            headers = {
+                "Authorization": f"Bearer {self.credentials.perplexity_api_key.get_secret_value()}",
+                "Content-Type": "application/json",
+            }
+
+            # Prepare parameters for the API request
+            data = self._prepare_api_parameters(input_data)
+
+            # Make the API request
+            response = requests.post(self.api_url, headers=headers, json=data)
+
+            # Check if request was successful
+            if response.status_code != 200:
+                raise ProcessingError(
+                    f"Perplexity AI API error: {response.status_code} - {response.text}"
+                )
+
+            # Parse the response
+            result = response.json()
+
+            # Extract content from the completion
+            content = result["choices"][0]["message"]["content"]
+
+            # Extract and process citations if available
+            citations = None
+            if "citations" in result:
+                citations = [
+                    PerplexityCitation(
+                        url=citation.get("url", ""),
+                        title=citation.get("title"),
+                        snippet=citation.get("snippet"),
+                    )
+                    for citation in result.get("citations", [])
+                ]
+
+            # Extract search queries if available
+            search_queries = None
+            if "usage" in result and "num_search_queries" in result["usage"]:
+                search_queries = result.get("search_queries", [])
+
+            # Create and return output
+            return PerplexityOutput(
+                response=content,
+                used_model=input_data.model,
+                usage=result.get("usage", {}),
+                citations=citations,
+                search_queries=search_queries,
+            )
+
+        except Exception as e:
+            if isinstance(e, ProcessingError):
+                raise e
+            raise ProcessingError(f"Perplexity AI processing failed: {str(e)}")
+
+
+class PerplexityProcessStreamError(Exception):
+    """Error raised during stream processing"""
+
+    pass
+
+
+class PerplexityStreamOutput(OutputSchema):
+    """Schema for streaming output tokens"""
+
+    token: str = Field(..., description="Text token")
+    finish_reason: Optional[str] = Field(
+        None, description="Why the completion finished"
+    )
+
+
+class PerplexityStreamingChatSkill(PerplexityChatSkill):
+    """Extension of PerplexityChatSkill that supports streaming responses"""
+
+    def process_stream(
+        self, input_data: PerplexityInput
+    ) -> Generator[PerplexityStreamOutput, None, None]:
+        """
+        Process the input and stream the response tokens.
+
+        Note: Perplexity AI API may not support true streaming. In that case, this
+        method will make a regular API call and yield the entire response at once.
+        """
+        try:
+            # Prepare headers with API key
+            headers = {
+                "Authorization": f"Bearer {self.credentials.perplexity_api_key.get_secret_value()}",
+                "Content-Type": "application/json",
+            }
+
+            # Prepare parameters for the API request, including stream=true if possible
+            data = self._prepare_api_parameters(input_data)
+            data["stream"] = True
+
+            # Make the API request
+            response = requests.post(
+                self.api_url, headers=headers, json=data, stream=True
+            )
+
+            # Check if request was successful
+            if response.status_code != 200:
+                raise PerplexityProcessStreamError(
+                    f"Perplexity AI API error: {response.status_code} - {response.text}"
+                )
+
+            # Process the streaming response if supported
+            for line in response.iter_lines():
+                if line:
+                    # Parse the response line
+                    try:
+                        # Remove 'data: ' prefix if present
+                        if line.startswith(b"data: "):
+                            line = line[6:]
+
+                        # Parse JSON
+                        import json
+
+                        chunk = json.loads(line)
+
+                        # Extract content
+                        if "choices" in chunk and len(chunk["choices"]) > 0:
+                            choice = chunk["choices"][0]
+                            if "delta" in choice and "content" in choice["delta"]:
+                                content = choice["delta"]["content"]
+                                if content:
+                                    yield PerplexityStreamOutput(
+                                        token=content,
+                                        finish_reason=choice.get("finish_reason"),
+                                    )
+                    except json.JSONDecodeError:
+                        # Skip non-JSON lines
+                        continue
+                    except Exception as e:
+                        raise PerplexityProcessStreamError(
+                            f"Error processing stream chunk: {str(e)}"
+                        )
+
+        except Exception as e:
+            if isinstance(e, PerplexityProcessStreamError):
+                raise ProcessingError(str(e))
+            raise ProcessingError(f"Perplexity AI streaming failed: {str(e)}")

airtrain/integrations/search/__init__.py

@@ -0,0 +1,21 @@
+"""
+Search integrations for AirTrain.
+
+This package provides integrations with various search providers.
+"""
+
+# Import specific search integrations as needed
+from .exa import (
+    ExaCredentials,
+    ExaSearchInputSchema,
+    ExaSearchOutputSchema,
+    ExaSearchSkill,
+)
+
+__all__ = [
+    # Exa Search
+    "ExaCredentials",
+    "ExaSearchInputSchema",
+    "ExaSearchOutputSchema",
+    "ExaSearchSkill",
+]

airtrain/integrations/search/exa/__init__.py

@@ -0,0 +1,23 @@
+"""
+Exa Search API integration.
+
+This module provides integration with the Exa search API for web searching capabilities.
+"""
+
+from .credentials import ExaCredentials
+from .schemas import (
+    ExaSearchInputSchema,
+    ExaSearchOutputSchema,
+    ExaContentConfig,
+    ExaSearchResult,
+)
+from .skills import ExaSearchSkill
+
+__all__ = [
+    "ExaCredentials",
+    "ExaSearchInputSchema",
+    "ExaSearchOutputSchema",
+    "ExaContentConfig",
+    "ExaSearchResult",
+    "ExaSearchSkill",
+]

airtrain/integrations/search/exa/credentials.py

@@ -0,0 +1,30 @@
+"""
+Credentials for Exa Search API.
+
+This module provides credential management for the Exa search API.
+"""
+
+from typing import Optional
+from pydantic import Field, SecretStr
+
+from airtrain.core.credentials import BaseCredentials
+
+
+class ExaCredentials(BaseCredentials):
+    """Credentials for accessing the Exa search API."""
+
+    exa_api_key: SecretStr = Field(
+        description="Exa search API key",
+    )
+
+    _required_credentials = {"exa_api_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate that the required credentials are present and valid."""
+        # First check that required credentials are present
+        await super().validate_credentials()
+
+        # In a production environment, we might want to make a test API call here
+        # to verify the API key is actually valid, but for now we'll just check
+        # that it's present
+        return True