airtrain 0.1.67__py3-none-any.whl → 0.1.68__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.67"
+__version__ = "0.1.68"
 
 import sys
 

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.67.dist-info → airtrain-0.1.68.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airtrain
-Version: 0.1.67
+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.67.dist-info → airtrain-0.1.68.dist-info}/RECORD

@@ -1,6 +1,6 @@
-airtrain/__init__.py,sha256=Wt0Bf2-pkLkX1tN2gepDmegHlmF7Xq80x3imMBmhEv4,3357
+airtrain/__init__.py,sha256=4kP3qhcib7cRor8EFs8syt7KyP9FxFEvDHfnzIRpinw,3357
 airtrain/__main__.py,sha256=EU8ffFmCdC1G-UcHHt0Oo3lB1PGqfC6kwzH39CnYSwU,72
-airtrain/__pycache__/__init__.cpython-313.pyc,sha256=SV1Rq3CsgY9QTiEgul7-YxCkdIEMymkD8KE9TVk5yxs,2759
+airtrain/__pycache__/__init__.cpython-313.pyc,sha256=WFA-vgd8933joKp50Rgj4Yt04TM-gTvQS0hRR-VtEdw,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=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/core/__pycache__/__init__.cpython-313.pyc,sha256=-xKNz5rm5DOdR4c9MtqdEkD8AQ5csdTsSDAQlNOT-N8,534
+airtrain/core/__pycache__/schemas.cpython-313.pyc,sha256=R0VqNuR747I0OUwuw03U8MWpLMpC7N_4wI__v9vhUkM,8327
+airtrain/core/__pycache__/skills.cpython-313.pyc,sha256=7K1QYJmfqtyXuBqRHOR_zQOV41d2WBprp3UznC1Yu4A,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
@@ -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=7Rn5byvKVhk-OQI87Z-iCEP9ujrOyqcMM5K1gv8I1CM,15604
+airtrain/integrations/openai/skills.py,sha256=olLCkF29oKjWividPtS9JQ2L9S2nyrD6ziT3Ac1y0iE,18087
 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
@@ -102,8 +102,8 @@ airtrain/tools/registry.py,sha256=K-1H5EipYcDNDx2jdpsEY9gjfV4aNCGI1pY2UsgSpC0,10
 airtrain/tools/search.py,sha256=MJNi17g6aBPSqbF0ChV8ZgMlzz_PoKSPAIpe_dazdt8,15081
 airtrain/tools/testing.py,sha256=q4ALEPRzukiadY6wFSPY7vA-T1o3XInLhXt18dsf6yY,4397
 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,,
+airtrain-0.1.68.dist-info/METADATA,sha256=Ee7I0VzSRztZCPWcsh5nFcuAq9HSaGspIJBpVgsZzI4,6503
+airtrain-0.1.68.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+airtrain-0.1.68.dist-info/entry_points.txt,sha256=rrJ36IUsyq6n1dSfTWXqVAgpQLPRWDfCqwd6_3B-G0U,52
+airtrain-0.1.68.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.68.dist-info/RECORD,,