airtrain 0.1.28__tar.gz → 0.1.31__tar.gz

{airtrain-0.1.28/airtrain.egg-info → airtrain-0.1.31}/PKG-INFO

@@ -1,10 +1,12 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.28
+Version: 0.1.31
 Summary: A platform for building and deploying AI agents with structured skills
 Home-page: https://github.com/rosaboyle/airtrain.dev
 Author: Dheeraj Pai
-Author-email: helloworldcmu@gmail.com
+Author-email: Dheeraj Pai <helloworldcmu@gmail.com>
+Project-URL: Homepage, https://github.com/rosaboyle/airtrain.dev
+Project-URL: Documentation, https://docs.airtrain.dev/
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
@@ -26,15 +28,29 @@ Requires-Dist: boto3>=1.36.6
 Requires-Dist: together>=1.3.13
 Requires-Dist: anthropic>=0.45.0
 Requires-Dist: groq>=0.15.0
+Requires-Dist: cerebras-cloud-sdk>=1.19.0
+Requires-Dist: google-genai>=1.0.0
+Requires-Dist: fireworks-ai>=0.15.12
+Requires-Dist: google-generativeai>=0.8.4
+Requires-Dist: click>=8.0.0
+Requires-Dist: rich>=13.3.1
+Requires-Dist: prompt-toolkit>=3.0.36
+Requires-Dist: colorama>=0.4.6
+Requires-Dist: typer>=0.9.0
+Provides-Extra: dev
+Requires-Dist: black>=24.10.0; extra == "dev"
+Requires-Dist: flake8>=7.1.1; extra == "dev"
+Requires-Dist: isort>=5.13.0; extra == "dev"
+Requires-Dist: mypy>=1.9.0; extra == "dev"
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: twine>=4.0.0; extra == "dev"
+Requires-Dist: build>=0.10.0; extra == "dev"
+Requires-Dist: types-PyYAML>=6.0; extra == "dev"
+Requires-Dist: types-requests>=2.31.0; extra == "dev"
+Requires-Dist: types-Markdown>=3.5.0; extra == "dev"
 Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
 Dynamic: home-page
-Dynamic: requires-dist
 Dynamic: requires-python
-Dynamic: summary
 
 # Airtrain
 
@@ -167,31 +183,3 @@ Contributions are welcome! Please feel free to submit a Pull Request.
 ## License
 
 This project is licensed under the MIT License - see the LICENSE file for details. 
-
-## Changelog
-
-
-
-## 0.1.28
-
-- Bug fix: reasoning to Fireworks structured output.
-- Added reasoning to Fireworks structured output.
-
-## 0.1.27
-
-- Added structured completion skills for Fireworks AI
-- Added Completion skills for Fireworks AI.
-- Added Combination skill for Groq and Fireworks AI.
-- Add completion streaming.
-- Added strcutured output streaming for Fireworks AI.
-
-## 0.1.23
-
-- Added conversation support for Deepseek, Togehter AI, Fireworks AI, Gemini, Groq, Cerebras and Sambanova.
-- Added Change Log
-
-
-## Notes
-
-The changelog format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
-and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). 

{airtrain-0.1.28 → airtrain-0.1.31}/airtrain/__init__.py

@@ -1,6 +1,6 @@
 """Airtrain - A platform for building and deploying AI agents with structured skills"""
 
-__version__ = "0.1.28"
+__version__ = "0.1.33"
 
 # Core imports
 from .core.skills import Skill, ProcessingError

airtrain-0.1.31/airtrain/cli/main.py

