airtrain 0.1.62__py3-none-any.whl → 0.1.67__py3-none-any.whl

airtrain/__init__.py

@@ -5,7 +5,7 @@ This library provides a flexible framework for building AI agents
 that can complete complex tasks using AI models, skills, and tools.
 """
 
-__version__ = "0.1.62"
+__version__ = "0.1.67"
 
 import sys
 

airtrain/integrations/fireworks/structured_completion_skills.py

@@ -161,6 +161,9 @@ class FireworksStructuredCompletionSkill(
                 data = response.json()
 
                 response_text = data["choices"][0]["text"]
+
+                response_text = response_text.split("</think>")[-1]
+
                 parsed_response = input_data.response_model.model_validate_json(
                     response_text
                 )

airtrain/integrations/openai/skills.py

@@ -1,12 +1,26 @@
-from typing import AsyncGenerator, List, Optional, Dict, TypeVar, Type, Generator, Union
+from typing import (
+    AsyncGenerator,
+    List,
+    Optional,
+    Dict,
+    TypeVar,
+    Type,
+    Generator,
+    Union,
+    Any,
+)
 from pydantic import Field, BaseModel
 from openai import OpenAI, AsyncOpenAI
 from openai.types.chat import ChatCompletionChunk
 import numpy as np
+import json
+from loguru import logger
 
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
 from .credentials import OpenAICredentials
+from ...core.credentials import get_credentials
+from ...core.errors import ValidationError
 
 
 class OpenAIInput(InputSchema):
@@ -35,6 +49,7 @@ class OpenAIInput(InputSchema):
         default=False,
         description="Whether to stream the response token by token",
     )
+    tools: Optional[List[Dict[str, Any]]] = None
 
 
 class OpenAIOutput(OutputSchema):
@@ -43,6 +58,8 @@ class OpenAIOutput(OutputSchema):
     response: str
     used_model: str
     usage: Dict[str, int]
+    raw_response: Optional[Any] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
 
 
 class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
@@ -64,6 +81,15 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
             organization=self.credentials.openai_organization_id,
         )
 
+    def validate_input(self, input_data: OpenAIInput) -> None:
+        """Validate the input before processing."""
+        if not input_data.model:
+            raise ValidationError("Model name is required")
+        if not input_data.user_input and not input_data.conversation_history:
+            raise ValidationError(
+                "Either user input or conversation history is required"
+            )
+
     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}]
@@ -77,6 +103,7 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
     def process_stream(self, input_data: OpenAIInput) -> Generator[str, None, None]:
         """Process the input and stream the response token by token."""
         try:
+            self.validate_input(input_data)
             messages = self._build_messages(input_data)
 
             stream = self.client.chat.completions.create(
@@ -92,45 +119,91 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
                     yield chunk.choices[0].delta.content
 
         except Exception as e:
+            logger.error(f"OpenAI streaming failed: {str(e)}")
             raise ProcessingError(f"OpenAI streaming failed: {str(e)}")
 
     def process(self, input_data: OpenAIInput) -> OpenAIOutput:
         """Process the input and return the complete response."""
         try:
+            self.validate_input(input_data)
+
             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)
+                return OpenAIOutput(
+                    response=response,
+                    used_model=input_data.model,
+                    usage={
+                        "total_tokens": 0,
+                        "prompt_tokens": 0,
+                        "completion_tokens": 0,
+                    },
+                )
             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,
-                used_model=input_data.model,
-                usage={
-                    "total_tokens": completion.usage.total_tokens,
-                    "prompt_tokens": completion.usage.prompt_tokens,
-                    "completion_tokens": completion.usage.completion_tokens,
-                },
-            )
+                # Create the completion parameters
+                completion_params = {
+                    "model": input_data.model,
+                    "messages": messages,
+                    "temperature": input_data.temperature,
+                    "max_tokens": input_data.max_tokens,
+                    "stream": False,
+                }
+
+                # Add tools if provided
+                if input_data.tools and len(input_data.tools) > 0:
+                    logger.info(f"Using {len(input_data.tools)} tools with OpenAI")
+                    completion_params["tools"] = input_data.tools
+
+                # Make the API call
+                completion = self.client.chat.completions.create(**completion_params)
+
+                # Extract response content
+                response = completion.choices[0].message.content or ""
+
+                # Extract tool calls if present
+                tool_calls = None
+                if (
+                    hasattr(completion.choices[0].message, "tool_calls")
+                    and completion.choices[0].message.tool_calls
+                ):
+                    tool_calls = []
+                    for tool_call in completion.choices[0].message.tool_calls:
+                        formatted_tool_call = {
+                            "id": tool_call.id,
+                            "name": tool_call.function.name,
+                            "arguments": tool_call.function.arguments,
+                        }
+                        tool_calls.append(formatted_tool_call)
+                    logger.info(
+                        f"Extracted {len(tool_calls)} tool calls from OpenAI response"
+                    )
+
+                return OpenAIOutput(
+                    response=response,
+                    used_model=input_data.model,
+                    usage={
+                        "total_tokens": completion.usage.total_tokens,
+                        "prompt_tokens": completion.usage.prompt_tokens,
+                        "completion_tokens": completion.usage.completion_tokens,
+                    },
+                    raw_response=completion,
+                    tool_calls=tool_calls,
+                )
 
         except Exception as e:
+            logger.error(f"OpenAI chat failed: {str(e)}")
             raise ProcessingError(f"OpenAI chat failed: {str(e)}")
 
     async def process_async(self, input_data: OpenAIInput) -> OpenAIOutput:
         """Async version of process method"""
         try:
+            self.validate_input(input_data)
             messages = self._build_messages(input_data)
             completion = await self.async_client.chat.completions.create(
                 model=input_data.model,
@@ -155,6 +228,7 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
     ) -> AsyncGenerator[str, None]:
         """Async version of stream processor"""
         try:
+            self.validate_input(input_data)
             messages = self._build_messages(input_data)
             stream = await self.async_client.chat.completions.create(
                 model=input_data.model,

airtrain/tools/__init__.py

@@ -21,6 +21,7 @@ from .network import ApiCallTool
 from .command import ExecuteCommandTool, FindFilesTool, TerminalNavigationTool
 from .search import SearchTermTool, WebSearchTool
 from .testing import RunPytestTool
+from .weather import WeatherTool
 
 __all__ = [
     # Base classes
@@ -42,4 +43,5 @@ __all__ = [
     "SearchTermTool",
     "WebSearchTool",
     "RunPytestTool",
+    "WeatherTool",
 ]

airtrain/tools/weather.py

@@ -0,0 +1,88 @@
+"""
+Weather tool for AirTrain.
+
+This module provides a tool for fetching weather data.
+For demonstration purposes, this returns mock data regardless of the location.
+"""
+
+import random
+from typing import Dict, Any
+from datetime import datetime
+from loguru import logger
+
+from .registry import StatelessTool, register_tool
+
+
+@register_tool("weather")
+class WeatherTool(StatelessTool):
+    """Tool for retrieving weather information for a location."""
+
+    def __init__(self):
+        self.name = "weather"
+        self.description = "Get current weather information for a location"
+        self.parameters = {
+            "type": "object",
+            "properties": {
+                "location": {
+                    "type": "string",
+                    "description": "City and country (e.g., 'New York, USA', 'Paris, France')",
+                }
+            },
+            "required": ["location"],
+            "additionalProperties": False,
+        }
+
+    def __call__(self, location: str) -> Dict[str, Any]:
+        """
+        Execute the weather tool to get mock weather data.
+
+        Args:
+            location: The city and country to get weather for
+
+        Returns:
+            A dictionary containing mock weather data
+        """
+        logger.info(f"Weather tool called with location: {location}")
+
+        # For demonstration purposes, return mock data regardless of input
+        temp_celsius = random.uniform(15.0, 30.0)
+        temp_fahrenheit = (temp_celsius * 9 / 5) + 32
+
+        conditions = random.choice(
+            ["Sunny", "Partly Cloudy", "Cloudy", "Light Rain", "Thunderstorms", "Clear"]
+        )
+
+        humidity = random.randint(30, 90)
+        wind_speed = random.uniform(0, 20)
+
+        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+        weather_data = {
+            "status": "success",
+            "location": location,
+            "timestamp": current_time,
+            "weather": {
+                "temperature": {
+                    "celsius": round(temp_celsius, 1),
+                    "fahrenheit": round(temp_fahrenheit, 1),
+                },
+                "conditions": conditions,
+                "humidity": humidity,
+                "wind_speed": round(wind_speed, 1),
+                "forecast": "This is mock data for demonstration purposes",
+            },
+        }
+
+        logger.info(f"Returning mock weather data for {location}")
+        return weather_data
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert tool to dictionary format for LLM function calling."""
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters,
+            },
+        }

