augment-sdk 0.1.1__py3-none-any.whl

augment/exceptions.py

@@ -0,0 +1,92 @@
+"""Custom exceptions for the Augment SDK"""
+
+
+class AugmentError(Exception):
+    """Base exception for all Augment SDK errors"""
+
+    def __init__(self, message: str, status_code: int = None):
+        self.message = message
+        self.status_code = status_code
+        super().__init__(self.message)
+
+
+class AuthenticationError(AugmentError):
+    """Raised when authentication fails"""
+
+    def __init__(self, message: str = "Authentication failed"):
+        super().__init__(message, status_code=401)
+
+
+class RateLimitError(AugmentError):
+    """Raised when rate limit is exceeded"""
+
+    def __init__(self, message: str = "Rate limit exceeded", retry_after: int = None):
+        super().__init__(message, status_code=429)
+        self.retry_after = retry_after
+
+
+class NotFoundError(AugmentError):
+    """Raised when a resource is not found"""
+
+    def __init__(self, message: str = "Resource not found"):
+        super().__init__(message, status_code=404)
+
+
+class ValidationError(AugmentError):
+    """Raised when request validation fails"""
+
+    def __init__(self, message: str = "Validation error"):
+        super().__init__(message, status_code=400)
+
+
+class AugmentCLIError(AugmentError):
+    """Raised when the auggie CLI command fails"""
+
+    def __init__(self, message: str, return_code: int = None, stderr: str = None):
+        super().__init__(message)
+        self.return_code = return_code
+        self.stderr = stderr
+
+
+class AugmentNotFoundError(AugmentError):
+    """Raised when auggie CLI is not found or accessible"""
+
+    def __init__(self, message: str = "auggie CLI not found"):
+        super().__init__(message)
+
+
+class AugmentWorkspaceError(AugmentError):
+    """Raised when workspace path is invalid or inaccessible"""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
+class AugmentParseError(AugmentError):
+    """Raised when parsing agent response fails"""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
+class AugmentJSONError(AugmentError):
+    """Raised when JSON parsing fails"""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
+class AugmentVerificationError(AugmentError):
+    """Raised when success criteria verification fails after max rounds"""
+
+    def __init__(
+        self,
+        message: str,
+        unmet_criteria: list = None,
+        issues: list = None,
+        rounds_attempted: int = None,
+    ):
+        super().__init__(message)
+        self.unmet_criteria = unmet_criteria or []
+        self.issues = issues or []
+        self.rounds_attempted = rounds_attempted

augment/function_tools.py

