airtrain 0.1.26__py3-none-any.whl → 0.1.28__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.26"
+__version__ = "0.1.28"
 
 # Core imports
 from .core.skills import Skill, ProcessingError

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=4096, 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/requests_skills.py

@@ -0,0 +1,152 @@
+from typing import List, Optional, Dict, Any, Generator
+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 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=4096, 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()}",
+        }
+
+    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)}")

airtrain/integrations/fireworks/skills.py

@@ -1,7 +1,9 @@
-from typing import List, Optional, Dict, Any
+from typing import List, Optional, Dict, Any, Generator
 from pydantic import Field
 import requests
 from loguru import logger
+from openai import OpenAI
+from openai.types.chat import ChatCompletionChunk
 
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
@@ -34,10 +36,14 @@ class FireworksInput(InputSchema):
     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",
+    )
 
 
 class FireworksOutput(OutputSchema):
-    """Schema for Fireworks AI output"""
+    """Schema for Fireworks AI chat output"""
 
     response: str = Field(..., description="Model's response text")
     used_model: str = Field(..., description="Model used for generation")
@@ -54,76 +60,71 @@ class FireworksChatSkill(Skill[FireworksInput, FireworksOutput]):
         """Initialize the skill with optional credentials"""
         super().__init__()
         self.credentials = credentials or FireworksCredentials.from_env()
-        self.base_url = "https://api.fireworks.ai/inference/v1"
+        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, str]]:
-        """
-        Build messages list from input data including conversation history.
-
-        Args:
-            input_data: The input data containing system prompt, conversation history, and user input
-
-        Returns:
-            List[Dict[str, str]]: List of messages in the format required by Fireworks AI
-        """
+        """Build messages list from input data including conversation history."""
         messages = [{"role": "system", "content": input_data.system_prompt}]
 
-        # Add conversation history if present
         if input_data.conversation_history:
             messages.extend(input_data.conversation_history)
 
-        # Add current user input
         messages.append({"role": "user", "content": input_data.user_input})
-
         return messages
 
-    def process(self, input_data: FireworksInput) -> FireworksOutput:
-        """Process the input using Fireworks AI API"""
+    def process_stream(self, input_data: FireworksInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
         try:
-            logger.info(f"Processing request with model {input_data.model}")
-
-            # Build messages using the helper method
             messages = self._build_messages(input_data)
 
-            # 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",
-                },
+            stream = self.client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+                stream=True,
             )
 
-            response.raise_for_status()
-            response_data = FireworksResponse(**response.json())
+            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)}")
 
-            logger.success("Successfully processed Fireworks AI request")
+    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)
+            else:
+                # For non-streaming, use regular completion
+                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,
+                )
+                response = completion.choices[0].message.content
 
             return FireworksOutput(
-                response=response_data.choices[0]["message"]["content"],
-                used_model=response_data.model,
+                response=response,
+                used_model=input_data.model,
                 usage={
-                    "prompt_tokens": response_data.usage.prompt_tokens,
-                    "completion_tokens": response_data.usage.completion_tokens,
-                    "total_tokens": response_data.usage.total_tokens,
+                    "total_tokens": completion.usage.total_tokens,
+                    "prompt_tokens": completion.usage.prompt_tokens,
+                    "completion_tokens": completion.usage.completion_tokens,
                 },
             )
 
         except Exception as e:
-            logger.exception(f"Fireworks AI processing failed: {str(e)}")
-            raise ProcessingError(f"Fireworks AI processing failed: {str(e)}")
+            raise ProcessingError(f"Fireworks chat failed: {str(e)}")

airtrain/integrations/fireworks/structured_completion_skills.py

@@ -0,0 +1,169 @@
+from typing import Type, TypeVar, Optional, List, Dict, Any, Generator
+from pydantic import BaseModel, 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
+
+ResponseT = TypeVar("ResponseT", bound=BaseModel)
+
+
+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=4096, description="Maximum tokens in response")
+    response_model: Type[ResponseT]
+    stream: bool = Field(
+        default=False,
+        description="Whether to stream the response",
+    )
+
+    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]
+
+
+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)}")

airtrain/integrations/fireworks/structured_requests_skills.py