{airtrain-0.1.62.dist-info → airtrain-0.1.67.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airtrain
-Version: 0.1.62
+Version: 0.1.67
 Summary: A platform for building and deploying AI agents with structured skills
 Home-page: https://github.com/rosaboyle/airtrain.dev
 Author: Dheeraj Pai

{airtrain-0.1.62.dist-info → airtrain-0.1.67.dist-info}/RECORD

@@ -1,6 +1,6 @@
-airtrain/__init__.py,sha256=gTO8GEuN1gIKjo5HlBOURfVnRtEuv61zF-HMxTXQYQk,3357
+airtrain/__init__.py,sha256=Wt0Bf2-pkLkX1tN2gepDmegHlmF7Xq80x3imMBmhEv4,3357
 airtrain/__main__.py,sha256=EU8ffFmCdC1G-UcHHt0Oo3lB1PGqfC6kwzH39CnYSwU,72
-airtrain/__pycache__/__init__.cpython-313.pyc,sha256=gNzuKeI25qKWqhRq6MRzf8Otl1qJEekuX_QU0HAx7bE,2759
+airtrain/__pycache__/__init__.cpython-313.pyc,sha256=SV1Rq3CsgY9QTiEgul7-YxCkdIEMymkD8KE9TVk5yxs,2759
 airtrain/agents/__init__.py,sha256=r6v5_bblxamRgiaCT8CVhyzaDdWohGM7sSjLgIUpA5s,795
 airtrain/agents/example_agent.py,sha256=0dCS8QXIvUYYkxwyOEMLMdlQ4KgMAerQ56r7NcYGqTw,11681
 airtrain/agents/groq_agent.py,sha256=s-f-cGSrgsR4Jlrv6xzeg2Z0LAEHByhFMYVbTvu-Bkg,10558
@@ -19,9 +19,9 @@ airtrain/core/__init__.py,sha256=9h7iKwTzZocCPc9bU6j8bA02BokteWIOcO1uaqGMcrk,254
 airtrain/core/credentials.py,sha256=J1jd8vLrOfet0GhLI1J44d35o7skjriBsMgpODmXwfo,5592
 airtrain/core/schemas.py,sha256=MMXrDviC4gRea_QaPpbjgO--B_UKxnD7YrxqZOLJZZU,7003
 airtrain/core/skills.py,sha256=kIkI1MwIzuAwyoxdnZ5MGOb70BOB-njbVb_csJEdVvc,9244
-airtrain/core/__pycache__/__init__.cpython-313.pyc,sha256=aHwrmX5Hf3s9pSa8V0PH3ytqldRjB_X9AWwA6IDpzyk,534
-airtrain/core/__pycache__/schemas.cpython-313.pyc,sha256=WpKRuARzTVGt_NDMApl8GNIQgF_jZcHV34z-KYs3g5g,8327
-airtrain/core/__pycache__/skills.cpython-313.pyc,sha256=4EppdsB6rISIJNF8CbyCPHXlMCX1JKETeKxOS8hENLw,11839
+airtrain/core/__pycache__/__init__.cpython-313.pyc,sha256=TQKjYutr-_x05pI-yyJShDLaNnCQMH-ZsGwSWBH_wi4,534
+airtrain/core/__pycache__/schemas.cpython-313.pyc,sha256=9OaxWCuPLgfeedvDnkKoEwNuLrxyRENjAOqnGAXpTK4,8327
+airtrain/core/__pycache__/skills.cpython-313.pyc,sha256=dUPAFMzeilfh7RYyQbZPaFnI8Tg0GTfI7CD4vGXOBtQ,11839
 airtrain/integrations/__init__.py,sha256=nGG3gZJ_sEaMzi_utsQCzExgnpHUjua8L9TyzofLQnw,2842
 airtrain/integrations/anthropic/__init__.py,sha256=K741w3v7fWsCknTo38ARqDL0D3HPlwDIvDuuBao9Tto,800
 airtrain/integrations/anthropic/credentials.py,sha256=hlTSw9HX66kYNaeQUtn0JjdZQBMNkzzFOJOoLOOzvcY,1246
@@ -45,7 +45,7 @@ airtrain/integrations/fireworks/list_models.py,sha256=o4fP0K3qstBopO7va2LysLp4_K
 airtrain/integrations/fireworks/models.py,sha256=yo4xtweSi4qQftg04r4naRddx3KjU9Jluzqf5C7V9f4,4626
 airtrain/integrations/fireworks/requests_skills.py,sha256=h6HRV5dGvV7t3zyjD-awW47RyeDbu8onNevhcgSSy94,8235
 airtrain/integrations/fireworks/skills.py,sha256=Ns1tXXTVtTeeVYadzm4dnmmOboo430WTMu2o56oWTDc,7156
-airtrain/integrations/fireworks/structured_completion_skills.py,sha256=airYakYWXzYRS9nfNfrH90N3eeN8YW7GaY3ygLSiBO8,6622
+airtrain/integrations/fireworks/structured_completion_skills.py,sha256=1cOt6aGYJGrO0Hy74jbLMzHGlgsUfynCSFLBf_6_Q9U,6692
 airtrain/integrations/fireworks/structured_requests_skills.py,sha256=uQR-nygtWmdGTwvU-aUdMNOMit_PiBVPYRa80ZloHLs,11852
 airtrain/integrations/fireworks/structured_skills.py,sha256=1wZ_7QDUhKWCSv_1lSEF6VnAqEeEA3jWHq7n0fWicgw,3897
 airtrain/integrations/google/__init__.py,sha256=ElwgcXfbg_gGMm6zbkMXCQPFKZUb-yTJk986o19A7Cs,214
@@ -63,7 +63,7 @@ airtrain/integrations/openai/chinese_assistant.py,sha256=F8bMeUUDly7BYG6wO648cAE
 airtrain/integrations/openai/credentials.py,sha256=NfRyp1QgEtgm8cxt2-BOLq-6d0X-Pcm80NnfHM8p0FY,1470
 airtrain/integrations/openai/list_models.py,sha256=vg8pZwLZ3F2Fx42X18WykpJOzZD9JG-2KJi49XWgSKo,4121
 airtrain/integrations/openai/models_config.py,sha256=W9mu_z9tCC4ZUKHSJ6Hk4X09TRZLqEhT7TtRY5JEk5g,8007
-airtrain/integrations/openai/skills.py,sha256=kEDe5q0Zlv_X-JGOYtb552ktb3aQQYVUYczVwMH0jxA,12823
+airtrain/integrations/openai/skills.py,sha256=7Rn5byvKVhk-OQI87Z-iCEP9ujrOyqcMM5K1gv8I1CM,15604
 airtrain/integrations/perplexity/__init__.py,sha256=asVQs-oVXVhnLmAOZ6l9R8KoigeHmw8D5OONb_aGgnY,1189
 airtrain/integrations/perplexity/credentials.py,sha256=5acl2DcF0dk7DQJWRwmFkBXPQ1zfKN4ARW31dP-4rpQ,1524
 airtrain/integrations/perplexity/list_models.py,sha256=iOOxQ50BfLc226ps0gSdFl-enxa-SmStiuU2vFTUhJ0,4298
@@ -94,15 +94,16 @@ airtrain/integrations/together/vision_models_config.py,sha256=m28HwYDk2Kup_J-a1F
 airtrain/telemetry/__init__.py,sha256=_xDHzSmQvRCqihT0QTmF7bS9yKEl2LaZ3Zq05hXaI3k,1108
 airtrain/telemetry/service.py,sha256=6IB2-FL3fGsYYgMIVQ9MNCp9UlSE5cVhLlB3VBYkAnY,5592
 airtrain/telemetry/views.py,sha256=qob1swyLNEk_UIpDi5VTwMDsEsGZhJheFQrGbP8T5hw,8115
-airtrain/tools/__init__.py,sha256=AauO_EEBcK09mkyCuvvVK_Oz0L5DOrnuySViWXCOt6Y,1021
+airtrain/tools/__init__.py,sha256=dL_CsjCD1uut0T-n6-LVXaH9k3Le2kvTIUDoQzjd1Hw,1073
 airtrain/tools/command.py,sha256=dxvs6RzppjWmkUe1oMtxOc7w2mFOGFFZ9Gylwnm37Sw,13355
 airtrain/tools/filesystem.py,sha256=-YYdHj_KeSWPYXeRhWhIX9s_KujVA1R5tF3r93zRVTU,6324
 airtrain/tools/network.py,sha256=YR0AtMXDXkhCsXcx7_t2d12ItnKY8XXTmyP1kdj2M4c,3883
 airtrain/tools/registry.py,sha256=K-1H5EipYcDNDx2jdpsEY9gjfV4aNCGI1pY2UsgSpC0,10246
 airtrain/tools/search.py,sha256=MJNi17g6aBPSqbF0ChV8ZgMlzz_PoKSPAIpe_dazdt8,15081
 airtrain/tools/testing.py,sha256=q4ALEPRzukiadY6wFSPY7vA-T1o3XInLhXt18dsf6yY,4397
-airtrain-0.1.62.dist-info/METADATA,sha256=fAOQQht25N594JsKGN6sFuiqjVMPjKMtly_RP6JW9co,6503
-airtrain-0.1.62.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-airtrain-0.1.62.dist-info/entry_points.txt,sha256=rrJ36IUsyq6n1dSfTWXqVAgpQLPRWDfCqwd6_3B-G0U,52
-airtrain-0.1.62.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
-airtrain-0.1.62.dist-info/RECORD,,
+airtrain/tools/weather.py,sha256=cOP79XF2GOHD_TKnwW7OA5DzykixugB06CzCQLIyONQ,2787
+airtrain-0.1.67.dist-info/METADATA,sha256=9xpWaDyUUyso_W_Z49Ff4hSXqKwM3zStRt61EvSMGng,6503
+airtrain-0.1.67.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+airtrain-0.1.67.dist-info/entry_points.txt,sha256=rrJ36IUsyq6n1dSfTWXqVAgpQLPRWDfCqwd6_3B-G0U,52
+airtrain-0.1.67.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.67.dist-info/RECORD,,