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

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

airtrain/integrations/fireworks/requests_skills.py

@@ -0,0 +1,207 @@
+from typing import List, Optional, Dict, Any, Generator, AsyncGenerator
+from pydantic import Field
+import requests
+import json
+from loguru import logger
+import aiohttp
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+
+
+class FireworksRequestInput(InputSchema):
+    """Schema for Fireworks AI chat input using requests"""
+
+    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",
+    )
+    model: str = Field(
+        default="accounts/fireworks/models/deepseek-r1",
+        description="Fireworks AI model to use",
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
+    )
+    max_tokens: int = Field(default=131072, description="Maximum tokens in response")
+    top_p: float = Field(
+        default=1.0, description="Top p sampling parameter", ge=0, le=1
+    )
+    top_k: int = Field(default=40, 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
+    )
+    stream: bool = Field(
+        default=False,
+        description="Whether to stream the response",
+    )
+
+
+class FireworksRequestOutput(OutputSchema):
+    """Schema for Fireworks AI chat output"""
+
+    response: str
+    used_model: str
+    usage: Dict[str, int]
+
+
+class FireworksRequestSkill(Skill[FireworksRequestInput, FireworksRequestOutput]):
+    """Skill for interacting with Fireworks AI models using requests"""
+
+    input_schema = FireworksRequestInput
+    output_schema = FireworksRequestOutput
+    BASE_URL = "https://api.fireworks.ai/inference/v1/chat/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()}",
+        }
+        self.stream_headers = {
+            "Accept": "text/event-stream",
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.credentials.fireworks_api_key.get_secret_value()}",
+        }
+
+    def _build_messages(
+        self, input_data: FireworksRequestInput
+    ) -> 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 _build_payload(self, input_data: FireworksRequestInput) -> Dict[str, Any]:
+        """Build the request payload."""
+        return {
+            "model": input_data.model,
+            "messages": self._build_messages(input_data),
+            "temperature": input_data.temperature,
+            "max_tokens": input_data.max_tokens,
+            "top_p": input_data.top_p,
+            "top_k": input_data.top_k,
+            "presence_penalty": input_data.presence_penalty,
+            "frequency_penalty": input_data.frequency_penalty,
+            "stream": input_data.stream,
+        }
+
+    def process_stream(
+        self, input_data: FireworksRequestInput
+    ) -> 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["choices"][0]["delta"].get("content"):
+                            yield data["choices"][0]["delta"]["content"]
+                    except json.JSONDecodeError:
+                        continue
+
+        except Exception as e:
+            raise ProcessingError(f"Fireworks streaming request failed: {str(e)}")
+
+    def process(self, input_data: FireworksRequestInput) -> FireworksRequestOutput:
+        """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_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]["message"]["content"]
+                usage = data["usage"]
+
+            return FireworksRequestOutput(
+                response=response_text, used_model=input_data.model, usage=usage
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Fireworks request failed: {str(e)}")
+
+    async def process_async(
+        self, input_data: FireworksRequestInput
+    ) -> FireworksRequestOutput:
+        """Async version of process method using aiohttp"""
+        try:
+            async with aiohttp.ClientSession() as session:
+                payload = self._build_payload(input_data)
+                async with session.post(
+                    self.BASE_URL, headers=self.headers, json=payload
+                ) as response:
+                    response.raise_for_status()
+                    data = await response.json()
+
+                    return FireworksRequestOutput(
+                        response=data["choices"][0]["message"]["content"],
+                        used_model=input_data.model,
+                        usage=data.get("usage", {}),
+                    )
+
+        except Exception as e:
+            raise ProcessingError(f"Async Fireworks request failed: {str(e)}")
+
+    async def process_stream_async(
+        self, input_data: FireworksRequestInput
+    ) -> AsyncGenerator[str, None]:
+        """Async version of stream processor using aiohttp"""
+        try:
+            async with aiohttp.ClientSession() as session:
+                payload = self._build_payload(input_data)
+                async with session.post(
+                    self.BASE_URL, headers=self.stream_headers, json=payload
+                ) as response:
+                    response.raise_for_status()
+
+                    async for line in response.content:
+                        if line.startswith(b"data: "):
+                            chunk = json.loads(line[6:].strip())
+                            if "choices" in chunk:
+                                content = (
+                                    chunk["choices"][0]
+                                    .get("delta", {})
+                                    .get("content", "")
+                                )
+                                if content:
+                                    yield content
+
+        except Exception as e:
+            raise ProcessingError(f"Async Fireworks streaming failed: {str(e)}")

airtrain/integrations/fireworks/skills.py

@@ -0,0 +1,181 @@
+from typing import List, Optional, Dict, Any, Generator, Union
+from pydantic import Field
+from openai import OpenAI
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+
+
+class FireworksInput(InputSchema):
+    """Schema for Fireworks AI 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, Any]] = Field(
+        default_factory=list,
+        description="List of previous conversation messages",
+    )
+    model: str = Field(
+        default="accounts/fireworks/models/deepseek-r1",
+        description="Fireworks AI 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"
+    )
+    context_length_exceeded_behavior: str = Field(
+        default="truncate", description="Behavior when context length is exceeded"
+    )
+    stream: bool = Field(
+        default=False,
+        description="Whether to stream the response token by token",
+    )
+    tools: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description=(
+            "A list of tools the model may use. "
+            "Currently only functions supported."
+        ),
+    )
+    tool_choice: Optional[Union[str, Dict[str, Any]]] = Field(
+        default=None,
+        description=(
+            "Controls which tool is called by the model. "
+            "'none', 'auto', or specific tool."
+        ),
+    )
+
+
+class FireworksOutput(OutputSchema):
+    """Schema for Fireworks 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(default_factory=dict, description="Usage statistics")
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        default=None, description="Tool calls generated by the model"
+    )
+
+
+class FireworksChatSkill(Skill[FireworksInput, FireworksOutput]):
+    """Skill for interacting with Fireworks AI models"""
+
+    input_schema = FireworksInput
+    output_schema = FireworksOutput
+
+    def __init__(self, credentials: Optional[FireworksCredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or FireworksCredentials.from_env()
+        self.client = OpenAI(
+            base_url="https://api.fireworks.ai/inference/v1",
+            api_key=self.credentials.fireworks_api_key.get_secret_value(),
+        )
+
+    def _build_messages(self, input_data: FireworksInput) -> List[Dict[str, Any]]:
+        """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: FireworksInput) -> 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"Fireworks streaming failed: {str(e)}")
+
+    def process(self, input_data: FireworksInput) -> FireworksOutput:
+        """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)
+                
+                # Create completion object for usage stats
+                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,
+                )
+            else:
+                # For non-streaming, use regular completion
+                messages = self._build_messages(input_data)
+                
+                # Prepare API call parameters
+                api_params = {
+                    "model": input_data.model,
+                    "messages": messages,
+                    "temperature": input_data.temperature,
+                    "max_tokens": input_data.max_tokens,
+                    "stream": False,
+                }
+                
+                # Add tools and tool_choice if provided
+                if input_data.tools:
+                    api_params["tools"] = input_data.tools
+                
+                if input_data.tool_choice:
+                    api_params["tool_choice"] = input_data.tool_choice
+                
+                completion = self.client.chat.completions.create(**api_params)
+                response = completion.choices[0].message.content or ""
+            
+            # Check for tool calls in the response
+            tool_calls = None
+            if (hasattr(completion.choices[0].message, "tool_calls") and
+                    completion.choices[0].message.tool_calls):
+                tool_calls = [
+                    {
+                        "id": tool_call.id,
+                        "type": tool_call.type,
+                        "function": {
+                            "name": tool_call.function.name,
+                            "arguments": tool_call.function.arguments
+                        }
+                    }
+                    for tool_call in completion.choices[0].message.tool_calls
+                ]
+
+            return FireworksOutput(
+                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,
+                },
+                tool_calls=tool_calls
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Fireworks chat failed: {str(e)}")