@@ -0,0 +1,265 @@
+"""
+Utilities for converting Python functions to tool schemas for agent function calling.
+"""
+
+import inspect
+import json
+from typing import Any, Callable, Dict, List, Optional, Union, get_args, get_origin
+
+
+def python_type_to_json_schema_type(py_type: Any) -> Dict[str, Any]:
+    """
+    Convert a Python type annotation to a JSON schema type.
+
+    Args:
+        py_type: Python type annotation
+
+    Returns:
+        JSON schema type definition
+    """
+    # Handle None/NoneType
+    if py_type is None or py_type is type(None):
+        return {"type": "null"}
+
+    # Get origin for generic types (List, Dict, Optional, etc.)
+    origin = get_origin(py_type)
+    args = get_args(py_type)
+
+    # Handle Optional[T] (Union[T, None])
+    if origin is Union:
+        # Filter out None types
+        non_none_args = [arg for arg in args if arg is not type(None)]
+        if len(non_none_args) == 1:
+            # This is Optional[T]
+            return python_type_to_json_schema_type(non_none_args[0])
+        else:
+            # Multiple union types - use anyOf
+            return {"anyOf": [python_type_to_json_schema_type(arg) for arg in args]}
+
+    # Handle List[T]
+    if origin is list or py_type is list:
+        if args:
+            return {"type": "array", "items": python_type_to_json_schema_type(args[0])}
+        return {"type": "array"}
+
+    # Handle Dict[K, V]
+    if origin is dict or py_type is dict:
+        schema = {"type": "object"}
+        if args and len(args) == 2:
+            # Dict[str, T] - we can specify value type
+            schema["additionalProperties"] = python_type_to_json_schema_type(args[1])
+        return schema
+
+    # Handle basic types
+    type_mapping = {
+        str: {"type": "string"},
+        int: {"type": "integer"},
+        float: {"type": "number"},
+        bool: {"type": "boolean"},
+        Any: {},  # No type constraint
+    }
+
+    if py_type in type_mapping:
+        return type_mapping[py_type]
+
+    # Default to object for unknown types
+    return {"type": "object"}
+
+
+def extract_param_description(
+    docstring: Optional[str], param_name: str
+) -> Optional[str]:
+    """
+    Extract parameter description from docstring.
+
+    Supports Google-style and NumPy-style docstrings.
+
+    Args:
+        docstring: Function docstring
+        param_name: Parameter name to find
+
+    Returns:
+        Parameter description if found, None otherwise
+    """
+    if not docstring:
+        return None
+
+    lines = docstring.split("\n")
+    in_args_section = False
+
+    for i, line in enumerate(lines):
+        stripped = line.strip()
+
+        # Check if we're entering Args/Parameters section
+        if stripped.lower() in ["args:", "arguments:", "parameters:", "params:"]:
+            in_args_section = True
+            continue
+
+        # Check if we're leaving Args section (new section starts)
+        if (
+            in_args_section
+            and stripped.endswith(":")
+            and not stripped.startswith(param_name)
+        ):
+            in_args_section = False
+            continue
+
+        # Look for parameter in Args section
+        if in_args_section:
+            # Google style: "param_name: description" or "param_name (type): description"
+            if stripped.startswith(f"{param_name}:") or stripped.startswith(
+                f"{param_name} ("
+            ):
+                # Extract description after colon
+                colon_idx = stripped.find(":", stripped.find(param_name))
+                if colon_idx != -1:
+                    desc = stripped[colon_idx + 1 :].strip()
+                    # Check if description continues on next lines
+                    j = i + 1
+                    while j < len(lines):
+                        next_line = lines[j].strip()
+                        # Stop if we hit another parameter or empty line
+                        if not next_line or ":" in next_line:
+                            break
+                        desc += " " + next_line
+                        j += 1
+                    return desc
+
+    return None
+
+
+def function_to_schema(func: Callable) -> Dict[str, Any]:
+    """
+    Convert a Python function to a JSON schema for agent function calling.
+
+    The function should have:
+    - Type hints for parameters
+    - A docstring with description and parameter descriptions
+
+    Args:
+        func: Python function to convert
+
+    Returns:
+        JSON schema dictionary with name, description, and parameters
+
+    Example:
+        >>> def get_weather(location: str, unit: str = "celsius") -> dict:
+        ...     '''Get weather for a location.
+        ...
+        ...     Args:
+        ...         location: City name or coordinates
+        ...         unit: Temperature unit (celsius or fahrenheit)
+        ...     '''
+        ...     pass
+        >>> schema = function_to_schema(get_weather)
+        >>> schema['name']
+        'get_weather'
+    """
+    sig = inspect.signature(func)
+    doc = inspect.getdoc(func)
+
+    # Extract function description (first line/paragraph of docstring)
+    description = ""
+    if doc:
+        # Get everything before Args/Parameters section
+        lines = doc.split("\n")
+        desc_lines = []
+        for line in lines:
+            stripped = line.strip()
+            if stripped.lower() in [
+                "args:",
+                "arguments:",
+                "parameters:",
+                "params:",
+                "returns:",
+                "return:",
+            ]:
+                break
+            desc_lines.append(line)
+        description = "\n".join(desc_lines).strip()
+
+    if not description:
+        description = f"Call the {func.__name__} function"
+
+    # Build parameters schema
+    properties = {}
+    required = []
+
+    for param_name, param in sig.parameters.items():
+        # Skip self, cls, *args, **kwargs
+        if param_name in ["self", "cls"] or param.kind in [
+            inspect.Parameter.VAR_POSITIONAL,
+            inspect.Parameter.VAR_KEYWORD,
+        ]:
+            continue
+
+        # Get type annotation
+        param_type = (
+            param.annotation if param.annotation != inspect.Parameter.empty else Any
+        )
+        param_schema = python_type_to_json_schema_type(param_type)
+
+        # Get parameter description from docstring
+        param_desc = extract_param_description(doc, param_name)
+        if param_desc:
+            param_schema["description"] = param_desc
+
+        properties[param_name] = param_schema
+
+        # Mark as required if no default value
+        if param.default == inspect.Parameter.empty:
+            required.append(param_name)
+
+    schema = {
+        "name": func.__name__,
+        "description": description,
+        "parameters": {
+            "type": "object",
+            "properties": properties,
+        },
+    }
+
+    if required:
+        schema["parameters"]["required"] = required
+
+    return schema
+
+
+def format_function_schemas_for_prompt(functions: List[Callable]) -> str:
+    """
+    Format function schemas as a string for inclusion in agent prompt.
+
+    Args:
+        functions: List of Python functions
+
+    Returns:
+        Formatted string describing available functions
+    """
+    if not functions:
+        return ""
+
+    schemas = [function_to_schema(func) for func in functions]
+
+    prompt = "You have access to the following functions that you can call:\n\n"
+
+    for schema in schemas:
+        prompt += f"Function: {schema['name']}\n"
+        prompt += f"Description: {schema['description']}\n"
+        prompt += f"Parameters: {json.dumps(schema['parameters'], indent=2)}\n\n"
+
+    prompt += """To call a function, include in your response:
+<function-call>
+{
+  "name": "function_name",
+  "arguments": {
+    "param1": "value1",
+    "param2": "value2"
+  }
+}
+</function-call>
+
+You can call multiple functions by including multiple <function-call> blocks.
+After calling functions, you will receive the results and can continue your response.
+"""
+
+    return prompt

