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

airtrain/integrations/openai/skills.py

@@ -0,0 +1,342 @@
+from typing import AsyncGenerator, List, Optional, Dict, TypeVar, Type, Generator, Union
+from pydantic import Field, BaseModel
+from openai import OpenAI, AsyncOpenAI
+from openai.types.chat import ChatCompletionChunk
+import numpy as np
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import OpenAICredentials
+
+
+class OpenAIInput(InputSchema):
+    """Schema for OpenAI chat input"""
+
+    user_input: str = Field(..., description="User's input text")
+    system_prompt: str = Field(
+        default="You are a helpful assistant.",
+        description="System prompt to guide the model's behavior",
+    )
+    conversation_history: List[Dict[str, str]] = Field(
+        default_factory=list,
+        description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
+    )
+    model: str = Field(
+        default="gpt-4o",
+        description="OpenAI model to use",
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+    max_tokens: Optional[int] = Field(
+        default=131072, description="Maximum tokens in response"
+    )
+    stream: bool = Field(
+        default=False,
+        description="Whether to stream the response token by token",
+    )
+
+
+class OpenAIOutput(OutputSchema):
+    """Schema for OpenAI chat output"""
+
+    response: str
+    used_model: str
+    usage: Dict[str, int]
+
+
+class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
+    """Skill for interacting with OpenAI models with async support"""
+
+    input_schema = OpenAIInput
+    output_schema = OpenAIOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+        self.async_client = AsyncOpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def _build_messages(self, input_data: OpenAIInput) -> List[Dict[str, str]]:
+        """Build messages list from input data including conversation history."""
+        messages = [{"role": "system", "content": input_data.system_prompt}]
+
+        if input_data.conversation_history:
+            messages.extend(input_data.conversation_history)
+
+        messages.append({"role": "user", "content": input_data.user_input})
+        return messages
+
+    def process_stream(self, input_data: OpenAIInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
+        try:
+            messages = self._build_messages(input_data)
+
+            stream = self.client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+                stream=True,
+            )
+
+            for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
+
+        except Exception as e:
+            raise ProcessingError(f"OpenAI streaming failed: {str(e)}")
+
+    def process(self, input_data: OpenAIInput) -> OpenAIOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                # For streaming, collect the entire response
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+            else:
+                # For non-streaming, use regular completion
+                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
+
+            return OpenAIOutput(
+                response=response,
+                used_model=input_data.model,
+                usage={
+                    "total_tokens": completion.usage.total_tokens,
+                    "prompt_tokens": completion.usage.prompt_tokens,
+                    "completion_tokens": completion.usage.completion_tokens,
+                },
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"OpenAI chat failed: {str(e)}")
+
+    async def process_async(self, input_data: OpenAIInput) -> OpenAIOutput:
+        """Async version of process method"""
+        try:
+            messages = self._build_messages(input_data)
+            completion = await self.async_client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+            return OpenAIOutput(
+                response=completion.choices[0].message.content,
+                used_model=completion.model,
+                usage={
+                    "total_tokens": completion.usage.total_tokens,
+                    "prompt_tokens": completion.usage.prompt_tokens,
+                    "completion_tokens": completion.usage.completion_tokens,
+                },
+            )
+        except Exception as e:
+            raise ProcessingError(f"OpenAI async chat failed: {str(e)}")
+
+    async def process_stream_async(
+        self, input_data: OpenAIInput
+    ) -> AsyncGenerator[str, None]:
+        """Async version of stream processor"""
+        try:
+            messages = self._build_messages(input_data)
+            stream = await self.async_client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+                stream=True,
+            )
+            async 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"OpenAI async streaming failed: {str(e)}")
+
+
+ResponseT = TypeVar("ResponseT", bound=BaseModel)
+
+
+class OpenAIParserInput(InputSchema):
+    """Schema for OpenAI structured output input"""
+
+    user_input: str
+    system_prompt: str = "You are a helpful assistant that provides structured data."
+    model: str = "gpt-4o"
+    temperature: float = 0.7
+    max_tokens: Optional[int] = None
+    response_model: Type[ResponseT]
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class OpenAIParserOutput(OutputSchema):
+    """Schema for OpenAI structured output"""
+
+    parsed_response: BaseModel
+    used_model: str
+    tokens_used: int
+
+
+class OpenAIParserSkill(Skill[OpenAIParserInput, OpenAIParserOutput]):
+    """Skill for getting structured responses from OpenAI"""
+
+    input_schema = OpenAIParserInput
+    output_schema = OpenAIParserOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def process(self, input_data: OpenAIParserInput) -> OpenAIParserOutput:
+        try:
+            # Use parse method instead of create
+            completion = self.client.beta.chat.completions.parse(
+                model=input_data.model,
+                messages=[
+                    {"role": "system", "content": input_data.system_prompt},
+                    {"role": "user", "content": input_data.user_input},
+                ],
+                response_format=input_data.response_model,
+            )
+
+            if completion.choices[0].message.parsed is None:
+                raise ProcessingError("Failed to parse response")
+
+            return OpenAIParserOutput(
+                parsed_response=completion.choices[0].message.parsed,
+                used_model=completion.model,
+                tokens_used=completion.usage.total_tokens,
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"OpenAI parsing failed: {str(e)}")
+
+
+class OpenAIEmbeddingsInput(InputSchema):
+    """Schema for OpenAI embeddings input"""
+
+    texts: Union[str, List[str]] = Field(
+        ..., description="Text or list of texts to generate embeddings for"
+    )
+    model: str = Field(
+        default="text-embedding-3-large", description="OpenAI embeddings model to use"
+    )
+    encoding_format: str = Field(
+        default="float", description="The format of the embeddings: 'float' or 'base64'"
+    )
+    dimensions: Optional[int] = Field(
+        default=None, description="Optional number of dimensions for the embeddings"
+    )
+
+
+class OpenAIEmbeddingsOutput(OutputSchema):
+    """Schema for OpenAI embeddings output"""
+
+    embeddings: List[List[float]] = Field(..., description="List of embeddings vectors")
+    used_model: str = Field(..., description="Model used for generating embeddings")
+    tokens_used: int = Field(..., description="Number of tokens used")
+
+
+class OpenAIEmbeddingsSkill(Skill[OpenAIEmbeddingsInput, OpenAIEmbeddingsOutput]):
+    """Skill for generating embeddings using OpenAI models"""
+
+    input_schema = OpenAIEmbeddingsInput
+    output_schema = OpenAIEmbeddingsOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+        self.async_client = AsyncOpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def process(self, input_data: OpenAIEmbeddingsInput) -> OpenAIEmbeddingsOutput:
+        """Generate embeddings for the input text(s)"""
+        try:
+            # Handle single text input
+            texts = (
+                [input_data.texts]
+                if isinstance(input_data.texts, str)
+                else input_data.texts
+            )
+
+            # Create embeddings
+            response = self.client.embeddings.create(
+                model=input_data.model,
+                input=texts,
+                encoding_format=input_data.encoding_format,
+                dimensions=input_data.dimensions,
+            )
+
+            # Extract embeddings
+            embeddings = [data.embedding for data in response.data]
+
+            return OpenAIEmbeddingsOutput(
+                embeddings=embeddings,
+                used_model=response.model,
+                tokens_used=response.usage.total_tokens,
+            )
+        except Exception as e:
+            raise ProcessingError(f"OpenAI embeddings generation failed: {str(e)}")
+
+    async def process_async(
+        self, input_data: OpenAIEmbeddingsInput
+    ) -> OpenAIEmbeddingsOutput:
+        """Async version of the embeddings generation"""
+        try:
+            # Handle single text input
+            texts = (
+                [input_data.texts]
+                if isinstance(input_data.texts, str)
+                else input_data.texts
+            )
+
+            # Create embeddings
+            response = await self.async_client.embeddings.create(
+                model=input_data.model,
+                input=texts,
+                encoding_format=input_data.encoding_format,
+                dimensions=input_data.dimensions,
+            )
+
+            # Extract embeddings
+            embeddings = [data.embedding for data in response.data]
+
+            return OpenAIEmbeddingsOutput(
+                embeddings=embeddings,
+                used_model=response.model,
+                tokens_used=response.usage.total_tokens,
+            )
+        except Exception as e:
+            raise ProcessingError(
+                f"OpenAI async embeddings generation failed: {str(e)}"
+            )

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

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