airtrain 0.1.13__py3-none-any.whl → 0.1.14__py3-none-any.whl

airtrain/__init__.py

@@ -1,6 +1,6 @@
 """Airtrain - A platform for building and deploying AI agents with structured skills"""
 
-__version__ = "0.1.13"
+__version__ = "0.1.14"
 
 # Core imports
 from .core.skills import Skill, ProcessingError

airtrain/integrations/fireworks/__init__.py

@@ -0,0 +1,11 @@
+"""Fireworks AI integration module"""
+
+from .credentials import FireworksCredentials
+from .skills import FireworksChatSkill, FireworksInput, FireworksOutput
+
+__all__ = [
+    "FireworksCredentials",
+    "FireworksChatSkill",
+    "FireworksInput",
+    "FireworksOutput",
+]

airtrain/integrations/fireworks/credentials.py

@@ -0,0 +1,18 @@
+from pydantic import SecretStr, BaseModel
+from typing import Optional
+import os
+
+
+class FireworksCredentials(BaseModel):
+    """Credentials for Fireworks AI API"""
+
+    fireworks_api_key: SecretStr
+
+    @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/models.py

@@ -0,0 +1,27 @@
+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

airtrain/integrations/fireworks/skills.py

@@ -0,0 +1,107 @@
+from typing import List, Optional, Dict, Any
+from pydantic import Field
+import requests
+from loguru import logger
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+from .models import FireworksMessage, FireworksResponse
+
+
+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",
+    )
+    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=None, description="Maximum tokens in response"
+    )
+    context_length_exceeded_behavior: str = Field(
+        default="truncate", description="Behavior when context length is exceeded"
+    )
+
+
+class FireworksOutput(OutputSchema):
+    """Schema for Fireworks AI 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")
+
+
+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.base_url = "https://api.fireworks.ai/inference/v1"
+
+    def process(self, input_data: FireworksInput) -> FireworksOutput:
+        """Process the input using Fireworks AI API"""
+        try:
+            logger.info(f"Processing request with model {input_data.model}")
+
+            # Prepare messages
+            messages = [
+                {"role": "system", "content": input_data.system_prompt},
+                {"role": "user", "content": input_data.user_input},
+            ]
+
+            # Prepare request payload
+            payload = {
+                "messages": messages,
+                "model": input_data.model,
+                "context_length_exceeded_behavior": input_data.context_length_exceeded_behavior,
+                "temperature": input_data.temperature,
+                "n": 1,
+                "response_format": {"type": "text"},
+                "stream": False,
+            }
+
+            if input_data.max_tokens:
+                payload["max_tokens"] = input_data.max_tokens
+
+            # Make API request
+            response = requests.post(
+                f"{self.base_url}/chat/completions",
+                json=payload,
+                headers={
+                    "Authorization": f"Bearer {self.credentials.fireworks_api_key.get_secret_value()}",
+                    "Content-Type": "application/json",
+                },
+            )
+
+            response.raise_for_status()
+            response_data = FireworksResponse(**response.json())
+
+            logger.success("Successfully processed Fireworks AI request")
+
+            return FireworksOutput(
+                response=response_data.choices[0]["message"]["content"],
+                used_model=response_data.model,
+                usage={
+                    "prompt_tokens": response_data.usage.prompt_tokens,
+                    "completion_tokens": response_data.usage.completion_tokens,
+                    "total_tokens": response_data.usage.total_tokens,
+                },
+            )
+
+        except Exception as e:
+            logger.exception(f"Fireworks AI processing failed: {str(e)}")
+            raise ProcessingError(f"Fireworks AI processing failed: {str(e)}")

{airtrain-0.1.13.dist-info → airtrain-0.1.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.13
+Version: 0.1.14
 Summary: A platform for building and deploying AI agents with structured skills
 Home-page: https://github.com/rosaboyle/airtrain.dev
 Author: Dheeraj Pai

{airtrain-0.1.13.dist-info → airtrain-0.1.14.dist-info}/RECORD

@@ -1,4 +1,4 @@
-airtrain/__init__.py,sha256=BwgzBp8NKID9GrpW2RUmCyV6Aj_wlKI9HJqHTe6blMw,2095
+airtrain/__init__.py,sha256=FRI-wOLxVAu3ECtPJGQ8ZLjDSM1PgoiDGvn8ctNtS_8,2095
 airtrain/contrib/__init__.py,sha256=pG-7mJ0pBMqp3Q86mIF9bo1PqoBOVSGlnEK1yY1U1ok,641
 airtrain/contrib/travel/__init__.py,sha256=clmBodw4nkTA-DsgjVGcXfJGPaWxIpCZDtdO-8RzL0M,811
 airtrain/contrib/travel/agents.py,sha256=tpQtZ0WUiXBuhvZtc2JlEam5TuR5l-Tndi14YyImDBM,8975
@@ -17,6 +17,10 @@ airtrain/integrations/aws/skills.py,sha256=TQiMXeXRRcJ14fe8Xi7Uk20iS6_INbcznuLGt
 airtrain/integrations/cerebras/__init__.py,sha256=zAD-qV38OzHhMCz1z-NvjjqcYEhURbm8RWTOKHNqbew,174
 airtrain/integrations/cerebras/credentials.py,sha256=IFkn8LxMAaOpvEWXDpb94VQGtqcDxQ7rZHKH-tX4Nuw,884
 airtrain/integrations/cerebras/skills.py,sha256=O9vwFzvv_tUOwFOVE8CszAQEac711eVYVUj_8dVMTpc,1596
+airtrain/integrations/fireworks/__init__.py,sha256=9pJvP0u1FJbNtB0oHa09mHVJLctELf_c27LOYyDk2ZI,271
+airtrain/integrations/fireworks/credentials.py,sha256=UpcwR9V5Hbk5sJbjFDJDbHMRqc90IQSqAvrtJCOvwEo,524
+airtrain/integrations/fireworks/models.py,sha256=F-MddbLCLAsTjwRr1l6IpJxOegyY4pD7jN9ySPiypSo,593
+airtrain/integrations/fireworks/skills.py,sha256=ZykowW8lMbTcZVJ0GO2Ut6E-u2-keXvE4F-_j-3JI4k,4074
 airtrain/integrations/google/__init__.py,sha256=INZFNOcNebz3m-Ggk07ZjmX0kNHIbTe_St9gBlZBki8,176
 airtrain/integrations/google/credentials.py,sha256=Mm4jNWF02rIf0_GuHLcUUPyLHC4NMRdF_iTCoVTQ0Bs,1033
 airtrain/integrations/google/skills.py,sha256=uwmgetl5Ien7fLOA5HIZdqoL6AZnexFDyzfsrGuJ1RU,1606
@@ -47,7 +51,7 @@ airtrain/integrations/together/rerank_skill.py,sha256=gjH24hLWCweWKPyyfKZMG3K_g9
 airtrain/integrations/together/schemas.py,sha256=pBMrbX67oxPCr-sg4K8_Xqu1DWbaC4uLCloVSascROg,1210
 airtrain/integrations/together/skills.py,sha256=UfLHnseZbA7R7q5dDco6mpV546Zfd3DTliZSrNkCL6Q,4518
 airtrain/integrations/together/vision_models_config.py,sha256=m28HwYDk2Kup_J-a1FtynIa2ZVcbl37kltfoHnK8zxs,1544
-airtrain-0.1.13.dist-info/METADATA,sha256=kvTff3-jQ_th6KyfKrSQ8EOjjq5SiyHwq802VFYKqCg,4536
-airtrain-0.1.13.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-airtrain-0.1.13.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
-airtrain-0.1.13.dist-info/RECORD,,
+airtrain-0.1.14.dist-info/METADATA,sha256=l4IPKLJ7Bf3gmZYSRPVEfz4oe1XGt_lWfvuZg68cNnE,4536
+airtrain-0.1.14.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+airtrain-0.1.14.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.14.dist-info/RECORD,,