augment/listener.py

@@ -0,0 +1,186 @@
+"""
+Agent listener interface for receiving notifications about agent activity.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+
+class AgentListener(ABC):
+    """
+    Interface for listening to agent events.
+
+    Implement this interface to receive notifications about what the agent is doing
+    during execution. This is useful for logging, debugging, or providing user feedback.
+
+    All methods are optional - you only need to implement the ones you care about.
+    """
+
+    def on_agent_message(self, message: str) -> None:
+        """
+        Called when the agent sends a complete message.
+
+        This is called with the agent's textual response, which may include
+        reasoning, explanations, or other context.
+
+        Args:
+            message: The complete message from the agent
+        """
+        pass
+
+    def on_tool_call(
+        self,
+        tool_call_id: str,
+        title: str,
+        kind: Optional[str] = None,
+        status: Optional[str] = None,
+    ) -> None:
+        """
+        Called when the agent makes a tool call.
+
+        Tools are things like:
+        - "view" - reading a file
+        - "str-replace-editor" - editing a file
+        - "launch-process" - running a command
+        - "save-file" - creating a new file
+        - "codebase-retrieval" - searching the codebase
+
+        Args:
+            tool_call_id: Unique identifier for this tool call
+            title: Human-readable description of what the tool is doing
+            kind: Category of tool (read, edit, delete, execute, etc.)
+            status: Current status (pending, in_progress, completed, failed)
+        """
+        pass
+
+    def on_tool_response(
+        self,
+        tool_call_id: str,
+        status: Optional[str] = None,
+        content: Optional[Any] = None,
+    ) -> None:
+        """
+        Called when a tool responds with results.
+
+        Args:
+            tool_call_id: Unique identifier for this tool call
+            status: Response status (completed, failed, etc.)
+            content: Response content/results from the tool
+        """
+        pass
+
+    def on_agent_thought(self, text: str) -> None:
+        """
+        Called when the agent shares its internal reasoning.
+
+        This provides insight into how the agent is thinking about the problem.
+
+        Args:
+            text: The thought content from the agent
+        """
+        pass
+
+    def on_function_call(
+        self,
+        function_name: str,
+        arguments: dict,
+    ) -> None:
+        """
+        Called when the agent calls a user-provided function.
+
+        Args:
+            function_name: Name of the function being called
+            arguments: Arguments passed to the function
+        """
+        pass
+
+    def on_function_result(
+        self,
+        function_name: str,
+        result: Any,
+        error: Optional[str] = None,
+    ) -> None:
+        """
+        Called when a user-provided function returns a result or error.
+
+        Args:
+            function_name: Name of the function that was called
+            result: The return value from the function (None if error)
+            error: Error message if the function raised an exception
+        """
+        pass
+
+
+class LoggingAgentListener(AgentListener):
+    """
+    Simple logging implementation of AgentListener.
+
+    Prints all events to stdout with timestamps and formatting.
+    Useful for debugging and understanding what the agent is doing.
+    """
+
+    def __init__(self, verbose: bool = False):
+        """
+        Initialize the logging listener.
+
+        Args:
+            verbose: If True, logs all events. If False, only logs major events.
+        """
+        self.verbose = verbose
+
+    def on_agent_message(self, message: str) -> None:
+        """Log agent messages."""
+        print(f"\n💬 Agent: {message[:200]}{'...' if len(message) > 200 else ''}")
+
+    def on_tool_call(
+        self,
+        tool_call_id: str,
+        title: str,
+        kind: Optional[str] = None,
+        status: Optional[str] = None,
+    ) -> None:
+        """Log tool calls."""
+        if self.verbose or status in ["pending", None]:
+            icon = "🔧" if kind != "read" else "📖"
+            print(f"{icon} Tool: {title} ({kind or 'unknown'})")
+
+    def on_tool_response(
+        self,
+        tool_call_id: str,
+        status: Optional[str] = None,
+        content: Optional[Any] = None,
+    ) -> None:
+        """Log tool responses."""
+        if self.verbose:
+            icon = "✅" if status == "completed" else "❌"
+            print(f"  {icon} Tool response: {status}")
+
+    def on_agent_thought(self, text: str) -> None:
+        """Log agent thoughts."""
+        if self.verbose:
+            print(f"💭 Thinking: {text[:100]}{'...' if len(text) > 100 else ''}")
+
+    def on_function_call(
+        self,
+        function_name: str,
+        arguments: dict,
+    ) -> None:
+        """Log function calls."""
+        print(
+            f"📞 Calling function: {function_name}({', '.join(f'{k}={v}' for k, v in arguments.items())})"
+        )
+
+    def on_function_result(
+        self,
+        function_name: str,
+        result: Any,
+        error: Optional[str] = None,
+    ) -> None:
+        """Log function results."""
+        if error:
+            print(f"  ❌ Function {function_name} failed: {error}")
+        else:
+            result_str = str(result)
+            if len(result_str) > 100:
+                result_str = result_str[:100] + "..."
+            print(f"  ✅ Function {function_name} returned: {result_str}")

augment/listener_adapter.py

@@ -0,0 +1,83 @@
+"""
+Adapter to bridge ACP AgentEventListener to Agent AgentListener.
+"""
+
+from typing import Any, Optional
+
+from .acp import AgentEventListener
+from .listener import AgentListener
+
+
+class AgentListenerAdapter(AgentEventListener):
+    """
+    Adapter that converts ACP events to AgentListener events.
+
+    This class implements AgentEventListener (the ACP interface) and forwards
+    relevant events to an AgentListener (the user-facing interface).
+    """
+
+    def __init__(self, agent_listener: Optional[AgentListener] = None):
+        """
+        Initialize the adapter.
+
+        Args:
+            agent_listener: The user's AgentListener to forward events to
+        """
+        self.agent_listener = agent_listener
+
+    def on_agent_message_chunk(self, text: str) -> None:
+        """
+        Called when the agent sends a message chunk (streaming).
+
+        We don't forward chunks to AgentListener since it doesn't support streaming.
+        We'll accumulate them and send the complete message in on_agent_message.
+        """
+        # AgentListener doesn't support streaming, so we ignore chunks
+        pass
+
+    def on_tool_call(
+        self,
+        tool_call_id: str,
+        title: str,
+        kind: Optional[str] = None,
+        status: Optional[str] = None,
+    ) -> None:
+        """
+        Called when the agent makes a tool call.
+
+        Forward to AgentListener if present.
+        """
+        if self.agent_listener:
+            self.agent_listener.on_tool_call(tool_call_id, title, kind, status)
+
+    def on_tool_response(
+        self,
+        tool_call_id: str,
+        status: Optional[str] = None,
+        content: Optional[Any] = None,
+    ) -> None:
+        """
+        Called when a tool responds with results.
+
+        Forward to AgentListener if present.
+        """
+        if self.agent_listener:
+            self.agent_listener.on_tool_response(tool_call_id, status, content)
+
+    def on_agent_thought(self, text: str) -> None:
+        """
+        Called when the agent shares its internal reasoning.
+
+        Forward to AgentListener if present.
+        """
+        if self.agent_listener:
+            self.agent_listener.on_agent_thought(text)
+
+    def on_agent_message(self, message: str) -> None:
+        """
+        Called when the agent finishes sending a complete message.
+
+        Forward to AgentListener if present.
+        """
+        if self.agent_listener:
+            self.agent_listener.on_agent_message(message)