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

airtrain/integrations/combined/groq_fireworks_skills.py

@@ -0,0 +1,126 @@
+from typing import Optional, Dict, Any, List
+from pydantic import Field
+import requests
+from groq import Groq
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from airtrain.integrations.fireworks.completion_skills import (
+    FireworksCompletionSkill,
+    FireworksCompletionInput,
+)
+
+
+class GroqFireworksInput(InputSchema):
+    """Schema for combined Groq and Fireworks input"""
+
+    user_input: str = Field(..., description="User's input text")
+    groq_model: str = Field(
+        default="mixtral-8x7b-32768", description="Groq model to use"
+    )
+    fireworks_model: str = Field(
+        default="accounts/fireworks/models/deepseek-r1",
+        description="Fireworks model to use",
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation"
+    )
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
+
+
+class GroqFireworksOutput(OutputSchema):
+    """Schema for combined Groq and Fireworks output"""
+
+    combined_response: str
+    groq_response: str
+    fireworks_response: str
+    used_models: Dict[str, str]
+    usage: Dict[str, Dict[str, int]]
+
+
+class GroqFireworksSkill(Skill[GroqFireworksInput, GroqFireworksOutput]):
+    """Skill combining Groq and Fireworks responses"""
+
+    input_schema = GroqFireworksInput
+    output_schema = GroqFireworksOutput
+
+    def __init__(
+        self,
+        groq_api_key: Optional[str] = None,
+        fireworks_skill: Optional[FireworksCompletionSkill] = None,
+    ):
+        """Initialize the skill with optional API keys"""
+        super().__init__()
+        self.groq_client = Groq(api_key=groq_api_key)
+        self.fireworks_skill = fireworks_skill or FireworksCompletionSkill()
+
+    def _get_groq_response(self, input_data: GroqFireworksInput) -> Dict[str, Any]:
+        """Get response from Groq"""
+        try:
+            completion = self.groq_client.chat.completions.create(
+                model=input_data.groq_model,
+                messages=[{"role": "user", "content": input_data.user_input}],
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+            return {
+                "response": completion.choices[0].message.content,
+                "usage": completion.usage.model_dump(),
+            }
+        except Exception as e:
+            raise ProcessingError(f"Groq request failed: {str(e)}")
+
+    def _get_fireworks_response(
+        self, groq_response: str, input_data: GroqFireworksInput
+    ) -> Dict[str, Any]:
+        """Get response from Fireworks"""
+        try:
+            formatted_prompt = (
+                f"<USER>{input_data.user_input}</USER>\n<ASSISTANT>{groq_response}"
+            )
+
+            fireworks_input = FireworksCompletionInput(
+                prompt=formatted_prompt,
+                model=input_data.fireworks_model,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+
+            result = self.fireworks_skill.process(fireworks_input)
+            return {"response": result.response, "usage": result.usage}
+        except Exception as e:
+            raise ProcessingError(f"Fireworks request failed: {str(e)}")
+
+    def process(self, input_data: GroqFireworksInput) -> GroqFireworksOutput:
+        """Process the input using both Groq and Fireworks"""
+        try:
+            # Get Groq response
+            groq_result = self._get_groq_response(input_data)
+
+            # Get Fireworks response
+            fireworks_result = self._get_fireworks_response(
+                groq_result["response"], input_data
+            )
+
+            # Combine responses in the required format
+            combined_response = (
+                f"<USER>{input_data.user_input}</USER>\n"
+                f"<ASSISTANT>{groq_result['response']} {fireworks_result['response']}"
+            )
+
+            return GroqFireworksOutput(
+                combined_response=combined_response,
+                groq_response=groq_result["response"],
+                fireworks_response=fireworks_result["response"],
+                used_models={
+                    "groq": input_data.groq_model,
+                    "fireworks": input_data.fireworks_model,
+                },
+                usage={
+                    "groq": groq_result["usage"],
+                    "fireworks": fireworks_result["usage"],
+                },
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Combined processing failed: {str(e)}")

airtrain/integrations/combined/list_models_factory.py

@@ -0,0 +1,210 @@
+from typing import Optional, Dict, Any, List
+from pydantic import Field
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from airtrain.core.credentials import BaseCredentials
+
+# Import existing list models skills
+from airtrain.integrations.openai.list_models import OpenAIListModelsSkill
+from airtrain.integrations.anthropic.list_models import AnthropicListModelsSkill
+from airtrain.integrations.together.list_models import TogetherListModelsSkill
+from airtrain.integrations.fireworks.list_models import FireworksListModelsSkill
+
+# Import credentials
+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:
+        arbitrary_types_allowed = True
+        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")
+
+
+# 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:
+            models = self.get_models()
+            return GenericListModelsOutput(models=models, provider=self.provider)
+        except Exception as e:
+            raise ProcessingError(f"Failed to list {self.provider} models: {str(e)}")
+
+
+# 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": "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-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)",
+            },
+        ]
+        return models
+
+
+# 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",
+            },
+        ]
+        return models
+
+
+# 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"},
+        ]
+        return models
+
+
+# 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,
+        "anthropic": AnthropicListModelsSkill,
+        "together": TogetherListModelsSkill,
+        "fireworks": FireworksListModelsSkill,
+        "groq": GroqListModelsSkill,
+        "cerebras": CerebrasListModelsSkill,
+        "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())

