airtrain 0.1.39__py3-none-any.whl → 0.1.41__py3-none-any.whl

airtrain/integrations/fireworks/structured_requests_skills.py

@@ -1,5 +1,5 @@
-from typing import Type, TypeVar, Optional, List, Dict, Any, Generator
-from pydantic import BaseModel, Field
+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
@@ -20,7 +20,7 @@ class FireworksStructuredRequestInput(InputSchema):
         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(
+    conversation_history: List[Dict[str, Any]] = Field(
         default_factory=list,
         description="List of previous conversation messages",
     )
@@ -34,8 +34,21 @@ class FireworksStructuredRequestInput(InputSchema):
     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",
+        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:
@@ -49,6 +62,9 @@ class FireworksStructuredRequestOutput(OutputSchema):
     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(
@@ -72,7 +88,7 @@ class FireworksStructuredRequestSkill(
 
     def _build_messages(
         self, input_data: FireworksStructuredRequestInput
-    ) -> List[Dict[str, str]]:
+    ) -> List[Dict[str, Any]]:
         """Build messages list from input data including conversation history."""
         messages = [{"role": "system", "content": input_data.system_prompt}]
 
@@ -86,24 +102,24 @@ class FireworksStructuredRequestSkill(
         self, input_data: FireworksStructuredRequestInput
     ) -> Dict[str, Any]:
         """Build the request payload."""
-        return {
+        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",
-                "schema": {
-                    **input_data.response_model.model_json_schema(),
-                    "required": [
-                        field
-                        for field, _ in input_data.response_model.model_fields.items()
-                    ],
-                },
-            },
+            "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]:
@@ -131,6 +147,10 @@ class FireworksStructuredRequestSkill(
                         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)
 
@@ -177,37 +197,94 @@ class FireworksStructuredRequestSkill(
 
                 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={},  # Usage stats not available in streaming mode
+                    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()
-                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
-                )
+                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=data["usage"],
-                    reasoning=reasoning,  # Add reasoning to output if present
+                    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:

airtrain/integrations/openai/__init__.py

@@ -5,6 +5,9 @@ from .skills import (
     OpenAIOutput,
     OpenAIParserInput,
     OpenAIParserOutput,
+    OpenAIEmbeddingsSkill,
+    OpenAIEmbeddingsInput,
+    OpenAIEmbeddingsOutput,
 )
 from .credentials import OpenAICredentials
 
@@ -16,4 +19,7 @@ __all__ = [
     "OpenAIParserOutput",
     "OpenAICredentials",
     "OpenAIOutput",
+    "OpenAIEmbeddingsSkill",
+    "OpenAIEmbeddingsInput",
+    "OpenAIEmbeddingsOutput",
 ]

airtrain/integrations/openai/models_config.py

@@ -11,6 +11,20 @@ class OpenAIModelConfig(NamedTuple):
 
 
 OPENAI_MODELS: Dict[str, OpenAIModelConfig] = {
+    "gpt-4.5-preview": OpenAIModelConfig(
+        display_name="GPT-4.5 Preview",
+        base_model="gpt-4.5-preview",
+        input_price=Decimal("75.00"),
+        cached_input_price=Decimal("37.50"),
+        output_price=Decimal("150.00"),
+    ),
+    "gpt-4.5-preview-2025-02-27": OpenAIModelConfig(
+        display_name="GPT-4.5 Preview (2025-02-27)",
+        base_model="gpt-4.5-preview",
+        input_price=Decimal("75.00"),
+        cached_input_price=Decimal("37.50"),
+        output_price=Decimal("150.00"),
+    ),
     "gpt-4o": OpenAIModelConfig(
         display_name="GPT-4 Optimized",
         base_model="gpt-4o",
@@ -25,69 +39,160 @@ OPENAI_MODELS: Dict[str, OpenAIModelConfig] = {
         cached_input_price=Decimal("1.25"),
         output_price=Decimal("10.00"),
     ),
-    "gpt-4o-2024-05-13": OpenAIModelConfig(
-        display_name="GPT-4 Optimized (2024-05-13)",
-        base_model="gpt-4o",
-        input_price=Decimal("5.00"),
+    "gpt-4o-audio-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Audio Preview",
+        base_model="gpt-4o-audio-preview",
+        input_price=Decimal("2.50"),
         cached_input_price=None,
-        output_price=Decimal("15.00"),
+        output_price=Decimal("10.00"),
     ),
     "gpt-4o-audio-preview-2024-12-17": OpenAIModelConfig(
-        display_name="GPT-4 Optimized Audio Preview",
+        display_name="GPT-4 Optimized Audio Preview (2024-12-17)",
         base_model="gpt-4o-audio-preview",
         input_price=Decimal("2.50"),
         cached_input_price=None,
         output_price=Decimal("10.00"),
     ),
-    "gpt-4o-realtime-preview-2024-12-17": OpenAIModelConfig(
+    "gpt-4o-realtime-preview": OpenAIModelConfig(
         display_name="GPT-4 Optimized Realtime Preview",
         base_model="gpt-4o-realtime-preview",
         input_price=Decimal("5.00"),
         cached_input_price=Decimal("2.50"),
         output_price=Decimal("20.00"),
     ),
-    "gpt-4o-mini-2024-07-18": OpenAIModelConfig(
+    "gpt-4o-realtime-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Realtime Preview (2024-12-17)",
+        base_model="gpt-4o-realtime-preview",
+        input_price=Decimal("5.00"),
+        cached_input_price=Decimal("2.50"),
+        output_price=Decimal("20.00"),
+    ),
+    "gpt-4o-mini": OpenAIModelConfig(
         display_name="GPT-4 Optimized Mini",
         base_model="gpt-4o-mini",
         input_price=Decimal("0.15"),
         cached_input_price=Decimal("0.075"),
         output_price=Decimal("0.60"),
     ),
-    "gpt-4o-mini-audio-preview-2024-12-17": OpenAIModelConfig(
+    "gpt-4o-mini-2024-07-18": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini (2024-07-18)",
+        base_model="gpt-4o-mini",
+        input_price=Decimal("0.15"),
+        cached_input_price=Decimal("0.075"),
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-mini-audio-preview": OpenAIModelConfig(
         display_name="GPT-4 Optimized Mini Audio Preview",
         base_model="gpt-4o-mini-audio-preview",
         input_price=Decimal("0.15"),
         cached_input_price=None,
         output_price=Decimal("0.60"),
     ),
-    "gpt-4o-mini-realtime-preview-2024-12-17": OpenAIModelConfig(
+    "gpt-4o-mini-audio-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Audio Preview (2024-12-17)",
+        base_model="gpt-4o-mini-audio-preview",
+        input_price=Decimal("0.15"),
+        cached_input_price=None,
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-mini-realtime-preview": OpenAIModelConfig(
         display_name="GPT-4 Optimized Mini Realtime Preview",
         base_model="gpt-4o-mini-realtime-preview",
         input_price=Decimal("0.60"),
         cached_input_price=Decimal("0.30"),
         output_price=Decimal("2.40"),
     ),
-    "o1-2024-12-17": OpenAIModelConfig(
+    "gpt-4o-mini-realtime-preview-2024-12-17": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Realtime Preview (2024-12-17)",
+        base_model="gpt-4o-mini-realtime-preview",
+        input_price=Decimal("0.60"),
+        cached_input_price=Decimal("0.30"),
+        output_price=Decimal("2.40"),
+    ),
+    "o1": OpenAIModelConfig(
         display_name="O1",
         base_model="o1",
         input_price=Decimal("15.00"),
         cached_input_price=Decimal("7.50"),
         output_price=Decimal("60.00"),
     ),
-    "o3-mini-2025-01-31": OpenAIModelConfig(
+    "o1-2024-12-17": OpenAIModelConfig(
+        display_name="O1 (2024-12-17)",
+        base_model="o1",
+        input_price=Decimal("15.00"),
+        cached_input_price=Decimal("7.50"),
+        output_price=Decimal("60.00"),
+    ),
+    "o3-mini": OpenAIModelConfig(
         display_name="O3 Mini",
         base_model="o3-mini",
         input_price=Decimal("1.10"),
         cached_input_price=Decimal("0.55"),
         output_price=Decimal("4.40"),
     ),
-    "o1-mini-2024-09-12": OpenAIModelConfig(
+    "o3-mini-2025-01-31": OpenAIModelConfig(
+        display_name="O3 Mini (2025-01-31)",
+        base_model="o3-mini",
+        input_price=Decimal("1.10"),
+        cached_input_price=Decimal("0.55"),
+        output_price=Decimal("4.40"),
+    ),
+    "o1-mini": OpenAIModelConfig(
         display_name="O1 Mini",
         base_model="o1-mini",
         input_price=Decimal("1.10"),
         cached_input_price=Decimal("0.55"),
         output_price=Decimal("4.40"),
     ),
+    "o1-mini-2024-09-12": OpenAIModelConfig(
+        display_name="O1 Mini (2024-09-12)",
+        base_model="o1-mini",
+        input_price=Decimal("1.10"),
+        cached_input_price=Decimal("0.55"),
+        output_price=Decimal("4.40"),
+    ),
+    "gpt-4o-mini-search-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Search Preview",
+        base_model="gpt-4o-mini-search-preview",
+        input_price=Decimal("0.15"),
+        cached_input_price=None,
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-mini-search-preview-2025-03-11": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Mini Search Preview (2025-03-11)",
+        base_model="gpt-4o-mini-search-preview",
+        input_price=Decimal("0.15"),
+        cached_input_price=None,
+        output_price=Decimal("0.60"),
+    ),
+    "gpt-4o-search-preview": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Search Preview",
+        base_model="gpt-4o-search-preview",
+        input_price=Decimal("2.50"),
+        cached_input_price=None,
+        output_price=Decimal("10.00"),
+    ),
+    "gpt-4o-search-preview-2025-03-11": OpenAIModelConfig(
+        display_name="GPT-4 Optimized Search Preview (2025-03-11)",
+        base_model="gpt-4o-search-preview",
+        input_price=Decimal("2.50"),
+        cached_input_price=None,
+        output_price=Decimal("10.00"),
+    ),
+    "computer-use-preview": OpenAIModelConfig(
+        display_name="Computer Use Preview",
+        base_model="computer-use-preview",
+        input_price=Decimal("3.00"),
+        cached_input_price=None,
+        output_price=Decimal("12.00"),
+    ),
+    "computer-use-preview-2025-03-11": OpenAIModelConfig(
+        display_name="Computer Use Preview (2025-03-11)",
+        base_model="computer-use-preview",
+        input_price=Decimal("3.00"),
+        cached_input_price=None,
+        output_price=Decimal("12.00"),
+    ),
 }
 
 

airtrain/integrations/openai/skills.py

@@ -1,7 +1,8 @@
-from typing import AsyncGenerator, List, Optional, Dict, TypeVar, Type, Generator
+from typing import AsyncGenerator, List, Optional, Dict, TypeVar, Type, Generator, Union
 from pydantic import Field, BaseModel
 from openai import OpenAI, AsyncOpenAI
 from openai.types.chat import ChatCompletionChunk
+import numpy as np
 
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
@@ -232,3 +233,110 @@ class OpenAIParserSkill(Skill[OpenAIParserInput, OpenAIParserOutput]):
 
         except Exception as e:
             raise ProcessingError(f"OpenAI parsing failed: {str(e)}")
+
+
+class OpenAIEmbeddingsInput(InputSchema):
+    """Schema for OpenAI embeddings input"""
+
+    texts: Union[str, List[str]] = Field(
+        ..., description="Text or list of texts to generate embeddings for"
+    )
+    model: str = Field(
+        default="text-embedding-3-large", description="OpenAI embeddings model to use"
+    )
+    encoding_format: str = Field(
+        default="float", description="The format of the embeddings: 'float' or 'base64'"
+    )
+    dimensions: Optional[int] = Field(
+        default=None, description="Optional number of dimensions for the embeddings"
+    )
+
+
+class OpenAIEmbeddingsOutput(OutputSchema):
+    """Schema for OpenAI embeddings output"""
+
+    embeddings: List[List[float]] = Field(..., description="List of embeddings vectors")
+    used_model: str = Field(..., description="Model used for generating embeddings")
+    tokens_used: int = Field(..., description="Number of tokens used")
+
+
+class OpenAIEmbeddingsSkill(Skill[OpenAIEmbeddingsInput, OpenAIEmbeddingsOutput]):
+    """Skill for generating embeddings using OpenAI models"""
+
+    input_schema = OpenAIEmbeddingsInput
+    output_schema = OpenAIEmbeddingsOutput
+
+    def __init__(self, credentials: Optional[OpenAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or OpenAICredentials.from_env()
+        self.client = OpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+        self.async_client = AsyncOpenAI(
+            api_key=self.credentials.openai_api_key.get_secret_value(),
+            organization=self.credentials.openai_organization_id,
+        )
+
+    def process(self, input_data: OpenAIEmbeddingsInput) -> OpenAIEmbeddingsOutput:
+        """Generate embeddings for the input text(s)"""
+        try:
+            # Handle single text input
+            texts = (
+                [input_data.texts]
+                if isinstance(input_data.texts, str)
+                else input_data.texts
+            )
+
+            # Create embeddings
+            response = self.client.embeddings.create(
+                model=input_data.model,
+                input=texts,
+                encoding_format=input_data.encoding_format,
+                dimensions=input_data.dimensions,
+            )
+
+            # Extract embeddings
+            embeddings = [data.embedding for data in response.data]
+
+            return OpenAIEmbeddingsOutput(
+                embeddings=embeddings,
+                used_model=response.model,
+                tokens_used=response.usage.total_tokens,
+            )
+        except Exception as e:
+            raise ProcessingError(f"OpenAI embeddings generation failed: {str(e)}")
+
+    async def process_async(
+        self, input_data: OpenAIEmbeddingsInput
+    ) -> OpenAIEmbeddingsOutput:
+        """Async version of the embeddings generation"""
+        try:
+            # Handle single text input
+            texts = (
+                [input_data.texts]
+                if isinstance(input_data.texts, str)
+                else input_data.texts
+            )
+
+            # Create embeddings
+            response = await self.async_client.embeddings.create(
+                model=input_data.model,
+                input=texts,
+                encoding_format=input_data.encoding_format,
+                dimensions=input_data.dimensions,
+            )
+
+            # Extract embeddings
+            embeddings = [data.embedding for data in response.data]
+
+            return OpenAIEmbeddingsOutput(
+                embeddings=embeddings,
+                used_model=response.model,
+                tokens_used=response.usage.total_tokens,
+            )
+        except Exception as e:
+            raise ProcessingError(
+                f"OpenAI async embeddings generation failed: {str(e)}"
+            )

airtrain/integrations/together/__init__.py

@@ -2,5 +2,18 @@
 
 from .credentials import TogetherAICredentials
 from .skills import TogetherAIChatSkill
+from .list_models import (
+    TogetherListModelsSkill,
+    TogetherListModelsInput,
+    TogetherListModelsOutput,
+)
+from .models import TogetherModel
 
-__all__ = ["TogetherAICredentials", "TogetherAIChatSkill"]
+__all__ = [
+    "TogetherAICredentials", 
+    "TogetherAIChatSkill",
+    "TogetherListModelsSkill",
+    "TogetherListModelsInput",
+    "TogetherListModelsOutput",
+    "TogetherModel",
+]

airtrain/integrations/together/list_models.py

@@ -0,0 +1,77 @@
+from typing import Optional
+import requests
+from pydantic import Field
+
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import TogetherAICredentials
+from .models import TogetherModel
+
+
+class TogetherListModelsInput(InputSchema):
+    """Schema for Together AI list models input"""
+    pass
+
+
+class TogetherListModelsOutput(OutputSchema):
+    """Schema for Together AI list models output"""
+    
+    data: list[TogetherModel] = Field(
+        default_factory=list,
+        description="List of Together AI models"
+    )
+    object: Optional[str] = Field(
+        default=None, 
+        description="Object type"
+    )
+
+
+class TogetherListModelsSkill(Skill[TogetherListModelsInput, TogetherListModelsOutput]):
+    """Skill for listing Together AI models"""
+
+    input_schema = TogetherListModelsInput
+    output_schema = TogetherListModelsOutput
+    
+    def __init__(self, credentials: Optional[TogetherAICredentials] = None):
+        """Initialize the skill with optional credentials"""
+        super().__init__()
+        self.credentials = credentials or TogetherAICredentials.from_env()
+        self.base_url = "https://api.together.xyz/v1"
+    
+    def process(
+        self, input_data: TogetherListModelsInput
+    ) -> TogetherListModelsOutput:
+        """Process the input and return a list of models."""
+        try:
+            # Build the URL
+            url = f"{self.base_url}/models"
+            
+            # Make the request
+            headers = {
+                "Authorization": (
+                    f"Bearer {self.credentials.together_api_key.get_secret_value()}"
+                ),
+                "accept": "application/json"
+            }
+            
+            response = requests.get(url, headers=headers)
+            response.raise_for_status()
+            
+            # Parse the response
+            result = response.json()
+            
+            # Convert the models to TogetherModel objects
+            models = []
+            for model_data in result.get("data", []):
+                models.append(TogetherModel(**model_data))
+            
+            # Return the output
+            return TogetherListModelsOutput(
+                data=models,
+                object=result.get("object")
+            )
+            
+        except requests.RequestException as e:
+            raise ProcessingError(f"Failed to list Together AI models: {str(e)}")
+        except Exception as e:
+            raise ProcessingError(f"Error listing Together AI models: {str(e)}")