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

airtrain/integrations/fireworks/structured_requests_skills.py

@@ -0,0 +1,291 @@
+from typing import Type, TypeVar, Optional, List, Dict, Any, Generator, Union
+from pydantic import BaseModel, Field, create_model
+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, 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: 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"
+    )
+    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 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
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        default=None, description="Tool calls generated by the model"
+    )
+
+
+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, 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 _build_payload(
+        self, input_data: FireworksStructuredRequestInput
+    ) -> Dict[str, Any]:
+        """Build the request payload."""
+        payload = {
+            "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"},
+        }
+
+        # Add tool-related parameters if provided
+        if input_data.tools:
+            payload["tools"] = input_data.tools
+        
+        if input_data.tool_choice:
+            payload["tool_choice"] = input_data.tool_choice
+
+        return payload
+
+    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 response with think tags
+            if not json_buffer:
+                # If no data was collected, raise error
+                raise ProcessingError("No data received from Fireworks API")
+
+            complete_response = "".join(json_buffer)
+            reasoning, json_str = self._parse_response_content(complete_response)
+
+            try:
+                parsed_response = input_data.response_model.model_validate_json(
+                    json_str
+                )
+                yield {"complete": parsed_response, "reasoning": reasoning}
+            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
+                reasoning = 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"]
+                        reasoning = chunk.get("reasoning")
+
+                if parsed_response is None:
+                    raise ProcessingError("Failed to parse streamed response")
+                
+                # Make a non-streaming call to get tool calls if tools were provided
+                tool_calls = None
+                if input_data.tools:
+                    # Create a non-streaming request to get tool calls
+                    non_stream_payload = self._build_payload(input_data)
+                    non_stream_payload["stream"] = False
+                    
+                    response = requests.post(
+                        self.BASE_URL,
+                        headers=self.headers,
+                        data=json.dumps(non_stream_payload),
+                    )
+                    response.raise_for_status()
+                    result = response.json()
+                    
+                    # Check for tool calls
+                    if (result["choices"][0]["message"].get("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 result["choices"][0]["message"]["tool_calls"]
+                        ]
+
+                return FireworksStructuredRequestOutput(
+                    parsed_response=parsed_response,
+                    used_model=input_data.model,
+                    usage={"total_tokens": 0},  # Can't get usage stats from streaming
+                    reasoning=reasoning,
+                    tool_calls=tool_calls,
+                )
+            else:
+                # For non-streaming, use regular request
+                payload = self._build_payload(input_data)
+                payload["stream"] = False  # Ensure it's not streaming
+
+                response = requests.post(
+                    self.BASE_URL, headers=self.headers, data=json.dumps(payload)
+                )
+                response.raise_for_status()
+                result = response.json()
+
+                # Get the content from the response
+                if "choices" not in result or not result["choices"]:
+                    raise ProcessingError("Invalid response format from Fireworks API")
+
+                content = result["choices"][0]["message"].get("content", "")
+                
+                # Check for tool calls
+                tool_calls = None
+                if (result["choices"][0]["message"].get("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 result["choices"][0]["message"]["tool_calls"]
+                    ]
+
+                # Parse the response content
+                reasoning, json_str = self._parse_response_content(content)
+                try:
+                    parsed_response = input_data.response_model.model_validate_json(
+                        json_str
+                    )
+                except Exception as e:
+                    raise ProcessingError(f"Failed to parse JSON response: {str(e)}")
+
+                return FireworksStructuredRequestOutput(
+                    parsed_response=parsed_response,
+                    used_model=input_data.model,
+                    usage={
+                        "total_tokens": result["usage"]["total_tokens"],
+                        "prompt_tokens": result["usage"]["prompt_tokens"],
+                        "completion_tokens": result["usage"]["completion_tokens"],
+                    },
+                    reasoning=reasoning,
+                    tool_calls=tool_calls,
+                )
+
+        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] = 131072
+    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/google/__init__.py

@@ -0,0 +1,7 @@
+"""Google Cloud integration module"""
+
+from .credentials import GoogleCloudCredentials
+from .skills import GoogleChatSkill
+# from .skills import VertexAISkill
+
+__all__ = ["GoogleCloudCredentials", "VertexAISkill"]

airtrain/integrations/google/credentials.py

@@ -0,0 +1,58 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+import google.genai as genai
+from google.cloud import storage
+import os
+
+# from google.cloud import storage
+
+
+class GoogleCloudCredentials(BaseCredentials):
+    """Google Cloud credentials"""
+
+    project_id: str = Field(..., description="Google Cloud Project ID")
+    service_account_key: SecretStr = Field(..., description="Service Account Key JSON")
+
+    _required_credentials = {"project_id", "service_account_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Google Cloud credentials"""
+        try:
+            # Initialize with service account key
+            storage_client = storage.Client.from_service_account_info(
+                self.service_account_key.get_secret_value()
+            )
+            # Test API call
+            storage_client.list_buckets(max_results=1)
+            return True
+        except Exception as e:
+            raise CredentialValidationError(
+                f"Invalid Google Cloud credentials: {str(e)}"
+            )
+
+
+class GeminiCredentials(BaseCredentials):
+    """Gemini API credentials"""
+
+    gemini_api_key: SecretStr = Field(..., description="Gemini API Key")
+
+    _required_credentials = {"gemini_api_key"}
+
+    @classmethod
+    def from_env(cls) -> "GeminiCredentials":
+        """Create credentials from environment variables"""
+        return cls(gemini_api_key=SecretStr(os.environ.get("GEMINI_API_KEY", "")))
+
+    async def validate_credentials(self) -> bool:
+        """Validate Gemini API credentials"""
+        try:
+            # Configure Gemini with API key
+            genai.configure(api_key=self.gemini_api_key.get_secret_value())
+
+            # Test API call with a simple model
+            model = genai.GenerativeModel("gemini-1.5-flash")
+            response = model.generate_content("test")
+
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid Gemini credentials: {str(e)}")

airtrain/integrations/google/skills.py

@@ -0,0 +1,122 @@
+from typing import List, Optional, Dict, Any
+from pydantic import Field
+import google.generativeai as genai
+from loguru import logger
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import GeminiCredentials
+
+
+class GoogleGenerationConfig(InputSchema):
+    """Schema for Google generation config"""
+
+    temperature: float = Field(
+        default=1.0, description="Temperature for response generation", ge=0, le=1
+    )
+    top_p: float = Field(
+        default=0.95, description="Top p sampling parameter", ge=0, le=1
+    )
+    top_k: int = Field(default=40, description="Top k sampling parameter")
+    max_output_tokens: int = Field(
+        default=8192, description="Maximum tokens in response"
+    )
+    response_mime_type: str = Field(
+        default="text/plain", description="Response MIME type"
+    )
+
+
+class GoogleInput(InputSchema):
+    """Schema for Google 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, str | List[Dict[str, str]]]] = Field(
+        default_factory=list,
+        description="List of conversation messages in Google's format",
+    )
+    model: str = Field(default="gemini-1.5-flash", description="Google model to use")
+    generation_config: GoogleGenerationConfig = Field(
+        default_factory=GoogleGenerationConfig, description="Generation configuration"
+    )
+
+
+class GoogleOutput(OutputSchema):
+    """Schema for Google 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")
+
+
+class GoogleChatSkill(Skill[GoogleInput, GoogleOutput]):
+    """Skill for Google chat"""
+
+    input_schema = GoogleInput
+    output_schema = GoogleOutput
+
+    def __init__(self, credentials: Optional[GeminiCredentials] = None):
+        super().__init__()
+        self.credentials = credentials or GeminiCredentials.from_env()
+        genai.configure(api_key=self.credentials.gemini_api_key.get_secret_value())
+
+    def _convert_history_format(
+        self, history: List[Dict[str, str]]
+    ) -> List[Dict[str, List[Dict[str, str]]]]:
+        """Convert standard history format to Google's format"""
+        google_history = []
+        for msg in history:
+            google_msg = {
+                "role": "user" if msg["role"] == "user" else "model",
+                "parts": [{"text": msg["content"]}],
+            }
+            google_history.append(google_msg)
+        return google_history
+
+    def process(self, input_data: GoogleInput) -> GoogleOutput:
+        try:
+            # Create generation config
+            generation_config = {
+                "temperature": input_data.generation_config.temperature,
+                "top_p": input_data.generation_config.top_p,
+                "top_k": input_data.generation_config.top_k,
+                "max_output_tokens": input_data.generation_config.max_output_tokens,
+                "response_mime_type": input_data.generation_config.response_mime_type,
+            }
+
+            # Initialize model
+            model = genai.GenerativeModel(
+                model_name=input_data.model,
+                generation_config=generation_config,
+                system_instruction=input_data.system_prompt,
+            )
+
+            # Convert history format if needed
+            history = (
+                input_data.conversation_history
+                if input_data.conversation_history
+                else self._convert_history_format([])
+            )
+
+            # Start chat session
+            chat = model.start_chat(history=history)
+
+            # Send message and get response
+            response = chat.send_message(input_data.user_input)
+
+            return GoogleOutput(
+                response=response.text,
+                used_model=input_data.model,
+                usage={
+                    "prompt_tokens": 0,
+                    "completion_tokens": 0,
+                    "total_tokens": 0,
+                },  # Google API doesn't provide usage stats
+            )
+
+        except Exception as e:
+            logger.exception(f"Google processing failed: {str(e)}")
+            raise ProcessingError(f"Google processing failed: {str(e)}")

airtrain/integrations/groq/__init__.py

@@ -0,0 +1,23 @@
+"""Groq integration module"""
+
+from .credentials import GroqCredentials
+from .skills import GroqChatSkill
+from .models_config import (
+    get_model_config,
+    get_default_model,
+    supports_tool_use,
+    supports_parallel_tool_use,
+    supports_json_mode,
+    GROQ_MODELS_CONFIG,
+)
+
+__all__ = [
+    "GroqCredentials", 
+    "GroqChatSkill",
+    "get_model_config",
+    "get_default_model",
+    "supports_tool_use",
+    "supports_parallel_tool_use",
+    "supports_json_mode",
+    "GROQ_MODELS_CONFIG",
+]

airtrain/integrations/groq/credentials.py

@@ -0,0 +1,24 @@
+from pydantic import Field, SecretStr
+from airtrain.core.credentials import BaseCredentials, CredentialValidationError
+from groq import Groq
+
+
+class GroqCredentials(BaseCredentials):
+    """Groq API credentials"""
+
+    groq_api_key: SecretStr = Field(..., description="Groq API key")
+
+    _required_credentials = {"groq_api_key"}
+
+    async def validate_credentials(self) -> bool:
+        """Validate Groq credentials"""
+        try:
+            client = Groq(api_key=self.groq_api_key.get_secret_value())
+            await client.chat.completions.create(
+                messages=[{"role": "user", "content": "Hi"}],
+                model="mixtral-8x7b-32768",
+                max_tokens=1,
+            )
+            return True
+        except Exception as e:
+            raise CredentialValidationError(f"Invalid Groq credentials: {str(e)}")