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

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

airtrain/integrations/fireworks/structured_completion_skills.py

@@ -0,0 +1,175 @@
+from typing import Any, Dict, Generator, List, Optional, Type, TypeVar
+from pydantic import BaseModel, Field
+import requests
+import json
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+
+ResponseT = TypeVar("ResponseT")
+
+
+class FireworksStructuredCompletionInput(InputSchema):
+    """Schema for Fireworks AI structured completion input"""
+
+    prompt: str = Field(..., description="Input prompt for completion")
+    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")
+    response_model: Type[ResponseT]
+    stream: bool = Field(
+        default=False,
+        description="Whether to stream the response token by token",
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class FireworksStructuredCompletionOutput(OutputSchema):
+    """Schema for Fireworks AI structured completion output"""
+
+    parsed_response: Any
+    used_model: str
+    usage: Dict[str, int]
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        default=None, 
+        description=(
+            "Tool calls are not applicable for completions, "
+            "included for compatibility"
+        )
+    )
+
+
+class FireworksStructuredCompletionSkill(
+    Skill[FireworksStructuredCompletionInput, FireworksStructuredCompletionOutput]
+):
+    """Skill for getting structured completion responses from Fireworks AI"""
+
+    input_schema = FireworksStructuredCompletionInput
+    output_schema = FireworksStructuredCompletionOutput
+    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: FireworksStructuredCompletionInput
+    ) -> Dict[str, Any]:
+        """Build the request payload."""
+        return {
+            "model": input_data.model,
+            "prompt": input_data.prompt,
+            "temperature": input_data.temperature,
+            "max_tokens": input_data.max_tokens,
+            "stream": input_data.stream,
+            "response_format": {
+                "type": "json_object",
+                "schema": {
+                    **input_data.response_model.model_json_schema(),
+                    "required": [
+                        field
+                        for field, _ in input_data.response_model.model_fields.items()
+                    ],
+                },
+            },
+        }
+
+    def process_stream(
+        self, input_data: FireworksStructuredCompletionInput
+    ) -> Generator[Dict[str, Any], 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()
+
+            json_buffer = []
+            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"):
+                            content = data["choices"][0]["text"]
+                            json_buffer.append(content)
+                            yield {"chunk": content}
+                    except json.JSONDecodeError:
+                        continue
+
+            # Once complete, parse the full JSON
+            complete_json = "".join(json_buffer)
+            try:
+                parsed_response = input_data.response_model.model_validate_json(
+                    complete_json
+                )
+                yield {"complete": parsed_response}
+            except Exception as e:
+                raise ProcessingError(f"Failed to parse JSON response: {str(e)}")
+
+        except Exception as e:
+            raise ProcessingError(f"Fireworks streaming request failed: {str(e)}")
+
+    def process(
+        self, input_data: FireworksStructuredCompletionInput
+    ) -> FireworksStructuredCompletionOutput:
+        """Process the input and return structured response."""
+        try:
+            if input_data.stream:
+                # For streaming, collect and parse the entire response
+                json_buffer = []
+                parsed_response = None
+
+                for chunk in self.process_stream(input_data):
+                    if "chunk" in chunk:
+                        json_buffer.append(chunk["chunk"])
+                    elif "complete" in chunk:
+                        parsed_response = chunk["complete"]
+
+                if parsed_response is None:
+                    raise ProcessingError("Failed to parse streamed response")
+
+                return FireworksStructuredCompletionOutput(
+                    parsed_response=parsed_response,
+                    used_model=input_data.model,
+                    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"]
+                parsed_response = input_data.response_model.model_validate_json(
+                    response_text
+                )
+
+                return FireworksStructuredCompletionOutput(
+                    parsed_response=parsed_response,
+                    used_model=input_data.model,
+                    usage=data["usage"],
+                )
+
+        except Exception as e:
+            raise ProcessingError(f"Fireworks structured completion failed: {str(e)}")