airtrain/integrations/fireworks/__init__.py

@@ -0,0 +1,21 @@
+"""Fireworks AI integration module"""
+
+from .credentials import FireworksCredentials
+from .skills import FireworksChatSkill, FireworksInput, FireworksOutput
+from .list_models import (
+    FireworksListModelsSkill,
+    FireworksListModelsInput,
+    FireworksListModelsOutput,
+)
+from .models import FireworksModel
+
+__all__ = [
+    "FireworksCredentials",
+    "FireworksChatSkill",
+    "FireworksInput",
+    "FireworksOutput",
+    "FireworksListModelsSkill",
+    "FireworksListModelsInput",
+    "FireworksListModelsOutput",
+    "FireworksModel",
+]

airtrain/integrations/fireworks/completion_skills.py

@@ -0,0 +1,147 @@
+from typing import List, Optional, Dict, Any, Generator, Union
+from pydantic import Field
+import requests
+import json
+from loguru import logger
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+
+
+class FireworksCompletionInput(InputSchema):
+    """Schema for Fireworks AI completion input using requests"""
+
+    prompt: str = Field(..., description="Input prompt for completion")
+    model: str = Field(
+        default="accounts/fireworks/models/deepseek-r1",
+        description="Fireworks AI model to use",
+    )
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+    top_p: float = Field(
+        default=1.0, description="Top p sampling parameter", ge=0, le=1
+    )
+    top_k: int = Field(default=50, description="Top k sampling parameter", ge=0)
+    presence_penalty: float = Field(
+        default=0.0, description="Presence penalty", ge=-2.0, le=2.0
+    )
+    frequency_penalty: float = Field(
+        default=0.0, description="Frequency penalty", ge=-2.0, le=2.0
+    )
+    repetition_penalty: float = Field(
+        default=1.0, description="Repetition penalty", ge=0.0
+    )
+    stop: Optional[Union[str, List[str]]] = Field(
+        default=None, description="Stop sequences"
+    )
+    echo: bool = Field(default=False, description="Echo the prompt in the response")
+    stream: bool = Field(default=False, description="Whether to stream the response")
+
+
+class FireworksCompletionOutput(OutputSchema):
+    """Schema for Fireworks AI completion output"""
+
+    response: str
+    used_model: str
+    usage: Dict[str, int]
+
+
+class FireworksCompletionSkill(
+    Skill[FireworksCompletionInput, FireworksCompletionOutput]
+):
+    """Skill for text completion using Fireworks AI"""
+
+    input_schema = FireworksCompletionInput
+    output_schema = FireworksCompletionOutput
+    BASE_URL = "https://api.fireworks.ai/inference/v1/completions"
+
+    def __init__(self, credentials: Optional[FireworksCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or FireworksCredentials.from_env()
+        self.headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.credentials.fireworks_api_key.get_secret_value()}",
+        }
+
+    def _build_payload(self, input_data: FireworksCompletionInput) -> Dict[str, Any]:
+        """Build the request payload."""
+        payload = {
+            "model": input_data.model,
+            "prompt": input_data.prompt,
+            "max_tokens": input_data.max_tokens,
+            "temperature": input_data.temperature,
+            "top_p": input_data.top_p,
+            "top_k": input_data.top_k,
+            "presence_penalty": input_data.presence_penalty,
+            "frequency_penalty": input_data.frequency_penalty,
+            "repetition_penalty": input_data.repetition_penalty,
+            "echo": input_data.echo,
+            "stream": input_data.stream,
+        }
+
+        if input_data.stop:
+            payload["stop"] = input_data.stop
+
+        return payload
+
+    def process_stream(
+        self, input_data: FireworksCompletionInput
+    ) -> Generator[str, None, None]:
+        """Process the input and stream the response."""
+        try:
+            payload = self._build_payload(input_data)
+            response = requests.post(
+                self.BASE_URL,
+                headers=self.headers,
+                data=json.dumps(payload),
+                stream=True,
+            )
+            response.raise_for_status()
+
+            for line in response.iter_lines():
+                if line:
+                    try:
+                        data = json.loads(line.decode("utf-8").removeprefix("data: "))
+                        if data.get("choices") and data["choices"][0].get("text"):
+                            yield data["choices"][0]["text"]
+                    except json.JSONDecodeError:
+                        continue
+
+        except Exception as e:
+            raise ProcessingError(f"Fireworks completion streaming failed: {str(e)}")
+
+    def process(
+        self, input_data: FireworksCompletionInput
+    ) -> FireworksCompletionOutput:
+        """Process the input and return completion 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_text = "".join(response_chunks)
+                usage = {}  # Usage stats not available in streaming mode
+            else:
+                # For non-streaming, use regular request
+                payload = self._build_payload(input_data)
+                response = requests.post(
+                    self.BASE_URL, headers=self.headers, data=json.dumps(payload)
+                )
+                response.raise_for_status()
+                data = response.json()
+
+                response_text = data["choices"][0]["text"]
+                usage = data["usage"]
+
+            return FireworksCompletionOutput(
+                response=response_text, used_model=input_data.model, usage=usage
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Fireworks completion failed: {str(e)}")

airtrain/integrations/fireworks/conversation_manager.py

@@ -0,0 +1,109 @@
+from typing import List, Dict, Optional
+from pydantic import BaseModel, Field
+from .skills import FireworksChatSkill, FireworksInput, FireworksOutput
+
+# TODO: Test this thing.
+
+
+class ConversationState(BaseModel):
+    """Model to track conversation state"""
+
+    messages: List[Dict[str, str]] = Field(
+        default_factory=list, description="List of conversation messages"
+    )
+    system_prompt: str = Field(
+        default="You are a helpful assistant.",
+        description="System prompt for the conversation",
+    )
+    model: str = Field(
+        default="accounts/fireworks/models/deepseek-r1",
+        description="Model being used for the conversation",
+    )
+    temperature: float = Field(default=0.7, description="Temperature setting")
+    max_tokens: Optional[int] = Field(default=131072, description="Max tokens setting")
+
+
+class FireworksConversationManager:
+    """Manager for handling conversation state with Fireworks AI"""
+
+    def __init__(
+        self,
+        skill: Optional[FireworksChatSkill] = None,
+        system_prompt: str = "You are a helpful assistant.",
+        model: str = "accounts/fireworks/models/deepseek-r1",
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+    ):
+        """
+        Initialize conversation manager.
+
+        Args:
+            skill: FireworksChatSkill instance (creates new one if None)
+            system_prompt: Initial system prompt
+            model: Model to use
+            temperature: Temperature setting
+            max_tokens: Max tokens setting
+        """
+        self.skill = skill or FireworksChatSkill()
+        self.state = ConversationState(
+            system_prompt=system_prompt,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+
+    def send_message(self, user_input: str) -> FireworksOutput:
+        """
+        Send a message and get response while maintaining conversation history.
+
+        Args:
+            user_input: User's message
+
+        Returns:
+            FireworksOutput: Model's response
+        """
+        # Create input with current conversation state
+        input_data = FireworksInput(
+            user_input=user_input,
+            system_prompt=self.state.system_prompt,
+            conversation_history=self.state.messages,
+            model=self.state.model,
+            temperature=self.state.temperature,
+            max_tokens=self.state.max_tokens,
+        )
+
+        # Get response
+        result = self.skill.process(input_data)
+
+        # Update conversation history
+        self.state.messages.extend(
+            [
+                {"role": "user", "content": user_input},
+                {"role": "assistant", "content": result.response},
+            ]
+        )
+
+        return result
+
+    def reset_conversation(self) -> None:
+        """Reset the conversation history while maintaining other settings"""
+        self.state.messages = []
+
+    def get_conversation_history(self) -> List[Dict[str, str]]:
+        """Get the current conversation history"""
+        return self.state.messages.copy()
+
+    def update_system_prompt(self, new_prompt: str) -> None:
+        """Update the system prompt for future messages"""
+        self.state.system_prompt = new_prompt
+
+    def save_state(self, file_path: str) -> None:
+        """Save conversation state to a file"""
+        with open(file_path, "w") as f:
+            f.write(self.state.model_dump_json(indent=2))
+
+    def load_state(self, file_path: str) -> None:
+        """Load conversation state from a file"""
+        with open(file_path, "r") as f:
+            data = f.read()
+            self.state = ConversationState.model_validate_json(data)

airtrain/integrations/fireworks/credentials.py

@@ -0,0 +1,26 @@
+from pydantic import SecretStr, BaseModel, Field
+from typing import Optional
+import os
+
+
+class FireworksCredentials(BaseModel):
+    """Credentials for Fireworks AI API"""
+
+    fireworks_api_key: SecretStr = Field(..., min_length=1)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the credentials."""
+        return f"FireworksCredentials(fireworks_api_key=SecretStr('**********'))"
+
+    def __str__(self) -> str:
+        """Return a string representation of the credentials."""
+        return self.__repr__()
+
+    @classmethod
+    def from_env(cls) -> "FireworksCredentials":
+        """Create credentials from environment variables"""
+        api_key = os.getenv("FIREWORKS_API_KEY")
+        if not api_key:
+            raise ValueError("FIREWORKS_API_KEY environment variable not set")
+
+        return cls(fireworks_api_key=api_key)