@@ -0,0 +1,53 @@
+import click
+from airtrain.integrations.openai.skills import OpenAIChatSkill, OpenAIInput
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def initialize_chat():
+    return OpenAIChatSkill()
+
+
+@click.group()
+def cli():
+    """Airtrain CLI - Your AI Agent Building Assistant"""
+    pass
+
+
+@cli.command()
+def chat():
+    """Start an interactive chat session with Airtrain"""
+    skill = initialize_chat()
+    click.echo("Welcome to Airtrain! I'm here to help you build your AI Agent.")
+    click.echo("Type 'exit' to end the conversation.\n")
+
+    while True:
+        user_input = click.prompt("You", type=str)
+
+        if user_input.lower() == "exit":
+            click.echo("\nGoodbye! Have a great day!")
+            break
+
+        try:
+            input_data = OpenAIInput(
+                user_input=user_input,
+                system_prompt="You are an AI assistant that helps users build their own AI agents. Be helpful and provide clear explanations.",
+                model="gpt-4o",
+                temperature=0.7,
+            )
+
+            result = skill.process(input_data)
+            click.echo(f"\nAirtrain: {result.response}\n")
+
+        except Exception as e:
+            click.echo(f"\nError: {str(e)}\n")
+
+
+def main():
+    cli()
+
+
+if __name__ == "__main__":
+    main()

{airtrain-0.1.28 → airtrain-0.1.31}/airtrain/integrations/anthropic/skills.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Dict, Any
+from typing import List, Optional, Dict, Any, Generator
 from pydantic import Field
 from anthropic import Anthropic
 import base64
@@ -35,6 +35,9 @@ class AnthropicInput(InputSchema):
         default_factory=list,
         description="List of image paths to include in the message",
     )
+    stream: bool = Field(
+        default=False, description="Whether to stream the response progressively"
+    )
 
 
 class AnthropicOutput(OutputSchema):
@@ -102,24 +105,49 @@ class AnthropicChatSkill(Skill[AnthropicInput, AnthropicOutput]):
 
         return messages
 
-    def process(self, input_data: AnthropicInput) -> AnthropicOutput:
+    def process_stream(self, input_data: AnthropicInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
         try:
-            # Build messages using the helper method
             messages = self._build_messages(input_data)
 
-            # Create chat completion with system prompt as a separate parameter
-            response = self.client.messages.create(
+            with self.client.beta.messages.stream(
                 model=input_data.model,
-                system=input_data.system_prompt,  # System prompt passed directly
+                system=input_data.system_prompt,
                 messages=messages,
                 max_tokens=input_data.max_tokens,
                 temperature=input_data.temperature,
-            )
+            ) as stream:
+                for chunk in stream.text_stream:
+                    yield chunk
+
+        except Exception as e:
+            logger.exception(f"Anthropic streaming failed: {str(e)}")
+            raise ProcessingError(f"Anthropic streaming failed: {str(e)}")
+
+    def process(self, input_data: AnthropicInput) -> AnthropicOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+                usage = {}  # Usage stats not available in streaming
+            else:
+                messages = self._build_messages(input_data)
+                response = self.client.messages.create(
+                    model=input_data.model,
+                    system=input_data.system_prompt,
+                    messages=messages,
+                    max_tokens=input_data.max_tokens,
+                    temperature=input_data.temperature,
+                )
+                usage = response.usage.model_dump() if response.usage else {}
 
             return AnthropicOutput(
                 response=response.content[0].text,
                 used_model=input_data.model,
-                usage=response.usage.model_dump(),
+                usage=usage,
             )
 
         except Exception as e:

{airtrain-0.1.28 → airtrain-0.1.31}/airtrain/integrations/cerebras/skills.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Dict, Any
+from typing import List, Optional, Dict, Any, Generator
 from pydantic import Field
 from cerebras.cloud.sdk import Cerebras
 from loguru import logger
@@ -27,6 +27,9 @@ class CerebrasInput(InputSchema):
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )
+    stream: bool = Field(
+        default=False, description="Whether to stream the response progressively"
+    )
 
 
 class CerebrasOutput(OutputSchema):
@@ -71,23 +74,52 @@ class CerebrasChatSkill(Skill[CerebrasInput, CerebrasOutput]):
 
         return messages
 
