voxagent 0.1.0__py3-none-any.whl

voxagent/providers/augment.py

@@ -0,0 +1,293 @@
+"""Augment (Auggie) CLI provider.
+
+This provider wraps the Auggie CLI directly using subprocess.
+It requires:
+1. The auggie CLI to be installed: brew install augment-cli
+2. Authentication via: auggie login OR a vault service providing tokens
+
+Models available:
+- haiku4.5: Fast and efficient responses
+- sonnet4.5: Great for everyday tasks
+- sonnet4: Legacy model
+- opus4.5: Best for complex tasks (Claude Opus 4.5)
+- gpt5: OpenAI GPT-5 legacy
+- gpt5.1: Strong reasoning and planning
+
+Vault Integration:
+When a vault_service is provided (any object implementing VaultProtocol with a
+get(key) method), the provider checks for an 'augment_access_token'. If found,
+it sets the AUGMENT_SESSION_AUTH environment variable when invoking the CLI.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+from voxagent.providers.base import (
+    AbortSignal,
+    ErrorChunk,
+    MessageEndChunk,
+    StreamChunk,
+    TextDeltaChunk,
+)
+from voxagent.providers.cli_base import CLINotFoundError, CLIProvider
+from voxagent.types import Message
+
+if TYPE_CHECKING:
+    pass  # No external type imports needed
+
+
+@runtime_checkable
+class VaultProtocol(Protocol):
+    """Protocol for vault services that can provide credentials.
+
+    This allows voxagent to work with any vault implementation without
+    depending on voxdomus directly.
+    """
+
+    def get(self, key: str) -> str:
+        """Get a credential value by key.
+
+        Args:
+            key: The credential key to retrieve.
+
+        Returns:
+            The credential value.
+
+        Raises:
+            KeyError: If the credential doesn't exist.
+        """
+        ...
+
+    def exists(self, key: str) -> bool:
+        """Check if a credential exists.
+
+        Args:
+            key: The credential key to check.
+
+        Returns:
+            True if the credential exists, False otherwise.
+        """
+        ...
+
+logger = logging.getLogger(__name__)
+
+# Vault key for the augment access token
+VAULT_KEY_ACCESS_TOKEN = "augment_access_token"
+
+
+class AugmentProvider(CLIProvider):
+    """Provider for Augment using the auggie CLI directly.
+
+    This provider spawns the auggie CLI as a subprocess to avoid
+    issues with the auggie-sdk's async task management.
+
+    Optionally integrates with voxDomus vault for token storage.
+    When a vault_service is provided and contains an access token,
+    the provider sets AUGMENT_SESSION_AUTH for the CLI subprocess.
+    """
+
+    CLI_NAME = "auggie"
+    ENV_KEY = "AUGMENT_API_TOKEN"
+
+    SUPPORTED_MODELS = [
+        "haiku4.5",
+        "sonnet4.5",
+        "sonnet4",
+        "opus4.5",
+        "gpt5",
+        "gpt5.1",
+    ]
+
+    def __init__(
+        self,
+        model: str = "sonnet4.5",
+        api_key: str | None = None,
+        base_url: str | None = None,
+        vault_service: VaultProtocol | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize Augment provider.
+
+        Args:
+            model: Model name (haiku4.5, sonnet4.5, opus4.5, etc.).
+            api_key: Optional API token (usually from auggie login).
+            base_url: Optional API URL override.
+            vault_service: Optional vault service (any object with a get(key) method).
+            **kwargs: Additional arguments.
+        """
+        super().__init__(model=model, api_key=api_key, base_url=base_url, **kwargs)
+        self._vault_service: VaultProtocol | None = vault_service
+        self._vault_token: str | None = None
+        self._vault_checked = False
+
+    @property
+    def name(self) -> str:
+        """Get the provider name."""
+        return "augment"
+
+    @property
+    def models(self) -> list[str]:
+        """Get supported models."""
+        return self.SUPPORTED_MODELS
+
+    @property
+    def supports_tools(self) -> bool:
+        """Auggie supports tools but we don't expose them."""
+        return False
+
+    @property
+    def context_limit(self) -> int:
+        """Approximate context limit."""
+        return 200000
+
+    def _build_cli_args(
+        self,
+        prompt: str,
+        system: str | None = None,
+    ) -> list[str]:
+        """Build auggie CLI arguments.
+
+        Uses --print mode for non-interactive one-shot execution.
+        """
+        args = ["--print", "--quiet", "--instruction", prompt]
+
+        if self._model:
+            args.extend(["--model", self._model])
+
+        return args
+
+    def _parse_output(self, stdout: str, stderr: str) -> str:
+        """Parse auggie CLI output."""
+        # auggie chat outputs the response directly
+        return stdout.strip()
+
+    def _get_vault_token(self) -> str | None:
+        """Get access token from vault if available.
+
+        Returns:
+            Access token string if found in vault, None otherwise.
+
+        Note:
+            Results are cached after first lookup.
+        """
+        if self._vault_checked:
+            return self._vault_token
+
+        self._vault_checked = True
+
+        if self._vault_service is None:
+            return None
+
+        try:
+            if self._vault_service.exists(VAULT_KEY_ACCESS_TOKEN):
+                self._vault_token = self._vault_service.get(VAULT_KEY_ACCESS_TOKEN)
+                logger.debug("Retrieved Augment token from vault")
+        except Exception as e:
+            logger.debug("Could not retrieve Augment token from vault: %s", e)
+
+        return self._vault_token
+
+    async def _run_cli(
+        self,
+        prompt: str,
+        system: str | None = None,
+    ) -> str:
+        """Run CLI command and return output.
+
+        Overrides base class to inject AUGMENT_SESSION_AUTH env var
+        if a vault token is available.
+
+        Args:
+            prompt: User prompt.
+            system: Optional system prompt.
+
+        Returns:
+            Parsed response text.
+
+        Raises:
+            Exception: If CLI execution fails.
+        """
+        cli_path = self._get_cli_path()
+        args = [cli_path] + self._build_cli_args(prompt, system)
+
+        logger.debug("Running CLI: %s", " ".join(args))
+
+        # Build environment with vault token if available
+        env = os.environ.copy()
+        vault_token = self._get_vault_token()
+        if vault_token:
+            env["AUGMENT_SESSION_AUTH"] = vault_token
+            logger.debug("Using vault token for AUGMENT_SESSION_AUTH")
+
+        proc = await asyncio.create_subprocess_exec(
+            *args,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=env,
+        )
+
+        stdout_bytes, stderr_bytes = await proc.communicate()
+        stdout = stdout_bytes.decode("utf-8", errors="replace")
+        stderr = stderr_bytes.decode("utf-8", errors="replace")
+
+        if proc.returncode != 0:
+            logger.warning("CLI exited with code %d: %s", proc.returncode, stderr)
+
+        return self._parse_output(stdout, stderr)
+
+    async def stream(
+        self,
+        messages: list[Message],
+        system: str | None = None,
+        tools: list[Any] | None = None,
+        abort_signal: AbortSignal | None = None,
+    ) -> AsyncIterator[StreamChunk]:
+        """Stream a response from Auggie CLI.
+
+        Note: The auggie CLI uses its own MCP tool configuration.
+        Tools passed from voxDomus are not used - auggie will use
+        tools from its own ~/.augment/settings.json config.
+        """
+        if tools:
+            logger.debug(
+                "Auggie CLI uses its own MCP tools - ignoring %d passed tools",
+                len(tools),
+            )
+
+        try:
+            prompt = self._messages_to_prompt(messages)
+            response = await self._run_cli(prompt, system)
+            if response:
+                yield TextDeltaChunk(delta=response)
+        except CLINotFoundError as e:
+            yield ErrorChunk(error=str(e))
+        except Exception as e:
+            yield ErrorChunk(error=f"Auggie CLI error: {e}")
+
+        yield MessageEndChunk()
+
+    async def complete(
+        self,
+        messages: list[Message],
+        system: str | None = None,
+        tools: list[Any] | None = None,
+    ) -> Message:
+        """Get a complete response from Auggie CLI."""
+        text_parts: list[str] = []
+
+        async for chunk in self.stream(messages, system, tools):
+            if isinstance(chunk, TextDeltaChunk):
+                text_parts.append(chunk.delta)
+            elif isinstance(chunk, ErrorChunk):
+                raise Exception(chunk.error)
+
+        return Message(role="assistant", content="".join(text_parts))
+
+
+__all__ = ["AugmentProvider"]
+

voxagent/providers/auth.py

@@ -0,0 +1,116 @@
+"""Authentication profile management with failover support.
+
+This module provides API credential profile management with cooldown
+and failure tracking for multi-provider LLM failover.
+"""
+
+from datetime import datetime, timedelta
+from enum import Enum
+
+from pydantic import BaseModel, ConfigDict
+
+
+class FailoverReason(str, Enum):
+    """Reasons for failover to a different API profile."""
+
+    AUTH_ERROR = "auth_error"
+    RATE_LIMIT = "rate_limit"
+    TIMEOUT = "timeout"
+    CONTEXT_OVERFLOW = "context_overflow"
+    MODEL_ERROR = "model_error"
+
+
+class FailoverError(Exception):
+    """Error that triggers failover to another profile."""
+
+    def __init__(self, reason: FailoverReason, message: str = "") -> None:
+        self.reason = reason
+        super().__init__(message or reason.value)
+
+
+class AuthProfile(BaseModel):
+    """API credential profile with cooldown and failure tracking."""
+
+    model_config = ConfigDict(validate_assignment=True)
+
+    id: str
+    api_key: str
+    provider: str
+    cooldown_until: datetime | None = None
+    failure_count: int = 0
+    disabled: bool = False
+
+    def is_available(self) -> bool:
+        """Check if profile is available (not in cooldown, not disabled)."""
+        if self.disabled:
+            return False
+        if self.cooldown_until is not None:
+            if datetime.now() < self.cooldown_until:
+                return False
+        return True
+
+    def set_cooldown(self, duration_seconds: int) -> None:
+        """Set cooldown until now + duration."""
+        self.cooldown_until = datetime.now() + timedelta(seconds=duration_seconds)
+
+    def record_failure(self) -> None:
+        """Increment failure count."""
+        self.failure_count += 1
+
+    def reset_failures(self) -> None:
+        """Reset failure count on success."""
+        self.failure_count = 0
+
+    def disable(self) -> None:
+        """Disable this profile."""
+        self.disabled = True
+
+
+class AuthProfileManager:
+    """Manages multiple API profiles with failover support."""
+
+    DEFAULT_MAX_FAILURES = 3
+    DEFAULT_RATE_LIMIT_COOLDOWN = 60  # seconds
+
+    def __init__(
+        self,
+        profiles: list[AuthProfile],
+        max_failures: int = DEFAULT_MAX_FAILURES,
+        rate_limit_cooldown: int = DEFAULT_RATE_LIMIT_COOLDOWN,
+    ) -> None:
+        self._profiles = list(profiles)
+        self._max_failures = max_failures
+        self._rate_limit_cooldown = rate_limit_cooldown
+
+    @property
+    def profiles(self) -> list[AuthProfile]:
+        return self._profiles
+
+    @property
+    def max_failures(self) -> int:
+        return self._max_failures
+
+    @property
+    def rate_limit_cooldown(self) -> int:
+        return self._rate_limit_cooldown
+
+    def get_available_profiles(self, provider: str | None = None) -> list[AuthProfile]:
+        """Get profiles that are not in cooldown and not disabled."""
+        available = [p for p in self._profiles if p.is_available()]
+        if provider is not None:
+            available = [p for p in available if p.provider == provider]
+        return available
+
+    def handle_failover(self, profile: AuthProfile, error: FailoverError) -> None:
+        """Update profile state based on failover error."""
+        if error.reason == FailoverReason.RATE_LIMIT:
+            profile.set_cooldown(self._rate_limit_cooldown)
+        elif error.reason == FailoverReason.AUTH_ERROR:
+            profile.record_failure()
+            if profile.failure_count >= self._max_failures:
+                profile.disable()
+
+    def record_success(self, profile: AuthProfile) -> None:
+        """Record successful use of profile."""
+        profile.reset_failures()
+

voxagent/providers/base.py

@@ -0,0 +1,268 @@
+"""Base provider abstract class and stream chunk types.
+
+This module defines the abstract interface for LLM providers:
+- StreamChunk types for streaming responses
+- AbortSignal for cancellation
+- BaseProvider ABC that all providers must implement
+"""
+
+import os
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+from voxagent.types import Message, ToolCall
+
+
+# =============================================================================
+# Stream Chunk Types
+# =============================================================================
+
+
+class TextDeltaChunk(BaseModel):
+    """A text delta chunk from streaming response.
+
+    Attributes:
+        type: Discriminator field, always "text_delta".
+        delta: The text content delta.
+    """
+
+    type: Literal["text_delta"] = "text_delta"
+    delta: str
+
+
+class ToolUseChunk(BaseModel):
+    """A tool use chunk from streaming response.
+
+    Attributes:
+        type: Discriminator field, always "tool_use".
+        tool_call: The tool call request.
+    """
+
+    type: Literal["tool_use"] = "tool_use"
+    tool_call: ToolCall
+
+
+class MessageEndChunk(BaseModel):
+    """A message end chunk signaling stream completion.
+
+    Attributes:
+        type: Discriminator field, always "message_end".
+    """
+
+    type: Literal["message_end"] = "message_end"
+
+
+class ErrorChunk(BaseModel):
+    """An error chunk from streaming response.
+
+    Attributes:
+        type: Discriminator field, always "error".
+        error: The error message.
+    """
+
+    type: Literal["error"] = "error"
+    error: str
+
+
+# Union type for all stream chunks
+StreamChunk = TextDeltaChunk | ToolUseChunk | MessageEndChunk | ErrorChunk
+
+
+# =============================================================================
+# Abort Signal
+# =============================================================================
+
+
+class AbortSignal:
+    """A signal for aborting async operations.
+
+    Attributes:
+        _aborted: Internal flag indicating if abort has been requested.
+        _reason: The reason for the abort.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the abort signal."""
+        self._aborted = False
+        self._reason: str = ""
+
+    @property
+    def aborted(self) -> bool:
+        """Check if abort has been requested."""
+        return self._aborted
+
+    @property
+    def reason(self) -> str:
+        """Get the reason for the abort."""
+        return self._reason
+
+    def abort(self, reason: str = "Aborted") -> None:
+        """Request abortion of the operation.
+
+        Args:
+            reason: The reason for aborting.
+        """
+        self._aborted = True
+        self._reason = reason
+
+
+# =============================================================================
+# Base Provider Abstract Class
+# =============================================================================
+
+
+class BaseProvider(ABC):
+    """Abstract base class for LLM providers.
+
+    Subclasses must implement all abstract properties and methods.
+    The ENV_KEY class variable should be set to the environment variable
+    name for the API key.
+    """
+
+    ENV_KEY: str = ""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the provider.
+
+        Args:
+            api_key: API key for authentication. Falls back to ENV_KEY env var.
+            base_url: Optional base URL for API requests.
+            **kwargs: Additional provider-specific arguments.
+        """
+        self._api_key = api_key
+        self._base_url = base_url
+        self._kwargs = kwargs
+
+    @property
+    def api_key(self) -> str | None:
+        """Get API key from constructor or environment variable."""
+        if self._api_key is not None:
+            return self._api_key
+        return os.environ.get(self.ENV_KEY) if self.ENV_KEY else None
+
+    @property
+    def base_url(self) -> str | None:
+        """Get the base URL for API requests."""
+        return self._base_url
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Get the provider name."""
+        ...
+
+    @property
+    @abstractmethod
+    def models(self) -> list[str]:
+        """Get the list of supported model names."""
+        ...
+
+    @property
+    @abstractmethod
+    def supports_tools(self) -> bool:
+        """Check if the provider supports tool/function calling."""
+        ...
+
+    @property
+    @abstractmethod
+    def supports_streaming(self) -> bool:
+        """Check if the provider supports streaming responses."""
+        ...
+
+    @property
+    @abstractmethod
+    def context_limit(self) -> int:
+        """Get the maximum context length in tokens."""
+        ...
+
+    @abstractmethod
+    async def stream(
+        self,
+        messages: list[Message],
+        system: str | None = None,
+        tools: list[Any] | None = None,
+        abort_signal: AbortSignal | None = None,
+    ) -> AsyncIterator[StreamChunk]:
+        """Stream a response from the provider.
+
+        Args:
+            messages: The conversation messages.
+            system: Optional system prompt.
+            tools: Optional list of tool definitions.
+            abort_signal: Optional signal to abort the stream.
+
+        Yields:
+            StreamChunk objects containing response data.
+        """
+        ...
+        # This yield is needed to make this an async generator
+        yield  # type: ignore[misc]
+
+    @abstractmethod
+    async def complete(
+        self,
+        messages: list[Message],
+        system: str | None = None,
+        tools: list[Any] | None = None,
+    ) -> Message:
+        """Get a complete response from the provider.
+
+        Args:
+            messages: The conversation messages.
+            system: Optional system prompt.
+            tools: Optional list of tool definitions.
+
+        Returns:
+            The assistant's response message.
+        """
+        ...
+
+    @abstractmethod
+    def count_tokens(
+        self,
+        messages: list[Message],
+        system: str | None = None,
+    ) -> int:
+        """Count tokens in the messages.
+
+        Args:
+            messages: The conversation messages.
+            system: Optional system prompt.
+
+        Returns:
+            The token count.
+        """
+        ...
+
+    def get_api_key(self, env_var_name: str) -> str | None:
+        """Get API key from constructor or specified environment variable.
+
+        Args:
+            env_var_name: The environment variable name to check.
+
+        Returns:
+            The API key or None if not found.
+        """
+        if self._api_key is not None:
+            return self._api_key
+        return os.environ.get(env_var_name)
+
+
+__all__ = [
+    "AbortSignal",
+    "BaseProvider",
+    "ErrorChunk",
+    "MessageEndChunk",
+    "StreamChunk",
+    "TextDeltaChunk",
+    "ToolUseChunk",
+]
+