airtrain 0.1.25__tar.gz → 0.1.27__tar.gz

{airtrain-0.1.25/airtrain.egg-info → airtrain-0.1.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.25
+Version: 0.1.27
 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,15 @@ This project is licensed under the MIT License - see the LICENSE file for detail
 ## Changelog
 
 
-## 0.1.15
+## 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.25 → airtrain-0.1.27}/airtrain/__init__.py

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

airtrain-0.1.27/airtrain/integrations/combined/groq_fireworks_skills.py

@@ -0,0 +1,126 @@
+from typing import Optional, Dict, Any, List
+from pydantic import Field
+import requests
+from groq import Groq
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from airtrain.integrations.fireworks.completion_skills import (
+    FireworksCompletionSkill,
+    FireworksCompletionInput,
+)
+
+
+class GroqFireworksInput(InputSchema):
+    """Schema for combined Groq and Fireworks input"""
+
+    user_input: str = Field(..., description="User's input text")
+    groq_model: str = Field(
+        default="mixtral-8x7b-32768", description="Groq model to use"
+    )
+    fireworks_model: str = Field(
+        default="accounts/fireworks/models/deepseek-r1",
+        description="Fireworks model to use",
+    )
+    temperature: float = Field(
+        default=0.7, description="Temperature for response generation"
+    )
+    max_tokens: int = Field(default=4096, description="Maximum tokens in response")
+
+
+class GroqFireworksOutput(OutputSchema):
+    """Schema for combined Groq and Fireworks output"""
+
+    combined_response: str
+    groq_response: str
+    fireworks_response: str
+    used_models: Dict[str, str]
+    usage: Dict[str, Dict[str, int]]
+
+
+class GroqFireworksSkill(Skill[GroqFireworksInput, GroqFireworksOutput]):
+    """Skill combining Groq and Fireworks responses"""
+
+    input_schema = GroqFireworksInput
+    output_schema = GroqFireworksOutput
+
+    def __init__(
+        self,
+        groq_api_key: Optional[str] = None,
+        fireworks_skill: Optional[FireworksCompletionSkill] = None,
+    ):
+        """Initialize the skill with optional API keys"""
+        super().__init__()
+        self.groq_client = Groq(api_key=groq_api_key)
+        self.fireworks_skill = fireworks_skill or FireworksCompletionSkill()
+
+    def _get_groq_response(self, input_data: GroqFireworksInput) -> Dict[str, Any]:
+        """Get response from Groq"""
+        try:
+            completion = self.groq_client.chat.completions.create(
+                model=input_data.groq_model,
+                messages=[{"role": "user", "content": input_data.user_input}],
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+            return {
+                "response": completion.choices[0].message.content,
+                "usage": completion.usage.model_dump(),
+            }
+        except Exception as e:
+            raise ProcessingError(f"Groq request failed: {str(e)}")
+
+    def _get_fireworks_response(
+        self, groq_response: str, input_data: GroqFireworksInput
+    ) -> Dict[str, Any]:
+        """Get response from Fireworks"""
+        try:
+            formatted_prompt = (
+                f"<USER>{input_data.user_input}</USER>\n<ASSISTANT>{groq_response}"
+            )
+
+            fireworks_input = FireworksCompletionInput(
+                prompt=formatted_prompt,
+                model=input_data.fireworks_model,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+
+            result = self.fireworks_skill.process(fireworks_input)
+            return {"response": result.response, "usage": result.usage}
+        except Exception as e:
+            raise ProcessingError(f"Fireworks request failed: {str(e)}")
+
+    def process(self, input_data: GroqFireworksInput) -> GroqFireworksOutput:
+        """Process the input using both Groq and Fireworks"""
+        try:
+            # Get Groq response
+            groq_result = self._get_groq_response(input_data)
+
+            # Get Fireworks response
+            fireworks_result = self._get_fireworks_response(
+                groq_result["response"], input_data
+            )
+
+            # Combine responses in the required format
+            combined_response = (
+                f"<USER>{input_data.user_input}</USER>\n"
+                f"<ASSISTANT>{groq_result['response']} {fireworks_result['response']}"
+            )
+
+            return GroqFireworksOutput(
+                combined_response=combined_response,
+                groq_response=groq_result["response"],
+                fireworks_response=fireworks_result["response"],
+                used_models={
+                    "groq": input_data.groq_model,
+                    "fireworks": input_data.fireworks_model,
+                },
+                usage={
+                    "groq": groq_result["usage"],
+                    "fireworks": fireworks_result["usage"],
+                },
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Combined processing failed: {str(e)}")

airtrain-0.1.27/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-0.1.27/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-0.1.25 → airtrain-0.1.27}/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)}")