-    def process(self, input_data: CerebrasInput) -> CerebrasOutput:
+    def process_stream(self, input_data: CerebrasInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
         try:
-            # Build messages using the helper method
             messages = self._build_messages(input_data)
 
-            # Create chat completion
-            response = self.client.chat.completions.create(
+            stream = self.client.chat.completions.create(
                 model=input_data.model,
                 messages=messages,
                 temperature=input_data.temperature,
                 max_tokens=input_data.max_tokens,
+                stream=True,
             )
 
+            for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
+
+        except Exception as e:
+            logger.exception(f"Cerebras streaming failed: {str(e)}")
+            raise ProcessingError(f"Cerebras streaming failed: {str(e)}")
+
+    def process(self, input_data: CerebrasInput) -> CerebrasOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+                usage = {}  # Usage stats not available in streaming
+            else:
+                messages = self._build_messages(input_data)
+                response = self.client.chat.completions.create(
+                    model=input_data.model,
+                    messages=messages,
+                    temperature=input_data.temperature,
+                    max_tokens=input_data.max_tokens,
+                )
+                usage = (
+                    response.usage.model_dump() if hasattr(response, "usage") else {}
+                )
+
             return CerebrasOutput(
                 response=response.choices[0].message.content,
                 used_model=input_data.model,
-                usage=response.usage.model_dump(),
+                usage=usage,
             )
 
         except Exception as e:

{airtrain-0.1.28 → airtrain-0.1.31}/airtrain/integrations/fireworks/requests_skills.py

@@ -1,8 +1,9 @@
-from typing import List, Optional, Dict, Any, Generator
+from typing import List, Optional, Dict, Any, Generator, AsyncGenerator
 from pydantic import Field
 import requests
 import json
 from loguru import logger
+import aiohttp
 
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
@@ -69,6 +70,11 @@ class FireworksRequestSkill(Skill[FireworksRequestInput, FireworksRequestOutput]
             "Content-Type": "application/json",
             "Authorization": f"Bearer {self.credentials.fireworks_api_key.get_secret_value()}",
         }
+        self.stream_headers = {
+            "Accept": "text/event-stream",
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.credentials.fireworks_api_key.get_secret_value()}",
+        }
 
     def _build_messages(
         self, input_data: FireworksRequestInput
@@ -150,3 +156,52 @@ class FireworksRequestSkill(Skill[FireworksRequestInput, FireworksRequestOutput]
 
         except Exception as e:
             raise ProcessingError(f"Fireworks request failed: {str(e)}")
+
+    async def process_async(
+        self, input_data: FireworksRequestInput
+    ) -> FireworksRequestOutput:
+        """Async version of process method using aiohttp"""
+        try:
+            async with aiohttp.ClientSession() as session:
+                payload = self._build_payload(input_data)
+                async with session.post(
+                    self.BASE_URL, headers=self.headers, json=payload
+                ) as response:
+                    response.raise_for_status()
+                    data = await response.json()
+
+                    return FireworksRequestOutput(
+                        response=data["choices"][0]["message"]["content"],
+                        used_model=input_data.model,
+                        usage=data.get("usage", {}),
+                    )
+
+        except Exception as e:
+            raise ProcessingError(f"Async Fireworks request failed: {str(e)}")
+
+    async def process_stream_async(
+        self, input_data: FireworksRequestInput
+    ) -> AsyncGenerator[str, None]:
+        """Async version of stream processor using aiohttp"""
+        try:
+            async with aiohttp.ClientSession() as session:
+                payload = self._build_payload(input_data)
+                async with session.post(
+                    self.BASE_URL, headers=self.stream_headers, json=payload
+                ) as response:
+                    response.raise_for_status()
+
+                    async for line in response.content:
+                        if line.startswith(b"data: "):
+                            chunk = json.loads(line[6:].strip())
+                            if "choices" in chunk:
+                                content = (
+                                    chunk["choices"][0]
+                                    .get("delta", {})
+                                    .get("content", "")
+                                )
+                                if content:
+                                    yield content
+
+        except Exception as e:
+            raise ProcessingError(f"Async Fireworks streaming failed: {str(e)}")

{airtrain-0.1.28 → airtrain-0.1.31}/airtrain/integrations/fireworks/structured_requests_skills.py

@@ -130,13 +130,15 @@ class FireworksStructuredRequestSkill(
                     except json.JSONDecodeError:
                         continue
 
-            # Once complete, parse the full JSON
-            complete_json = "".join(json_buffer)
+            # Once complete, parse the full response with think tags
+            complete_response = "".join(json_buffer)
+            reasoning, json_str = self._parse_response_content(complete_response)
+
             try:
                 parsed_response = input_data.response_model.model_validate_json(
-                    complete_json
+                    json_str
                 )
-                yield {"complete": parsed_response}
+                yield {"complete": parsed_response, "reasoning": reasoning}
             except Exception as e:
                 raise ProcessingError(f"Failed to parse JSON response: {str(e)}")
 
@@ -164,12 +166,14 @@ class FireworksStructuredRequestSkill(
                 # 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")
@@ -178,6 +182,7 @@ class FireworksStructuredRequestSkill(
                     parsed_response=parsed_response,
                     used_model=input_data.model,
                     usage={},  # Usage stats not available in streaming mode
+                    reasoning=reasoning,
                 )
             else:
                 # For non-streaming, use regular request

{airtrain-0.1.28 → airtrain-0.1.31}/airtrain/integrations/groq/skills.py

@@ -1,4 +1,4 @@
-from typing import Optional, Dict, Any, List
+from typing import Generator, Optional, Dict, Any, List
 from pydantic import Field
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
@@ -18,11 +18,16 @@ class GroqInput(InputSchema):
         default_factory=list,
         description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
     )
-    model: str = Field(default="mixtral-8x7b", description="Groq model to use")
+    model: str = Field(
+        default="deepseek-r1-distill-llama-70b-specdec", description="Groq model to use"
+    )
     max_tokens: int = Field(default=1024, description="Maximum tokens in response")
     temperature: float = Field(
         default=0.7, description="Temperature for response generation", ge=0, le=1
     )
+    stream: bool = Field(
+        default=False, description="Whether to stream the response progressively"
+    )
 
 
 class GroqOutput(OutputSchema):
@@ -30,7 +35,9 @@ class GroqOutput(OutputSchema):
 
     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")
+    usage: Dict[str, Any] = Field(
+        default_factory=dict, description="Usage statistics from the API"
+    )
 
 
 class GroqChatSkill(Skill[GroqInput, GroqOutput]):
@@ -65,23 +72,53 @@ class GroqChatSkill(Skill[GroqInput, GroqOutput]):
 
         return messages
 
-    def process(self, input_data: GroqInput) -> GroqOutput:
+    def process_stream(self, input_data: GroqInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
         try:
-            # Build messages using the helper method
             messages = self._build_messages(input_data)
 
-            # Create chat completion
-            response = self.client.chat.completions.create(
+            stream = self.client.chat.completions.create(
                 model=input_data.model,
                 messages=messages,
                 temperature=input_data.temperature,
                 max_tokens=input_data.max_tokens,
+                stream=True,
             )
 
+            for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
+
+        except Exception as e:
+            raise ProcessingError(f"Groq streaming failed: {str(e)}")
+
+    def process(self, input_data: GroqInput) -> GroqOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+                usage = {}  # Usage stats not available in streaming
+            else:
+                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
+                usage = {
+                    "total_tokens": completion.usage.total_tokens,
+                    "prompt_tokens": completion.usage.prompt_tokens,
+                    "completion_tokens": completion.usage.completion_tokens,
+                }
+
             return GroqOutput(
-                response=response.choices[0].message.content,
-                used_model=input_data.model,
-                usage=response.usage.model_dump(),
+                response=response, used_model=input_data.model, usage=usage
             )
 
         except Exception as e:

{airtrain-0.1.28 → airtrain-0.1.31}/airtrain/integrations/openai/skills.py

@@ -1,9 +1,6 @@
-from typing import List, Optional, Dict, Any, TypeVar, Type, Generator
+from typing import AsyncGenerator, List, Optional, Dict, TypeVar, Type, Generator
 from pydantic import Field, BaseModel
-from openai import OpenAI
-import base64
-from pathlib import Path
-from loguru import logger
+from openai import OpenAI, AsyncOpenAI
 from openai.types.chat import ChatCompletionChunk
 
 from airtrain.core.skills import Skill, ProcessingError
@@ -48,7 +45,7 @@ class OpenAIOutput(OutputSchema):
 
 
 class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
-    """Skill for interacting with OpenAI models"""
+    """Skill for interacting with OpenAI models with async support"""
 
     input_schema = OpenAIInput
     output_schema = OpenAIOutput
@@ -61,6 +58,10 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
             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 _build_messages(self, input_data: OpenAIInput) -> List[Dict[str, str]]:
         """Build messages list from input data including conversation history."""
@@ -126,6 +127,47 @@ class OpenAIChatSkill(Skill[OpenAIInput, OpenAIOutput]):
         except Exception as e:
             raise ProcessingError(f"OpenAI chat failed: {str(e)}")
 
+    async def process_async(self, input_data: OpenAIInput) -> OpenAIOutput:
+        """Async version of process method"""
+        try:
+            messages = self._build_messages(input_data)
+            completion = await self.async_client.chat.completions.create(
+                model=input_data.model,
+                messages=messages,
+                temperature=input_data.temperature,
+                max_tokens=input_data.max_tokens,
+            )
+            return OpenAIOutput(
+                response=completion.choices[0].message.content,
+                used_model=completion.model,
+                usage={
+                    "total_tokens": completion.usage.total_tokens,
+                    "prompt_tokens": completion.usage.prompt_tokens,
+                    "completion_tokens": completion.usage.completion_tokens,
+                },
+            )
+        except Exception as e:
+            raise ProcessingError(f"OpenAI async chat failed: {str(e)}")
+
+    async def process_stream_async(
+        self, input_data: OpenAIInput
+    ) -> AsyncGenerator[str, None]:
+        """Async version of stream processor"""
+        try:
+            messages = self._build_messages(input_data)
+            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,
+            )
+            async for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
+        except Exception as e:
+            raise ProcessingError(f"OpenAI async streaming failed: {str(e)}")
+
 
 ResponseT = TypeVar("ResponseT", bound=BaseModel)
 

{airtrain-0.1.28 → airtrain-0.1.31}/airtrain/integrations/sambanova/skills.py

@@ -1,4 +1,4 @@
-from typing import Optional, Dict, Any, List
+from typing import Optional, Dict, Any, List, Generator
 from pydantic import Field
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
@@ -28,6 +28,9 @@ class SambanovaInput(InputSchema):
     top_p: float = Field(
         default=0.1, description="Top p sampling parameter", ge=0, le=1
     )
+    stream: bool = Field(
+        default=False, description="Whether to stream the response progressively"
+    )
 
 
 class SambanovaOutput(OutputSchema):
@@ -73,24 +76,53 @@ class SambanovaChatSkill(Skill[SambanovaInput, SambanovaOutput]):
 
         return messages
 
-    def process(self, input_data: SambanovaInput) -> SambanovaOutput:
+    def process_stream(self, input_data: SambanovaInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
         try:
-            # Build messages using the helper method
             messages = self._build_messages(input_data)
 
-            # Create chat completion
-            response = self.client.chat.completions.create(
+            stream = self.client.chat.completions.create(
                 model=input_data.model,
                 messages=messages,
                 temperature=input_data.temperature,
                 max_tokens=input_data.max_tokens,
                 top_p=input_data.top_p,
+                stream=True,
             )
 
+            for chunk in stream:
+                if chunk.choices[0].delta.content is not None:
+                    yield chunk.choices[0].delta.content
+
+        except Exception as e:
+            raise ProcessingError(f"Sambanova streaming failed: {str(e)}")
+
+    def process(self, input_data: SambanovaInput) -> SambanovaOutput:
+        """Process the input and return the complete response."""
+        try:
+            if input_data.stream:
+                response_chunks = []
+                for chunk in self.process_stream(input_data):
+                    response_chunks.append(chunk)
+                response = "".join(response_chunks)
+                usage = {}  # Usage stats not available in streaming
+            else:
+                messages = self._build_messages(input_data)
+                response = self.client.chat.completions.create(
+                    model=input_data.model,
+                    messages=messages,
+                    temperature=input_data.temperature,
+                    max_tokens=input_data.max_tokens,
+                    top_p=input_data.top_p,
+                )
+                usage = (
+                    response.usage.model_dump() if hasattr(response, "usage") else {}
+                )
+
             return SambanovaOutput(
                 response=response.choices[0].message.content,
                 used_model=input_data.model,
-                usage=response.usage.model_dump(),
+                usage=usage,
             )
 
         except Exception as e: