airtrain 0.1.66__tar.gz → 0.1.68__tar.gz

{airtrain-0.1.66 → airtrain-0.1.68}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airtrain
-Version: 0.1.66
+Version: 0.1.68
 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.66 → airtrain-0.1.68}/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.66"
+__version__ = "0.1.68"
 
 import sys
 

{airtrain-0.1.66 → airtrain-0.1.68}/airtrain/integrations/openai/skills.py

@@ -16,11 +16,13 @@ import numpy as np
 import json
 from loguru import logger
 
-from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.skills import (
+    Skill,
+    ProcessingError,
+    InputValidationError as ValidationError,
+)
 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):
@@ -106,13 +108,23 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
             self.validate_input(input_data)
             messages = self._build_messages(input_data)
 
-            stream = self.client.chat.completions.create(
-                model=input_data.model,
-                messages=messages,
-                temperature=input_data.temperature,
-                max_tokens=input_data.max_tokens,
-                stream=True,
-            )
+            # Create the completion parameters for streaming
+            completion_params = {
+                "model": input_data.model,
+                "messages": messages,
+                "temperature": input_data.temperature,
+                "max_tokens": input_data.max_tokens,
+                "stream": True,
+            }
+
+            # 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 streaming"
+                )
+                completion_params["tools"] = input_data.tools
+
+            stream = self.client.chat.completions.create(**completion_params)
 
             for chunk in stream:
                 if chunk.choices[0].delta.content is not None:
@@ -205,22 +217,56 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
         try:
             self.validate_input(input_data)
             messages = self._build_messages(input_data)
+
+            # 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 async")
+                completion_params["tools"] = input_data.tools
+
             completion = await self.async_client.chat.completions.create(
-                model=input_data.model,
-                messages=messages,
-                temperature=input_data.temperature,
-                max_tokens=input_data.max_tokens,
+                **completion_params
             )
+
+            # 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 async response"
+                )
+
             return OpenAIOutput(
-                response=completion.choices[0].message.content,
+                response=completion.choices[0].message.content or "",
                 used_model=completion.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 async chat failed: {str(e)}")
             raise ProcessingError(f"OpenAI async chat failed: {str(e)}")
 
     async def process_stream_async(
@@ -230,17 +276,32 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
         try:
             self.validate_input(input_data)
             messages = self._build_messages(input_data)
+
+            # Create the completion parameters for streaming
+            completion_params = {
+                "model": input_data.model,
+                "messages": messages,
+                "temperature": input_data.temperature,
+                "max_tokens": input_data.max_tokens,
+                "stream": True,
+            }
+
+            # 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 async streaming"
+                )
+                completion_params["tools"] = input_data.tools
+
             stream = await self.async_client.chat.completions.create(
-                model=input_data.model,
-                messages=messages,
-                temperature=input_data.temperature,
-                max_tokens=input_data.max_tokens,
-                stream=True,
+                **completion_params
             )
+
             async for chunk in stream:
                 if chunk.choices[0].delta.content is not None:
                     yield chunk.choices[0].delta.content
         except Exception as e:
+            logger.error(f"OpenAI async streaming failed: {str(e)}")
             raise ProcessingError(f"OpenAI async streaming failed: {str(e)}")
 
 
@@ -306,6 +367,7 @@ class OpenAIParserSkill(Skill[OpenAIParserInput, OpenAIParserOutput]):
             )
 
         except Exception as e:
+            logger.error(f"OpenAI parsing failed: {str(e)}")
             raise ProcessingError(f"OpenAI parsing failed: {str(e)}")
 
 
@@ -380,6 +442,7 @@ class OpenAIEmbeddingsSkill(Skill[OpenAIEmbeddingsInput, OpenAIEmbeddingsOutput]
                 tokens_used=response.usage.total_tokens,
             )
         except Exception as e:
+            logger.error(f"OpenAI embeddings generation failed: {str(e)}")
             raise ProcessingError(f"OpenAI embeddings generation failed: {str(e)}")
 
     async def process_async(
@@ -411,6 +474,7 @@ class OpenAIEmbeddingsSkill(Skill[OpenAIEmbeddingsInput, OpenAIEmbeddingsOutput]
                 tokens_used=response.usage.total_tokens,
             )
         except Exception as e:
+            logger.error(f"OpenAI async embeddings generation failed: {str(e)}")
             raise ProcessingError(
                 f"OpenAI async embeddings generation failed: {str(e)}"
             )

{airtrain-0.1.66 → airtrain-0.1.68}/airtrain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airtrain
-Version: 0.1.66
+Version: 0.1.68
 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.66 → airtrain-0.1.68}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "airtrain"
-version = "0.1.66"
+version = "0.1.68"
 description = "A platform for building and deploying AI agents with structured skills"
 readme = "README.md"
 requires-python = ">=3.8"