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

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/sambanova/__init__.py

@@ -0,0 +1,6 @@
+"""Sambanova integration module"""
+
+from .credentials import SambanovaCredentials
+from .skills import SambanovaChatSkill
+
+__all__ = ["SambanovaCredentials", "SambanovaChatSkill"]

airtrain/integrations/sambanova/credentials.py

@@ -0,0 +1,20 @@
+from pydantic import Field, SecretStr, HttpUrl
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+
+
+class SambanovaCredentials(BaseCredentials):
+    """SambaNova credentials"""
+
+    sambanova_api_key: SecretStr = Field(..., description="SambaNova API key")
+    sambanova_endpoint_url: HttpUrl = Field(..., description="SambaNova API endpoint")
+
+    _required_credentials = {"sambanova_api_key", "sambanova_endpoint_url"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate SambaNova credentials"""
+        try:
+            # Implement SambaNova-specific validation
+            # This would depend on their API client implementation
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid SambaNova credentials: {str(e)}")

airtrain/integrations/sambanova/skills.py

@@ -0,0 +1,129 @@
+from typing import Optional, Dict, Any, List, Generator
+from pydantic import Field
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import SambanovaCredentials
+import openai
+
+
+class SambanovaInput(InputSchema):
+    """Schema for Sambanova 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="DeepSeek-R1-Distill-Llama-70B", description="Sambanova model to use"
+    )
+    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
+    )
+    top_p: float = Field(
+        default=0.1, description="Top p sampling parameter", ge=0, le=1
+    )
+    stream: bool = Field(
+        default=False, description="Whether to stream the response progressively"
+    )
+
+
+class SambanovaOutput(OutputSchema):
+    """Schema for Sambanova 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 SambanovaChatSkill(Skill[SambanovaInput, SambanovaOutput]):
+    """Skill for Sambanova chat"""
+
+    input_schema = SambanovaInput
+    output_schema = SambanovaOutput
+
+    def __init__(self, credentials: Optional[SambanovaCredentials] = None):
+        super().__init__()
+        self.credentials = credentials or SambanovaCredentials.from_env()
+        self.client = openai.OpenAI(
+            api_key=self.credentials.sambanova_api_key.get_secret_value(),
+            base_url="https://api.sambanova.ai/v1",
+        )
+
+    def _build_messages(self, input_data: SambanovaInput) -> 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 Sambanova
+        """
+        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: SambanovaInput) -> 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,
+                top_p=input_data.top_p,
+                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"Sambanova streaming failed: {str(e)}")
+
+    def process(self, input_data: SambanovaInput) -> SambanovaOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+                usage = {}  # Usage stats not available in streaming
+            else:
+                messages = self._build_messages(input_data)
+                response = self.client.chat.completions.create(
+                    model=input_data.model,
+                    messages=messages,
+                    temperature=input_data.temperature,
+                    max_tokens=input_data.max_tokens,
+                    top_p=input_data.top_p,
+                )
+                usage = (
+                    response.usage.model_dump() if hasattr(response, "usage") else {}
+                )
+
+            return SambanovaOutput(
+                response=response.choices[0].message.content,
+                used_model=input_data.model,
+                usage=usage,
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Sambanova processing 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

airtrain/integrations/search/exa/schemas.py

@@ -0,0 +1,114 @@
+"""
+Schemas for Exa Search API.
+
+This module defines the input and output schemas for the Exa search API.
+"""
+
+from typing import Dict, List, Optional, Any, Union, bool
+from pydantic import BaseModel, Field, HttpUrl
+
+from airtrain.core.schemas import InputSchema, OutputSchema
+
+
+class ExaContentConfig(BaseModel):
+    """Configuration for the content to be returned by Exa search."""
+
+    text: bool = Field(default=True, description="Whether to return text content.")
+    extractedText: Optional[bool] = Field(
+        default=None, description="Whether to return extracted text content."
+    )
+    embedded: Optional[bool] = Field(
+        default=None, description="Whether to return embedded content."
+    )
+    links: Optional[bool] = Field(
+        default=None, description="Whether to return links from the content."
+    )
+    screenshot: Optional[bool] = Field(
+        default=None, description="Whether to return screenshots of the content."
+    )
+    highlighted: Optional[bool] = Field(
+        default=None, description="Whether to return highlighted text."
+    )
+
+
+class ExaSearchInputSchema(InputSchema):
+    """Input schema for Exa search API."""
+
+    query: str = Field(description="The search query to execute.")
+    numResults: Optional[int] = Field(
+        default=None, description="Number of results to return."
+    )
+    contents: Optional[ExaContentConfig] = Field(
+        default_factory=ExaContentConfig,
+        description="Configuration for the content to be returned.",
+    )
+    highlights: Optional[dict] = Field(
+        default=None, description="Highlighting configuration for search results."
+    )
+    useAutoprompt: Optional[bool] = Field(
+        default=None, description="Whether to use autoprompt for the search."
+    )
+    type: Optional[str] = Field(default=None, description="Type of search to perform.")
+    includeDomains: Optional[List[str]] = Field(
+        default=None, description="List of domains to include in the search."
+    )
+    excludeDomains: Optional[List[str]] = Field(
+        default=None, description="List of domains to exclude from the search."
+    )
+
+
+class ExaModerationConfig(BaseModel):
+    """Moderation configuration returned in search results."""
+
+    llamaguardS1: Optional[bool] = None
+    llamaguardS3: Optional[bool] = None
+    llamaguardS4: Optional[bool] = None
+    llamaguardS12: Optional[bool] = None
+    domainBlacklisted: Optional[bool] = None
+    domainBlacklistedMedia: Optional[bool] = None
+
+
+class ExaHighlight(BaseModel):
+    """Highlight information for a search result."""
+
+    text: str
+    score: float
+
+
+class ExaSearchResult(BaseModel):
+    """Individual search result from Exa."""
+
+    id: str
+    url: str
+    title: Optional[str] = None
+    text: Optional[str] = None
+    extractedText: Optional[str] = None
+    embedded: Optional[Dict[str, Any]] = None
+    score: float
+    published: Optional[str] = None
+    author: Optional[str] = None
+    highlights: Optional[List[ExaHighlight]] = None
+    robotsAllowed: Optional[bool] = None
+    moderationConfig: Optional[ExaModerationConfig] = None
+    urls: Optional[List[str]] = None
+
+
+class ExaCostDetails(BaseModel):
+    """Cost details for an Exa search request."""
+
+    total: float
+    search: Dict[str, float]
+    contents: Dict[str, float]
+
+
+class ExaSearchOutputSchema(OutputSchema):
+    """Output schema for Exa search API."""
+
+    results: List[ExaSearchResult] = Field(description="List of search results.")
+    query: str = Field(description="The original search query.")
+    autopromptString: Optional[str] = Field(
+        default=None, description="Autoprompt string used for the search if enabled."
+    )
+    costDollars: Optional[ExaCostDetails] = Field(
+        default=None, description="Cost details for the search request."
+    )