airtrain 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

airtrain/integrations/together/skills.py

@@ -0,0 +1,305 @@
+from typing import Optional, Dict, Any, List, Generator, Union
+from pydantic import Field, validator
+from airtrain.core.skills import Skill, ProcessingError
+from airtrain.core.schemas import InputSchema, OutputSchema
+from .credentials import TogetherAICredentials
+from .models import TogetherAIImageInput, TogetherAIImageOutput, GeneratedImage
+from .models_config import get_max_completion_tokens
+from pathlib import Path
+import base64
+import time
+from together import Together
+
+
+class TogetherAIInput(InputSchema):
+    """Schema for Together AI input"""
+
+    user_input: str = Field(..., description="User's input text")
+    system_prompt: str = Field(
+        default="You are a helpful assistant.",
+        description="System prompt to guide the model's behavior",
+    )
+    conversation_history: List[Dict[str, str]] = Field(
+        default_factory=list,
+        description=(
+            "List of previous conversation messages in "
+            "[{'role': 'user|assistant', 'content': 'message'}] format"
+        ),
+    )
+    model: str = Field(
+        default="meta-llama/Llama-3.1-8B-Instruct", 
+        description="Together AI model to use"
+    )
+    max_tokens: int = Field(
+        default=4096, 
+        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"
+    )
+    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."
+        ),
+    )
+    
+    @validator('max_tokens')
+    def validate_max_tokens(cls, v, values):
+        """Validate that max_tokens doesn't exceed the model's limit."""
+        if 'model' in values:
+            model_id = values['model']
+            max_limit = get_max_completion_tokens(model_id)
+            if v > max_limit:
+                return max_limit
+        return v
+
+
+class TogetherAIOutput(OutputSchema):
+    """Schema for Together AI output"""
+
+    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")
+    tool_calls: Optional[List[Dict[str, Any]]] = Field(
+        default=None, description="Tool calls generated by the model"
+    )
+
+
+class TogetherAIChatSkill(Skill[TogetherAIInput, TogetherAIOutput]):
+    """Skill for Together AI chat"""
+
+    input_schema = TogetherAIInput
+    output_schema = TogetherAIOutput
+
+    def __init__(self, credentials: Optional[TogetherAICredentials] = None):
+        super().__init__()
+        self.credentials = credentials or TogetherAICredentials.from_env()
+        self.client = Together(
+            api_key=self.credentials.together_api_key.get_secret_value()
+        )
+
+    def _build_messages(self, input_data: TogetherAIInput) -> List[Dict[str, str]]:
+        """
+        Build messages list from input data including conversation history.
+
+        Args:
+            input_data: The input data containing system prompt, conversation history, and user input
+
+        Returns:
+            List[Dict[str, str]]: List of messages in the format required by Together AI
+        """
+        messages = [{"role": "system", "content": input_data.system_prompt}]
+
+        # Add conversation history if present
+        if input_data.conversation_history:
+            messages.extend(input_data.conversation_history)
+
+        # Add current user input
+        messages.append({"role": "user", "content": input_data.user_input})
+
+        return messages
+
+    def process_stream(self, input_data: TogetherAIInput) -> Generator[str, None, None]:
+        """Process the input and stream the response token by token."""
+        try:
+            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,
+            )
+
+            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"Together AI streaming failed: {str(e)}")
+
+    def process(self, input_data: TogetherAIInput) -> TogetherAIOutput:
+        """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
+                tool_calls = None  # Tool calls not available in streaming
+            else:
+                messages = self._build_messages(input_data)
+                
+                # Prepare API call parameters
+                api_params = {
+                    "model": input_data.model,
+                    "messages": messages,
+                    "temperature": input_data.temperature,
+                    "max_tokens": input_data.max_tokens,
+                    "stream": False,
+                }
+                
+                # Add tools and tool_choice if provided
+                if input_data.tools:
+                    api_params["tools"] = input_data.tools
+                
+                if input_data.tool_choice:
+                    api_params["tool_choice"] = input_data.tool_choice
+                
+                try:
+                    completion = self.client.chat.completions.create(**api_params)
+                    response = completion.choices[0].message.content or ""
+                    
+                    # Extract usage information
+                    usage = (
+                        completion.usage.model_dump()
+                        if hasattr(completion, "usage")
+                        else {}
+                    )
+                    
+                    # Check for tool calls in the response
+                    tool_calls = None
+                    if (
+                        hasattr(completion.choices[0].message, "tool_calls") 
+                        and completion.choices[0].message.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 completion.choices[0].message.tool_calls
+                        ]
+                except Exception as api_error:
+                    # Provide more specific error messages for common tool call issues
+                    error_message = str(api_error)
+                    if "tool_call_id" in error_message:
+                        raise ProcessingError(
+                            "Tool call error: Missing tool_call_id in conversation history. "
+                            "When adding tool responses to conversation history, "
+                            "make sure each message with role='tool' includes a 'tool_call_id' field."
+                        )
+                    elif "messages" in error_message and "tool" in error_message:
+                        raise ProcessingError(
+                            "Tool call error: Invalid message format in conversation history. "
+                            f"Original error: {error_message}"
+                        )
+                    else:
+                        # Re-raise with original error message
+                        raise ProcessingError(f"Together AI API error: {error_message}")
+
+            return TogetherAIOutput(
+                response=response,
+                used_model=input_data.model,
+                usage=usage,
+                tool_calls=tool_calls
+            )
+
+        except ProcessingError:
+            # Re-raise ProcessingError without modification
+            raise
+        except Exception as e:
+            raise ProcessingError(f"Together AI processing failed: {str(e)}")
+
+
+class TogetherAIImageSkill(Skill[TogetherAIImageInput, TogetherAIImageOutput]):
+    """Skill for Together AI image generation"""
+
+    input_schema = TogetherAIImageInput
+    output_schema = TogetherAIImageOutput
+
+    def __init__(self, credentials: Optional[TogetherAICredentials] = None):
+        super().__init__()
+        self.credentials = credentials or TogetherAICredentials.from_env()
+        self.client = Together(
+            api_key=self.credentials.together_api_key.get_secret_value()
+        )
+
+    def process(self, input_data: TogetherAIImageInput) -> TogetherAIImageOutput:
+        try:
+            start_time = time.time()
+
+            # Generate images
+            response = self.client.images.generate(
+                prompt=input_data.prompt,
+                model=input_data.model,
+                steps=input_data.steps,
+                n=input_data.n,
+                size=input_data.size,
+                negative_prompt=input_data.negative_prompt,
+                seed=input_data.seed,
+            )
+
+            # Calculate total time
+            total_time = time.time() - start_time
+
+            # Convert response to our output format
+            generated_images = [
+                GeneratedImage(
+                    b64_json=img.b64_json,
+                    seed=getattr(img, "seed", None),
+                    finish_reason=getattr(img, "finish_reason", None),
+                )
+                for img in response.data
+            ]
+
+            return TogetherAIImageOutput(
+                images=generated_images,
+                model=input_data.model,
+                prompt=input_data.prompt,
+                total_time=total_time,
+                usage=getattr(response, "usage", {}),
+            )
+
+        except Exception as e:
+            raise ProcessingError(f"Together AI image generation failed: {str(e)}")
+
+    def save_images(
+        self, output: TogetherAIImageOutput, output_dir: Path
+    ) -> List[Path]:
+        """
+        Save generated images to disk
+
+        Args:
+            output (TogetherAIImageOutput): Generation output containing images
+            output_dir (Path): Directory to save images
+
+        Returns:
+            List[Path]: List of paths to saved images
+        """
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        saved_paths = []
+        for i, img in enumerate(output.images):
+            output_path = output_dir / f"image_{i}.png"
+            image_data = base64.b64decode(img.b64_json)
+
+            with open(output_path, "wb") as f:
+                f.write(image_data)
+
+            saved_paths.append(output_path)
+
+        return saved_paths

airtrain/integrations/together/vision_models_config.py

@@ -0,0 +1,49 @@
+from typing import Dict, NamedTuple
+
+
+class VisionModelConfig(NamedTuple):
+    organization: str
+    display_name: str
+    context_length: int
+
+
+TOGETHER_VISION_MODELS: Dict[str, VisionModelConfig] = {
+    "meta-llama/Llama-Vision-Free": VisionModelConfig(
+        organization="Meta",
+        display_name="(Free) Llama 3.2 11B Vision Instruct Turbo",
+        context_length=131072,
+    ),
+    "meta-llama/Llama-3.2-11B-Vision-Instruct-Turbo": VisionModelConfig(
+        organization="Meta",
+        display_name="Llama 3.2 11B Vision Instruct Turbo",
+        context_length=131072,
+    ),
+    "meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo": VisionModelConfig(
+        organization="Meta",
+        display_name="Llama 3.2 90B Vision Instruct Turbo",
+        context_length=131072,
+    ),
+}
+
+
+def get_vision_model_config(model_id: str) -> VisionModelConfig:
+    """Get vision model configuration by model ID"""
+    if model_id not in TOGETHER_VISION_MODELS:
+        raise ValueError(f"Model {model_id} not found in Together AI vision models")
+    return TOGETHER_VISION_MODELS[model_id]
+
+
+def list_vision_models_by_organization(
+    organization: str,
+) -> Dict[str, VisionModelConfig]:
+    """Get all vision models for a specific organization"""
+    return {
+        model_id: config
+        for model_id, config in TOGETHER_VISION_MODELS.items()
+        if config.organization.lower() == organization.lower()
+    }
+
+
+def get_default_vision_model() -> str:
+    """Get the default vision model ID"""
+    return "meta-llama/Llama-Vision-Free"

airtrain/telemetry/__init__.py

@@ -0,0 +1,38 @@
+"""
+Airtrain Telemetry
+
+This package provides telemetry functionality for Airtrain usage.
+Telemetry is enabled by default to help improve the library and can be disabled by 
+setting AIRTRAIN_TELEMETRY_ENABLED=false in your environment variables or .env file.
+"""
+
+from airtrain.telemetry.service import ProductTelemetry
+from airtrain.telemetry.views import (
+    AgentRunTelemetryEvent,
+    AgentStepTelemetryEvent,
+    AgentEndTelemetryEvent,
+    ModelInvocationTelemetryEvent,
+    ErrorTelemetryEvent,
+    UserFeedbackTelemetryEvent,
+    SkillInitTelemetryEvent,
+    SkillProcessTelemetryEvent,
+    PackageInstallTelemetryEvent,
+    PackageImportTelemetryEvent,
+)
+
+__all__ = [
+    "ProductTelemetry",
+    "AgentRunTelemetryEvent",
+    "AgentStepTelemetryEvent",
+    "AgentEndTelemetryEvent",
+    "ModelInvocationTelemetryEvent",
+    "ErrorTelemetryEvent",
+    "UserFeedbackTelemetryEvent",
+    "SkillInitTelemetryEvent",
+    "SkillProcessTelemetryEvent",
+    "PackageInstallTelemetryEvent",
+    "PackageImportTelemetryEvent",
+]
+
+# Create a singleton instance for easy import
+telemetry = ProductTelemetry() 

airtrain/telemetry/service.py

@@ -0,0 +1,167 @@
+import logging
+import os
+import platform
+import sys
+import uuid
+from pathlib import Path
+
+from dotenv import load_dotenv
+from posthog import Posthog
+
+from airtrain.telemetry.views import BaseTelemetryEvent
+
+load_dotenv()
+
+logger = logging.getLogger(__name__)
+
+# Enhanced event settings to collect more data
+POSTHOG_EVENT_SETTINGS = {
+    'process_person_profile': True,
+    'enable_sent_at': True,  # Add timing information
+    'capture_performance': True,  # Collect performance data
+    'capture_pageview': True,  # More detailed usage tracking
+}
+
+
+def singleton(cls):
+    """Singleton decorator for classes."""
+    instances = {}
+
+    def get_instance(*args, **kwargs):
+        if cls not in instances:
+            instances[cls] = cls(*args, **kwargs)
+        return instances[cls]
+
+    return get_instance
+
+
+@singleton
+class ProductTelemetry:
+    """
+    Service for capturing telemetry data from Airtrain usage.
+
+    Telemetry is enabled by default but can be disabled by setting 
+    AIRTRAIN_TELEMETRY_ENABLED=false in your environment.
+    """
+
+    USER_ID_PATH = str(
+        Path.home() / '.cache' / 'airtrain' / 'telemetry_user_id'
+    )
+    # API key for PostHog
+    PROJECT_API_KEY = 'phc_1pLNkG3QStYEXIz0CAPQaOGpcmxpE3CJXhE1HANWgIz'
+    HOST = 'https://us.i.posthog.com'
+    UNKNOWN_USER_ID = 'UNKNOWN'
+
+    _curr_user_id = None
+
+    def __init__(self) -> None:
+        telemetry_disabled = os.getenv('AIRTRAIN_TELEMETRY_ENABLED', 'true').lower() == 'false'
+        self.debug_logging = os.getenv('AIRTRAIN_LOGGING_LEVEL', 'info').lower() == 'debug'
+        
+        # System information to include with telemetry
+        self.system_info = {
+            'os': platform.system(),
+            'os_version': platform.version(),
+            'python_version': f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
+            'platform': platform.platform(),
+            'machine': platform.machine(),
+            'hostname': platform.node(),
+            'username': os.getlogin() if hasattr(os, 'getlogin') else 'unknown'
+        }
+        isBeta = True  # TODO: remove this once out of beta
+        if telemetry_disabled and not isBeta:
+            self._posthog_client = None
+        else:
+            if not isBeta:
+                logging.info(
+                    'Telemetry enabled. To disable, set '
+                    'AIRTRAIN_TELEMETRY_ENABLED=false in your environment.'
+                )
+            if isBeta:
+                logging.info(
+                    'You are currently in beta. Telemetry is enabled by default.'
+                )
+            self._posthog_client = Posthog(
+                project_api_key=self.PROJECT_API_KEY,
+                host=self.HOST,
+                disable_geoip=False  # Collect geographical data
+            )
+
+            # Set debug mode if enabled
+            if self.debug_logging:
+                self._posthog_client.debug = True
+
+            # Identify user more specifically
+            self._posthog_client.identify(
+                self.user_id,
+                {
+                    **self.system_info,
+                    'first_seen': True
+                }
+            )
+
+            # Silence posthog's logging only if debug is off
+            if not self.debug_logging:
+                posthog_logger = logging.getLogger('posthog')
+                posthog_logger.disabled = True
+
+        if self._posthog_client is None:
+            logger.debug('Telemetry disabled')
+
+    def capture(self, event: BaseTelemetryEvent) -> None:
+        """Capture a telemetry event and send it to PostHog if telemetry is enabled."""
+        if self._posthog_client is None:
+            return
+
+        # Add system information to all events
+        enhanced_properties = {
+            **event.properties,
+            **POSTHOG_EVENT_SETTINGS,
+            **self.system_info
+        }
+
+        if self.debug_logging:
+            logger.debug(f'Telemetry event: {event.name} {enhanced_properties}')
+        self._direct_capture(event, enhanced_properties)
+
+    def _direct_capture(self, event: BaseTelemetryEvent, enhanced_properties: dict) -> None:
+        """
+        Send the event to PostHog. Should not be thread blocking because posthog handles it.
+        """
+        if self._posthog_client is None:
+            return
+
+        try:
+            self._posthog_client.capture(
+                self.user_id,
+                event.name,
+                enhanced_properties
+            )
+        except Exception as e:
+            logger.error(f'Failed to send telemetry event {event.name}: {e}')
+
+    @property
+    def user_id(self) -> str:
+        """
+        Get the user ID for telemetry.
+        Creates a new one if it doesn't exist.
+        """
+        if self._curr_user_id:
+            return self._curr_user_id
+
+        # File access may fail due to permissions or other reasons.
+        # We don't want to crash so we catch all exceptions.
+        try:
+            if not os.path.exists(self.USER_ID_PATH):
+                os.makedirs(os.path.dirname(self.USER_ID_PATH), exist_ok=True)
+                with open(self.USER_ID_PATH, 'w') as f:
+                    # Use a more identifiable ID prefix
+                    new_user_id = f"airtrain-user-{uuid.uuid4()}"
+                    f.write(new_user_id)
+                self._curr_user_id = new_user_id
+            else:
+                with open(self.USER_ID_PATH, 'r') as f:
+                    self._curr_user_id = f.read()
+        except Exception:
+            self._curr_user_id = self.UNKNOWN_USER_ID
+        return self._curr_user_id