@@ -0,0 +1,209 @@
+from typing import Type, TypeVar, Optional, List, Dict, Any, Generator
+from pydantic import BaseModel, Field
+import requests
+import json
+from loguru import logger
+import re
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+
+ResponseT = TypeVar("ResponseT", bound=BaseModel)
+
+
+class FireworksStructuredRequestInput(InputSchema):
+    """Schema for Fireworks AI structured output input using requests"""
+
+    user_input: str = Field(..., description="User's input text")
+    system_prompt: str = Field(
+        default="You are a helpful assistant that provides structured data.",
+        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=4096, description="Maximum tokens in response")
+    response_model: Type[ResponseT]
+    stream: bool = Field(
+        default=False,
+        description="Whether to stream the response",
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class FireworksStructuredRequestOutput(OutputSchema):
+    """Schema for Fireworks AI structured output"""
+
+    parsed_response: Any
+    used_model: str
+    usage: Dict[str, int]
+    reasoning: Optional[str] = None
+
+
+class FireworksStructuredRequestSkill(
+    Skill[FireworksStructuredRequestInput, FireworksStructuredRequestOutput]
+):
+    """Skill for getting structured responses from Fireworks AI using requests"""
+
+    input_schema = FireworksStructuredRequestInput
+    output_schema = FireworksStructuredRequestOutput
+    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()}",
+        }
+
+    def _build_messages(
+        self, input_data: FireworksStructuredRequestInput
+    ) -> 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: FireworksStructuredRequestInput
+    ) -> 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,
+            "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: FireworksStructuredRequestInput
+    ) -> 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["choices"][0]["delta"].get("content"):
+                            content = data["choices"][0]["delta"]["content"]
+                            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 _parse_response_content(self, content: str) -> tuple[Optional[str], str]:
+        """Parse response content to extract reasoning and JSON."""
+        # Extract reasoning if present
+        reasoning_match = re.search(r"<think>(.*?)</think>", content, re.DOTALL)
+        reasoning = reasoning_match.group(1).strip() if reasoning_match else None
+
+        # Extract JSON
+        json_match = re.search(r"</think>\s*(\{.*\})", content, re.DOTALL)
+        json_str = json_match.group(1).strip() if json_match else content
+
+        return reasoning, json_str
+
+    def process(
+        self, input_data: FireworksStructuredRequestInput
+    ) -> FireworksStructuredRequestOutput:
+        """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 FireworksStructuredRequestOutput(
+                    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_content = data["choices"][0]["message"]["content"]
+
+                # Parse the response content to extract reasoning and JSON
+                reasoning, json_str = self._parse_response_content(response_content)
+
+                # Parse the JSON string into the specified model
+                parsed_response = input_data.response_model.model_validate_json(
+                    json_str
+                )
+
+                return FireworksStructuredRequestOutput(
+                    parsed_response=parsed_response,
+                    used_model=input_data.model,
+                    usage=data["usage"],
+                    reasoning=reasoning,  # Add reasoning to output if present
+                )
+
+        except Exception as e:
+            raise ProcessingError(f"Fireworks structured request failed: {str(e)}")

airtrain/integrations/fireworks/structured_skills.py

@@ -0,0 +1,102 @@
+from typing import Type, TypeVar, Optional, List, Dict, Any
+from pydantic import BaseModel, Field
+from openai import OpenAI
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import FireworksCredentials
+import re
+
+# Generic type variable for Pydantic response models
+ResponseT = TypeVar("ResponseT", bound=BaseModel)
+
+
+class FireworksParserInput(InputSchema):
+    """Schema for Fireworks structured output input"""
+
+    user_input: str
+    system_prompt: str = "You are a helpful assistant that provides structured data."
+    model: str = "accounts/fireworks/models/deepseek-r1"
+    temperature: float = 0.7
+    max_tokens: Optional[int] = None
+    response_model: Type[ResponseT]
+    conversation_history: List[Dict[str, str]] = Field(
+        default_factory=list,
+        description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class FireworksParserOutput(OutputSchema):
+    """Schema for Fireworks structured output"""
+
+    parsed_response: BaseModel
+    used_model: str
+    tokens_used: int
+    reasoning: Optional[str] = None
+
+
+class FireworksParserSkill(Skill[FireworksParserInput, FireworksParserOutput]):
+    """Skill for getting structured responses from Fireworks"""
+
+    input_schema = FireworksParserInput
+    output_schema = FireworksParserOutput
+
+    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 process(self, input_data: FireworksParserInput) -> FireworksParserOutput:
+        try:
+            # Build messages list including conversation history
+            messages = [{"role": "system", "content": input_data.system_prompt}]
+
+            # Add conversation history if present
+            if input_data.conversation_history:
+                messages.extend(input_data.conversation_history)
+
+            # Add current user input
+            messages.append({"role": "user", "content": input_data.user_input})
+
+            # Make API call with JSON schema
+            completion = self.client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                response_format={
+                    "type": "json_object",
+                    "schema": input_data.response_model.model_json_schema(),
+                },
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+
+            response_content = completion.choices[0].message.content
+
+            # Extract reasoning if present
+            reasoning_match = re.search(
+                r"<think>(.*?)</think>", response_content, re.DOTALL
+            )
+            reasoning = reasoning_match.group(1).strip() if reasoning_match else None
+
+            # Extract JSON
+            json_match = re.search(r"</think>\s*(\{.*\})", response_content, re.DOTALL)
+            json_str = json_match.group(1).strip() if json_match else response_content
+
+            # Parse the response into the specified model
+            parsed_response = input_data.response_model.parse_raw(json_str)
+
+            return FireworksParserOutput(
+                parsed_response=parsed_response,
+                used_model=completion.model,
+                tokens_used=completion.usage.total_tokens,
+                reasoning=reasoning,
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Fireworks parsing failed: {str(e)}")

airtrain/integrations/openai/skills.py

@@ -1,9 +1,10 @@
-from typing import List, Optional, Dict, Any, TypeVar, Type
+from typing import List, Optional, Dict, Any, TypeVar, Type, Generator
 from pydantic import Field, BaseModel
 from openai import OpenAI
 import base64
 from pathlib import Path
 from loguru import logger
+from openai.types.chat import ChatCompletionChunk
 
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
@@ -18,40 +19,36 @@ class OpenAIInput(InputSchema):
         default="You are a helpful assistant.",
         description="System prompt to guide the model's behavior",
     )
-    model: str = Field(default="gpt-4o", description="OpenAI model to use")
-    max_tokens: int = Field(default=8192, description="Maximum tokens in response")
-    temperature: float = Field(
-        default=0.2, description="Temperature for response generation", ge=0, le=1
+    conversation_history: List[Dict[str, str]] = Field(
+        default_factory=list,
+        description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
+    )
+    model: str = Field(
+        default="gpt-4o",
+        description="OpenAI model to use",
     )
-    images: Optional[List[Path]] = Field(
-        default=None,
-        description="Optional list of image paths to include in the message",
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation", ge=0, le=1
     )
-    functions: Optional[List[Dict[str, Any]]] = Field(
-        default=None,
-        description="Optional function definitions for function calling",
+    max_tokens: Optional[int] = Field(
+        default=None, description="Maximum tokens in response"
     )
-    function_call: Optional[str] = Field(
-        default=None,
-        description="Controls function calling behavior",
+    stream: bool = Field(
+        default=False,
+        description="Whether to stream the response token by token",
     )
 
 
 class OpenAIOutput(OutputSchema):
     """Schema for OpenAI chat output"""
 
-    response: str = Field(..., description="Model's response text")
-    used_model: str = Field(..., description="Model used for generation")
-    usage: Dict[str, Any] = Field(
-        default_factory=dict, description="Usage statistics from the API"
-    )
-    function_call: Optional[Dict[str, Any]] = Field(
-        default=None, description="Function call information if applicable"
-    )
+    response: str
+    used_model: str
+    usage: Dict[str, int]
 
 
 class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
-    """Skill for interacting with OpenAI's models"""
+    """Skill for interacting with OpenAI models"""
 
     input_schema = OpenAIInput
     output_schema = OpenAIOutput
@@ -65,82 +62,69 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
             organization=self.credentials.openai_organization_id,
         )
 
-    def _encode_image(self, image_path: Path) -> Dict[str, Any]:
-        """Convert image to base64 for API consumption"""
-        try:
-            if not image_path.exists():
-                raise FileNotFoundError(f"Image file not found: {image_path}")
-
-            with open(image_path, "rb") as img_file:
-                encoded = base64.b64encode(img_file.read()).decode()
-                return {
-                    "type": "image_url",
-                    "image_url": {"url": f"data:image/jpeg;base64,{encoded}"},
-                }
-        except Exception as e:
-            logger.error(f"Failed to encode image {image_path}: {str(e)}")
-            raise ProcessingError(f"Image encoding failed: {str(e)}")
+    def _build_messages(self, input_data: OpenAIInput) -> List[Dict[str, str]]:
+        """Build messages list from input data including conversation history."""
+        messages = [{"role": "system", "content": input_data.system_prompt}]
 
-    def process(self, input_data: OpenAIInput) -> OpenAIOutput:
-        """Process the input using OpenAI's API"""
-        try:
-            logger.info(f"Processing request with model {input_data.model}")
+        if input_data.conversation_history:
+            messages.extend(input_data.conversation_history)
 
-            # Prepare message content
-            content = []
+        messages.append({"role": "user", "content": input_data.user_input})
+        return messages
 
-            # Add text content
-            content.append({"type": "text", "text": input_data.user_input})
-
-            # Add images if provided
-            if input_data.images:
-                logger.debug(f"Processing {len(input_data.images)} images")
-                for image_path in input_data.images:
-                    content.append(self._encode_image(image_path))
-
-            # Prepare messages
-            messages = [
-                {"role": "system", "content": input_data.system_prompt},
-                {"role": "user", "content": content},
-            ]
-
-            # Create completion parameters
-            params = {
-                "model": input_data.model,
-                "messages": messages,
-                "temperature": input_data.temperature,
-                "max_tokens": input_data.max_tokens,
-            }
+    def process_stream(self, input_data: OpenAIInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
+        try:
+            messages = self._build_messages(input_data)
 
-            # Add function calling if provided
-            if input_data.functions:
-                params["functions"] = input_data.functions
-                params["function_call"] = input_data.function_call
+            stream = self.client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+                stream=True,
+            )
 
-            # Create chat completion
-            response = self.client.chat.completions.create(**params)
+            for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
 
-            # Extract function call if present
-            function_call = None
-            if response.choices[0].message.function_call:
-                function_call = response.choices[0].message.function_call.model_dump()
+        except Exception as e:
+            raise ProcessingError(f"OpenAI streaming failed: {str(e)}")
 
