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

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)

airtrain/integrations/fireworks/list_models.py

@@ -0,0 +1,128 @@
+from typing import Optional, List
+import requests
+from pydantic import Field
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+from .models import FireworksModel
+
+
+class FireworksListModelsInput(InputSchema):
+    """Schema for Fireworks AI list models input"""
+
+    account_id: str = Field(..., description="The Account Id")
+    page_size: Optional[int] = Field(
+        default=50, 
+        description=(
+            "The maximum number of models to return. The maximum page_size is 200, "
+            "values above 200 will be coerced to 200."
+        ),
+        le=200
+    )
+    page_token: Optional[str] = Field(
+        default=None,
+        description=(
+            "A page token, received from a previous ListModels call. Provide this "
+            "to retrieve the subsequent page. When paginating, all other parameters "
+            "provided to ListModels must match the call that provided the page token."
+        )
+    )
+    filter: Optional[str] = Field(
+        default=None,
+        description=(
+            "Only model satisfying the provided filter (if specified) will be "
+            "returned. See https://google.aip.dev/160 for the filter grammar."
+        )
+    )
+    order_by: Optional[str] = Field(
+        default=None,
+        description=(
+            "A comma-separated list of fields to order by. e.g. \"foo,bar\" "
+            "The default sort order is ascending. To specify a descending order for a "
+            "field, append a \" desc\" suffix. e.g. \"foo desc,bar\" "
+            "Subfields are specified with a \".\" character. e.g. \"foo.bar\" "
+            "If not specified, the default order is by \"name\"."
+        )
+    )
+
+
+class FireworksListModelsOutput(OutputSchema):
+    """Schema for Fireworks AI list models output"""
+
+    models: List[FireworksModel] = Field(
+        default_factory=list, 
+        description="List of Fireworks models"
+    )
+    next_page_token: Optional[str] = Field(
+        default=None, 
+        description="Token for retrieving the next page of results"
+    )
+    total_size: Optional[int] = Field(
+        default=None, 
+        description="Total number of models available"
+    )
+
+
+class FireworksListModelsSkill(
+    Skill[FireworksListModelsInput, FireworksListModelsOutput]
+):
+    """Skill for listing Fireworks AI models"""
+
+    input_schema = FireworksListModelsInput
+    output_schema = FireworksListModelsOutput
+    
+    def __init__(self, credentials: Optional[FireworksCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or FireworksCredentials.from_env()
+        self.base_url = "https://api.fireworks.ai/v1"
+    
+    def process(
+        self, input_data: FireworksListModelsInput
+    ) -> FireworksListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            # Build the URL
+            url = f"{self.base_url}/accounts/{input_data.account_id}/models"
+            
+            # Prepare query parameters
+            params = {}
+            if input_data.page_size:
+                params["pageSize"] = input_data.page_size
+            if input_data.page_token:
+                params["pageToken"] = input_data.page_token
+            if input_data.filter:
+                params["filter"] = input_data.filter
+            if input_data.order_by:
+                params["orderBy"] = input_data.order_by
+            
+            # Make the request
+            headers = {
+                "Authorization": (
+                    f"Bearer {self.credentials.fireworks_api_key.get_secret_value()}"
+                )
+            }
+            
+            response = requests.get(url, headers=headers, params=params)
+            response.raise_for_status()
+            
+            # Parse the response
+            result = response.json()
+            
+            # Convert the models to FireworksModel objects
+            models = []
+            for model_data in result.get("models", []):
+                models.append(FireworksModel(**model_data))
+            
+            # Return the output
+            return FireworksListModelsOutput(
+                models=models,
+                next_page_token=result.get("nextPageToken"),
+                total_size=result.get("totalSize")
+            )
+            
+        except requests.RequestException as e:
+            raise ProcessingError(f"Failed to list Fireworks models: {str(e)}")
+        except Exception as e:
+            raise ProcessingError(f"Error listing Fireworks models: {str(e)}") 

airtrain/integrations/fireworks/models.py

@@ -0,0 +1,139 @@
+from typing import List, Optional, Dict, Any
+from pydantic import Field, BaseModel
+
+
+class FireworksMessage(BaseModel):
+    """Schema for Fireworks chat message"""
+
+    content: str
+    role: str = Field(..., pattern="^(system|user|assistant)$")
+
+
+class FireworksUsage(BaseModel):
+    """Schema for Fireworks API usage statistics"""
+
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class FireworksResponse(BaseModel):
+    """Schema for Fireworks API response"""
+
+    id: str
+    choices: List[Dict[str, Any]]
+    created: int
+    model: str
+    usage: FireworksUsage
+
+
+class FireworksModelStatus(BaseModel):
+    """Schema for Fireworks model status"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksModelBaseDetails(BaseModel):
+    """Schema for Fireworks base model details"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksPeftDetails(BaseModel):
+    """Schema for Fireworks PEFT details"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksConversationConfig(BaseModel):
+    """Schema for Fireworks conversation configuration"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksModelDeployedRef(BaseModel):
+    """Schema for Fireworks deployed model reference"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksDeprecationDate(BaseModel):
+    """Schema for Fireworks deprecation date"""
+    # This would be filled with actual fields from the API response
+
+
+class FireworksModel(BaseModel):
+    """Schema for a Fireworks model"""
+    
+    name: str
+    display_name: Optional[str] = None
+    description: Optional[str] = None
+    create_time: Optional[str] = None
+    created_by: Optional[str] = None
+    state: Optional[str] = None
+    status: Optional[Dict[str, Any]] = None
+    kind: Optional[str] = None
+    github_url: Optional[str] = None
+    hugging_face_url: Optional[str] = None
+    base_model_details: Optional[Dict[str, Any]] = None
+    peft_details: Optional[Dict[str, Any]] = None
+    teft_details: Optional[Dict[str, Any]] = None
+    public: Optional[bool] = None
+    conversation_config: Optional[Dict[str, Any]] = None
+    context_length: Optional[int] = None
+    supports_image_input: Optional[bool] = None
+    supports_tools: Optional[bool] = None
+    imported_from: Optional[str] = None
+    fine_tuning_job: Optional[str] = None
+    default_draft_model: Optional[str] = None
+    default_draft_token_count: Optional[int] = None
+    precisions: Optional[List[str]] = None
+    deployed_model_refs: Optional[List[Dict[str, Any]]] = None
+    cluster: Optional[str] = None
+    deprecation_date: Optional[Dict[str, Any]] = None
+    calibrated: Optional[bool] = None
+    tunable: Optional[bool] = None
+    supports_lora: Optional[bool] = None
+    use_hf_apply_chat_template: Optional[bool] = None
+
+
+class ListModelsInput(BaseModel):
+    """Schema for listing Fireworks models input"""
+    
+    account_id: str = Field(..., description="The Account Id")
+    page_size: Optional[int] = Field(
+        default=50, 
+        description=(
+            "The maximum number of models to return. The maximum page_size is 200, "
+            "values above 200 will be coerced to 200."
+        ),
+        le=200
+    )
+    page_token: Optional[str] = Field(
+        default=None,
+        description=(
+            "A page token, received from a previous ListModels call. Provide this "
+            "to retrieve the subsequent page. When paginating, all other parameters "
+            "provided to ListModels must match the call that provided the page token."
+        )
+    )
+    filter: Optional[str] = Field(
+        default=None,
+        description=(
+            "Only model satisfying the provided filter (if specified) will be "
+            "returned. See https://google.aip.dev/160 for the filter grammar."
+        )
+    )
+    order_by: Optional[str] = Field(
+        default=None,
+        description=(
+            "A comma-separated list of fields to order by. e.g. \"foo,bar\" "
+            "The default sort order is ascending. To specify a descending order for a "
+            "field, append a \" desc\" suffix. e.g. \"foo desc,bar\" "
+            "Subfields are specified with a \".\" character. e.g. \"foo.bar\" "
+            "If not specified, the default order is by \"name\"."
+        )
+    )
+
+
+class ListModelsOutput(BaseModel):
+    """Schema for listing Fireworks models output"""
+    
+    models: List[FireworksModel]
+    next_page_token: Optional[str] = None
+    total_size: Optional[int] = None