-            logger.success("Successfully processed OpenAI request")
+    def process(self, input_data: OpenAIInput) -> OpenAIOutput:
+        """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)
+            else:
+                # For non-streaming, use regular completion
+                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,
+                )
+                response = completion.choices[0].message.content
 
             return OpenAIOutput(
-                response=response.choices[0].message.content or "",
-                used_model=response.model,
+                response=response,
+                used_model=input_data.model,
                 usage={
-                    "prompt_tokens": response.usage.prompt_tokens,
-                    "completion_tokens": response.usage.completion_tokens,
-                    "total_tokens": response.usage.total_tokens,
+                    "total_tokens": completion.usage.total_tokens,
+                    "prompt_tokens": completion.usage.prompt_tokens,
+                    "completion_tokens": completion.usage.completion_tokens,
                 },
-                function_call=function_call,
             )
 
         except Exception as e:
-            logger.exception(f"OpenAI processing failed: {str(e)}")
-            raise ProcessingError(f"OpenAI processing failed: {str(e)}")
+            raise ProcessingError(f"OpenAI chat failed: {str(e)}")
 
 
 ResponseT = TypeVar("ResponseT", bound=BaseModel)

{airtrain-0.1.26.dist-info → airtrain-0.1.28.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.26
+Version: 0.1.28
 Summary: A platform for building and deploying AI agents with structured skills
 Home-page: https://github.com/rosaboyle/airtrain.dev
 Author: Dheeraj Pai
@@ -171,7 +171,21 @@ This project is licensed under the MIT License - see the LICENSE file for detail
 ## Changelog
 
 
-## 0.1.32
+
+## 0.1.28
+
+- Bug fix: reasoning to Fireworks structured output.
+- Added reasoning to Fireworks structured output.
+
+## 0.1.27
+
+- Added structured completion skills for Fireworks AI
+- Added Completion skills for Fireworks AI.
+- Added Combination skill for Groq and Fireworks AI.
+- Add completion streaming.
+- Added strcutured output streaming for Fireworks AI.
+
+## 0.1.23
 
 - Added conversation support for Deepseek, Togehter AI, Fireworks AI, Gemini, Groq, Cerebras and Sambanova.
 - Added Change Log

{airtrain-0.1.26.dist-info → airtrain-0.1.28.dist-info}/RECORD

@@ -1,4 +1,4 @@
-airtrain/__init__.py,sha256=yB_7KwsrnELZUpWHGM8FMlOeY_7cSWHwGk1e8eZT76Y,2099
+airtrain/__init__.py,sha256=ZUseQpkG-otzfnOZRjQZHpJYG6w2R-c-k872-UMaRqI,2099
 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
@@ -18,10 +18,15 @@ airtrain/integrations/cerebras/__init__.py,sha256=zAD-qV38OzHhMCz1z-NvjjqcYEhURb
 airtrain/integrations/cerebras/credentials.py,sha256=KDEH4r8FGT68L9p34MLZWK65wq_a703pqIF3ODaSbts,694
 airtrain/integrations/cerebras/skills.py,sha256=Ksggq_s5wHWlf_xQIOO8MFoNTYV0cm9SHZ7GESOd2YE,3527
 airtrain/integrations/fireworks/__init__.py,sha256=9pJvP0u1FJbNtB0oHa09mHVJLctELf_c27LOYyDk2ZI,271
+airtrain/integrations/fireworks/completion_skills.py,sha256=G657xWd7izLOxXq7RmqdupBF4DHqXQgXuhQ-MW7mtqc,5613
 airtrain/integrations/fireworks/conversation_manager.py,sha256=m6VEHijqpYEYawkKhuHtb8RQxw4kxGWFWdbSK6zGuro,3704
 airtrain/integrations/fireworks/credentials.py,sha256=UpcwR9V5Hbk5sJbjFDJDbHMRqc90IQSqAvrtJCOvwEo,524
 airtrain/integrations/fireworks/models.py,sha256=F-MddbLCLAsTjwRr1l6IpJxOegyY4pD7jN9ySPiypSo,593
-airtrain/integrations/fireworks/skills.py,sha256=-6zfe5eooygDYL1cyrk4PUT2TLEmAT2ZRfH2cr42DaU,4960
+airtrain/integrations/fireworks/requests_skills.py,sha256=zPIR70l0KdSGmA5WyDEopFAdKHSPltAttJBzWyHu6Bk,5878
+airtrain/integrations/fireworks/skills.py,sha256=OB4epD4CSTxExUCC1oMJ_8rHLOoftlxf0FUoIVrd4mA,5163
+airtrain/integrations/fireworks/structured_completion_skills.py,sha256=IXG4gsZDSfuscrmKIHfnyHkBaCV7zlPInaWXb95iC5k,6428
+airtrain/integrations/fireworks/structured_requests_skills.py,sha256=5FllptmuUewS1OX3Jf0vK6kQjMV4QJmisExbs0ElWPY,8150
+airtrain/integrations/fireworks/structured_skills.py,sha256=BZaLqSOTC11QdZ4kDORS4JnwF_YXBAa-IiwQ5dJiHXw,3895
 airtrain/integrations/google/__init__.py,sha256=ElwgcXfbg_gGMm6zbkMXCQPFKZUb-yTJk986o19A7Cs,214
 airtrain/integrations/google/credentials.py,sha256=KSvWNqW8Mjr4MkysRvUqlrOSGdShNIe5u2OPO6vRrWY,2047
 airtrain/integrations/google/skills.py,sha256=ytsoksCY4qbfRO9Brnxhc2694fAj0ytnHX20SXS_FOM,4547
@@ -35,7 +40,7 @@ airtrain/integrations/openai/__init__.py,sha256=K-NY2_T1T6SEOgkpbUA55cWvK2nr2NOJ
 airtrain/integrations/openai/chinese_assistant.py,sha256=MMhv4NBOoEQ0O22ZZtP255rd5ajHC9l6FPWIjpqxBOA,1581
 airtrain/integrations/openai/credentials.py,sha256=NfRyp1QgEtgm8cxt2-BOLq-6d0X-Pcm80NnfHM8p0FY,1470
 airtrain/integrations/openai/models_config.py,sha256=bzosqqpDy2AJxu2vGdk2H4voqEGlv7LORR6fpJLhNic,3962
-airtrain/integrations/openai/skills.py,sha256=Olg9-6f_p2XgkVwwcB9tvjAMApmM2EK81i8LP4qVVvs,7676
+airtrain/integrations/openai/skills.py,sha256=SEWpwfWPsDEmf7IcSzWZuEzb13YT4gzQIHO8_O1bmN4,6936
 airtrain/integrations/sambanova/__init__.py,sha256=dp_263iOckM_J9pOEvyqpf3FrejD6-_x33r0edMCTe0,179
 airtrain/integrations/sambanova/credentials.py,sha256=JyN8sbMCoXuXAjim46aI3LTicBijoemS7Ao0rn4yBJU,824
 airtrain/integrations/sambanova/skills.py,sha256=SDFY-ZzhOEIxQgTkQJzX9gN7UDqqnCBJdK7I2JydIoY,3625
@@ -52,7 +57,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=mUoHc2r5TYQi5iGzwz2aDuUeROGq7teCtNrOlNApef4,6276
 airtrain/integrations/together/vision_models_config.py,sha256=m28HwYDk2Kup_J-a1FtynIa2ZVcbl37kltfoHnK8zxs,1544
-airtrain-0.1.26.dist-info/METADATA,sha256=MCoC4vRc3ZcGrwJEJgWhn2qULmOzwZ7MIWx9eblF23s,4882
-airtrain-0.1.26.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-airtrain-0.1.26.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
-airtrain-0.1.26.dist-info/RECORD,,
+airtrain-0.1.28.dist-info/METADATA,sha256=svDrTc0qzBYsuYKGR0VKg7AZSIEMQpUht1M2csmWGH0,5243
+airtrain-0.1.28.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+airtrain-0.1.28.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.28.dist-info/RECORD,,