Skip to main content

Python SDK for building AI-powered coding workflows with the Augment CLI agent (auggie)

Project description

Auggie Python SDK

PyPI version Python Versions License: MIT

A Python SDK for interacting with the Augment CLI agent (auggie) programmatically. Build AI-powered coding workflows with type-safe responses, function calling, and real-time event streaming.

What is Augment?

Augment is an AI-powered coding assistant that helps developers write, refactor, and understand code. The Augment Python SDK allows you to programmatically interact with the Augment CLI agent, enabling you to:

  • Automate complex coding workflows
  • Build custom AI-powered development tools
  • Integrate AI assistance into your existing Python applications
  • Create structured, type-safe interactions with AI coding agents

Features

  • Simple API: Clean, intuitive interface for agent interactions
  • Typed Results: Get structured data back from the agent with full type safety
  • Type Inference: Let the agent choose the appropriate type from a list of options
  • Function Calling: Provide Python functions the agent can call during execution
  • Success Criteria: Iteratively verify and correct work against quality standards
  • Event Listeners: Monitor agent activity with AgentListener interface
  • Automatic Retries: Parsing failures are automatically retried with feedback to the agent
  • Session Management: Maintain conversation continuity with context managers
  • Model Selection: Choose which AI model to use
  • Message Access: Get the agent's reasoning via agent.last_model_answer

Installation

Install from PyPI:

pip install auggie-sdk

Prerequisites

  • Python 3.10 or higher
  • Augment CLI (auggie) installed and authenticated
    • Install with: npm install -g @augmentcode/auggie@prerelease
    • Login with: auggie login

Optional Dependencies

For development and testing:

# Install with test dependencies
pip install auggie-sdk[test]

# Install with development dependencies
pip install auggie-sdk[dev]

For detailed installation instructions, including platform-specific guides and troubleshooting, see the examples README.

Quick Start

Basic Usage

from auggie_sdk import Auggie

# Create an agent
agent = Auggie()

# Automatic type inference (agent chooses the best type)
result = agent.run("What is 15 + 27?")
print(f"Answer: {result} (type: {type(result).__name__})")  # Answer: 42 (type: int)

# Explicit typed response
result = agent.run("What is 15 + 27?", return_type=int)
print(f"Answer: {result}")  # Answer: 42

Resource Management

The agent automatically cleans up processes when it goes out of scope, but you can also manage cleanup explicitly:

# Option 1: Context manager (recommended for explicit cleanup)
with Auggie() as agent:
    result = agent.run("What is 15 + 27?")
    print(result)
# Processes automatically cleaned up here

# Option 2: Manual cleanup
agent = Auggie()
try:
    result = agent.run("What is 15 + 27?")
finally:
    agent.close()  # Explicitly cleanup

# Option 3: Automatic cleanup (default)
agent = Auggie()
result = agent.run("What is 15 + 27?")
# Processes cleaned up when agent goes out of scope

Session Management

By default, each call is independent:

from auggie_sdk import Auggie

agent = Auggie()

# Each call creates a fresh session
agent.run("Create a function called add_numbers")
agent.run("Test that function")  # ❌ Won't remember add_numbers!

Use session context manager for conversation continuity:

agent = Auggie()

# Calls within a session share context
with agent.session() as session:
    session.run("Create a function called add_numbers")
    session.run("Test that function")  # ✅ Remembers add_numbers!
    session.run("Add error handling")  # ✅ Still remembers!

Model Selection

# List available models
from auggie_sdk import Auggie

models = Auggie.get_available_models()
for model in models:
    print(f"{model.name} [{model.id}]")
    print(f"  {model.description}")

# Use specific AI model (use model.id from get_available_models())
agent = Auggie(model="sonnet4.5")

# Or with workspace and model
agent = Auggie(
    workspace_root="/path/to/project",
    model="haiku4.5"  # Fast and efficient
)

Automatic Type Inference

When you don't specify a return_type, the agent automatically infers the best type from common types (int, float, bool, str, list, dict):

# Agent automatically determines the type
result = agent.run("What is 2 + 2?")
print(f"Result: {result}, Type: {type(result).__name__}")  # Result: 4, Type: int

result = agent.run("List the primary colors")
print(f"Result: {result}, Type: {type(result).__name__}")  # Result: ['red', 'blue', 'yellow'], Type: list

result = agent.run("Is Python statically typed?")
print(f"Result: {result}, Type: {type(result).__name__}")  # Result: False, Type: bool

Typed Results

from dataclasses import dataclass
from enum import Enum

@dataclass
class Task:
    title: str
    priority: str
    estimated_hours: int

class Priority(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"

# Get structured data back with explicit type
task = agent.run("Create a task: 'Fix login bug', high priority, 8 hours", return_type=Task)
print(f"Task: {task.title}, Priority: {task.priority}")

# Works with lists of objects too
tasks = agent.run("Create 3 example tasks", return_type=list[Task])
for task in tasks:
    print(f"- {task.title}")

# Access the agent's reasoning
print(f"Agent's explanation: {agent.last_model_answer}")

Function Calling

Provide Python functions that the agent can call during execution:

def get_weather(location: str, unit: str = "celsius") -> dict:
    """
    Get the current weather for a location.

    Args:
        location: City name or coordinates
        unit: Temperature unit (celsius or fahrenheit)
    """
    # Your implementation here
    return {"temp": 22, "condition": "sunny"}

def calculate_distance(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    """
    Calculate distance between two coordinates.

    Args:
        lat1: Latitude of first point
        lon1: Longitude of first point
        lat2: Latitude of second point
        lon2: Longitude of second point
    """
    # Your implementation here
    return 100.5

# Agent can call these functions as needed
result = agent.run(
    "What's the weather in San Francisco and how far is it from Los Angeles?",
    return_type=str,
    functions=[get_weather, calculate_distance]
)

Requirements for functions:

  • Must have type hints for parameters
  • Should have docstrings with parameter descriptions
  • The agent will automatically call them when needed

Event Listeners

Monitor what the agent is doing with the AgentListener interface. The listener receives events from the underlying ACP (Agent Client Protocol) layer, including tool calls, agent messages, and function calls:

from auggie_sdk import Auggie, AgentListener, LoggingAgentListener

# Use the built-in logging listener
listener = LoggingAgentListener(verbose=True)
agent = Auggie(listener=listener)

result = agent.run(
    "What's the weather in San Francisco?",
    return_type=dict,
    functions=[get_weather]
)
# Prints: 📞 Calling function: get_weather(location=San Francisco)
#         ✅ Function get_weather returned: {'temp': 72, ...}

# Or create a custom listener
class MyListener(AgentListener):
    def on_function_call(self, function_name: str, arguments: dict) -> None:
        print(f"Agent is calling {function_name}")

    def on_function_result(self, function_name: str, result, error=None) -> None:
        if error:
            print(f"Function failed: {error}")
        else:
            print(f"Function returned: {result}")

agent = Auggie(listener=MyListener())

Available listener methods:

  • on_agent_message(message) - Agent sends a message
  • on_function_call(function_name, arguments) - Agent calls a function
  • on_function_result(function_name, result, error) - Function returns
  • on_tool_call(tool_call_id, title, kind, status) - Agent uses a tool (view, edit, etc.)
  • on_tool_response(tool_call_id, status, content) - Tool responds
  • on_agent_thought(text) - Agent shares internal reasoning

All methods are optional - only implement the ones you need.

Automatic Retries for Parsing Failures

When requesting typed results, the agent automatically retries if parsing fails:

# If the agent's response can't be parsed, it will automatically retry
# up to max_retries times (default: 3)
result = agent.run("What is 2 + 2?", return_type=int)

# Customize retry behavior
result = agent.run(
    "Parse this complex data",
    return_type=MyDataClass,
    max_retries=5  # Try up to 5 times
)

How it works:

  1. Agent sends the instruction
  2. If parsing fails, agent is told about the error and asked to fix the output
  3. This continues until success or max_retries is exhausted
  4. The retry happens in the same session, so the agent has full context

Success Criteria

Ensure the agent iteratively verifies and corrects its work against quality standards:

# Agent will iteratively work and verify until all criteria are met
agent.run(
    "Create a Python function to calculate fibonacci numbers",
    success_criteria=[
        "Function has type hints for all parameters and return value",
        "Function has a comprehensive docstring with examples",
        "Function handles edge cases (n=0, n=1, negative numbers)",
        "Code follows PEP 8 style guidelines",
    ],
    max_verification_rounds=3  # Optional: control max iterations (default: 3)
)

How it works (iterative loop):

  1. Agent works on the task
  2. Agent verifies all success criteria
  3. If all criteria met → Done! ✅
  4. If not all met → Agent receives specific feedback about what's wrong
  5. Agent fixes the identified issues
  6. Go back to step 2
  7. Repeat until all criteria met or max_verification_rounds reached

Benefits:

  • Ensures quality standards are automatically met
  • Agent receives structured feedback on what needs fixing
  • Iterative improvement - agent can make multiple passes
  • Reduces need for manual review
  • Works great for code generation, documentation, and refactoring tasks

Exception handling: If criteria are not met after max_verification_rounds, an AugmentVerificationError is raised with details about unmet criteria and specific issues.

See also: The success criteria example at examples/python-sdk/success_criteria_example.py for detailed examples and best practices.

Supported Types

Automatic type inference (when no return_type specified):

  • int, float, bool, str, list, dict

Explicit types (when return_type is specified):

  • Built-in types: int, float, str, bool, list, dict
  • Dataclasses: Any Python dataclass
  • Enums: Python enums
  • Generic types: list[SomeClass], list[int], etc.

API Reference

Auggie Class

class Auggie:
    def __init__(
        self,
        workspace_root: Optional[Union[str, Path]] = None,
        model: Optional[str] = None,  # Use model.id from get_available_models()
        listener: Optional[AgentListener] = None,
        timeout: int = 180,
    ):
        """
        Initialize an agent instance.

        Args:
            workspace_root: Path to workspace root (defaults to current directory)
            model: AI model ID (e.g., "sonnet4.5", "haiku4.5").
                   Use Auggie.get_available_models() to see all options.
            listener: Optional listener for agent events
            timeout: Default timeout in seconds for run() calls (default: 180)
        """

    def run(
        self,
        instruction: str,
        return_type: Optional[Type[T]] = None,
        timeout: Optional[int] = None,
        max_retries: int = 3,
        functions: Optional[List[Callable]] = None,
        success_criteria: Optional[List[str]] = None,
        max_verification_rounds: int = 3,
    ) -> Union[T, Any]:
        """
        Execute an instruction and return response.

        If return_type is None, automatically infers the type and returns the result.
        If return_type is specified, returns the parsed result of that type.
        Use type(result) to inspect the inferred type when return_type is None.
        """

    def session(self, session_id: Optional[str] = None) -> Auggie:
        """Create a session context for conversation continuity."""

    @staticmethod
    def get_available_models() -> List[Model]:
        """Get the list of available AI models."""

Auggie.get_available_models()

Get the list of available AI models for your account.

from auggie_sdk import Auggie, Model

models = Auggie.get_available_models()
# Returns: List[Model]

# Each Model has:
# - id: str          # Model identifier (e.g., "sonnet4.5")
# - name: str        # Human-readable name (e.g., "Claude Sonnet 4.5")
# - description: str # Additional info (e.g., "Anthropic Claude Sonnet 4.5, 200k context")

Session Usage Patterns

# Single task with multiple related steps
with agent.session() as session:
    session.run("Create the main function")
    session.run("Add error handling")
    session.run("Write comprehensive tests")

# Continue the same task automatically
with agent.session() as session:  # Resumes last session
    session.run("Add more features to the function")

# Work on different concerns with explicit session IDs
with agent.session("backend-work") as backend:
    backend.run("Create API endpoints")

with agent.session("frontend-work") as frontend:
    frontend.run("Create React components")

# Return to backend work
with agent.session("backend-work") as backend:
    backend.run("Add authentication to the API")

Testing with Mock Clients

For testing purposes, you can provide your own ACPClient implementation to Auggie:

from auggie_sdk import Auggie
from auggie_sdk.acp import ACPClient

class MockACPClient(ACPClient):
    """A mock client for testing."""

    def __init__(self):
        self._running = False

    def start(self) -> None:
        self._running = True

    def stop(self) -> None:
        self._running = False

    def send_message(self, message: str, timeout: float = 30.0) -> str:
        # Return mock responses in the expected format with <augment-agent-result> tags
        return "<augment-agent-result>42</augment-agent-result>"

    def clear_context(self) -> None:
        pass

    @property
    def is_running(self) -> bool:
        return self._running

# Use the mock client in tests
mock_client = MockACPClient()
agent = Auggie(acp_client=mock_client)

# Now agent.run() will use your mock client
response = agent.run("Test instruction", return_type=int)
print(response)  # 42

This is useful for:

  • Unit testing code that uses the Agent without calling the real CLI
  • Controlling responses for predictable testing
  • Faster test execution
  • Testing without authentication or network access

Prompt to Code Converter

Convert complex prompts into structured SDK programs with the prompt_to_code.py tool:

# Convert a prompt file to an SDK program
python prompt_to_code.py my_prompt.txt

# With custom output file
python prompt_to_code.py my_prompt.txt --output my_workflow.py

# With custom model
python prompt_to_code.py my_prompt.txt --model sonnet4.5

Why convert prompts to SDK programs?

  • ✅ Better control over workflow execution
  • ✅ Type safety with Python's type system
  • ✅ Debugging capabilities with standard Python tools
  • ✅ Reusability - run the same workflow multiple times
  • ✅ Maintainability - easier to modify and extend

Example: Given a prompt like:

Analyze all Python files in src/, identify security issues,
create a report, and generate fixes for critical issues.

The tool generates a complete Python program with:

  • Proper imports and dataclasses
  • Session management for context
  • Typed results for decision-making
  • Loops for iteration
  • Error handling

See docs/PROMPT_TO_CODE.md for detailed documentation.

Examples

See the examples/python-sdk/ directory in the auggie repo for more usage examples:

Event Listeners

The SDK supports real-time event listeners to observe what the agent is doing:

from auggie_sdk import Auggie
from auggie_sdk.acp import AgentEventListener

class MyListener(AgentEventListener):
    def on_agent_message_chunk(self, text: str) -> None:
        """Called when agent sends response chunks (streaming)."""
        print(f"{text}", end="", flush=True)

    def on_agent_message(self, message: str) -> None:
        """Called when agent finishes sending complete message."""
        print(f"\n[Complete: {len(message)} chars]")

    def on_tool_call(self, tool_call_id, title, kind=None, status=None):
        """Called when agent makes a tool call (read file, edit, etc.)."""
        print(f"\n🔧 Using tool: {title}")

    def on_tool_response(self, tool_call_id, status=None, content=None):
        """Called when tool responds."""
        if status == "completed":
            print("✅ Done!")

    def on_agent_thought(self, text: str) -> None:
        """Called when agent shares its reasoning."""
        print(f"💭 Thinking: {text}")

# Use the listener
listener = MyListener()
agent = Auggie(listener=listener)
response = agent.run("Read the README and summarize it")

For detailed explanation of all events, see: docs/AGENT_EVENT_LISTENER.md

Error Handling

from auggie_sdk.exceptions import AugmentCLIError, AugmentParseError

try:
    result = agent.run("What is the color blue?", return_type=int)
except AugmentParseError as e:
    print(f"Could not parse as int: {e}")
except AugmentCLIError as e:
    print(f"CLI error: {e}")

Key Features

Session Management

The Agent class uses the ACP (Agent Client Protocol) client internally for better performance. By default, each run() call creates a fresh session. Use the session() context manager to maintain conversation continuity:

from auggie_sdk import Auggie

agent = Auggie()

# Use session context for conversation continuity
with agent.session() as session:
    session.run("Create a function")
    session.run("Test it")  # Remembers the function!
    session.run("Optimize it")  # Still remembers!

Real-Time Streaming (Optional)

You can optionally provide an event listener to receive real-time updates:

from auggie_sdk import Auggie
from auggie_sdk.acp import AgentEventListener

class MyListener(AgentEventListener):
    def on_agent_message_chunk(self, text: str):
        print(text, end="", flush=True)

agent = Auggie(listener=MyListener())
agent.run("Create a hello world function")  # See real-time output!

ACP Client (Advanced)

The SDK includes ACP (Agent Client Protocol) clients for more advanced use cases. ACP clients maintain a long-running connection to agents, providing better performance and real-time streaming of responses.

Key Difference: Unlike the Auggie class which spawns a new process per request, ACP clients maintain a single persistent session. All messages sent to the client automatically share context - no need for explicit session management!

Available ACP Clients

  1. AuggieACPClient - For Augment CLI (default)
  2. ClaudeCodeACPClient - For Claude Code via Anthropic's API

Augment CLI (AuggieACPClient)

from auggie_sdk.acp import AuggieACPClient

# Create client with model and workspace configuration
client = AuggieACPClient(
    model="sonnet4.5",
    workspace_root="/path/to/workspace"
)

# Start the agent (creates a persistent session)
client.start()

# Send messages - they all share the same session context!
response1 = client.send_message("Create a function called add_numbers")
print(response1)

response2 = client.send_message("Now test that function")  # Remembers add_numbers!
print(response2)

# Stop the agent
client.stop()

Claude Code (ClaudeCodeACPClient)

Use Claude Code directly via Anthropic's API:

from auggie_sdk.acp import ClaudeCodeACPClient

# Requires: npm install -g @zed-industries/claude-code-acp
# And: export ANTHROPIC_API_KEY=...

client = ClaudeCodeACPClient(
    model="sonnet4.5",
    workspace_root="/path/to/workspace"
)

client.start()
response = client.send_message("Write a Python function to calculate fibonacci")
print(response)
client.stop()

See Claude Code Client Documentation for details.

ACP with Context Manager

from auggie_sdk.acp import AuggieACPClient

# Automatically starts and stops
with AuggieACPClient(model="sonnet4.5") as client:
    response = client.send_message("What is 2 + 2?")
    print(response)

ACP with Event Listener

from auggie_sdk.acp import AuggieACPClient, AgentEventListener

class MyListener(AgentEventListener):
    def on_agent_message_chunk(self, text: str):
        print(f"Agent: {text}", end="", flush=True)

    def on_tool_call(self, tool_name: str, tool_input: dict):
        print(f"\n[Tool: {tool_name}]")

client = AuggieACPClient(
    model="sonnet4.5",
    listener=MyListener()
)
client.start()
response = client.send_message("Create a hello world function")
client.stop()

ACP Client Parameters

  • workspace_root: Workspace root directory (optional, defaults to current directory)
  • model: AI model ID (e.g., "sonnet4.5", "haiku4.5")
  • listener: Event listener for real-time updates (optional)
  • cli_path: Path to the Augment CLI (optional, auto-detected)

Auggie vs ACP Client Comparison

When to use Auggie (subprocess-based):

  • Simple one-off requests
  • Don't need real-time streaming
  • Want explicit session control

When to use AuggieACPClient (long-running):

  • Multiple related requests
  • Want automatic session continuity
  • Need real-time streaming and events
  • Better performance (no subprocess overhead)
# Agent - each call is independent by default
agent = Auggie()
agent.run("Create a function")
agent.run("Test it")  # ❌ Doesn't remember the function

# Agent - need explicit session for context
with agent.session() as session:
    session.run("Create a function")
    session.run("Test it")  # ✅ Remembers the function

# ACP Client - automatic session continuity
client = AuggieACPClient()
client.start()
client.send_message("Create a function")
client.send_message("Test it")  # ✅ Automatically remembers!
client.stop()

See acp_example_usage.py for more examples.

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

Support

License

MIT License - see the LICENSE file for details.

Links

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

auggie_sdk-0.1.10.tar.gz (88.4 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

auggie_sdk-0.1.10-py3-none-any.whl (65.3 kB view details)

Uploaded Python 3

File details

Details for the file auggie_sdk-0.1.10.tar.gz.

File metadata

  • Download URL: auggie_sdk-0.1.10.tar.gz
  • Upload date:
  • Size: 88.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for auggie_sdk-0.1.10.tar.gz
Algorithm Hash digest
SHA256 770b86f066bd5cdd2d68e153ff80403aff9449367565f9d6d7dbf21e3f99df34
MD5 8f467ca57a3257550345b1275cb83e98
BLAKE2b-256 c97fc070d8a99a0db47da7637e4236fa47f72ca8f10ba14c99146c4299ea0a84

See more details on using hashes here.

Provenance

The following attestation bundles were made for auggie_sdk-0.1.10.tar.gz:

Publisher: publish.yml on augmentcode/auggie-sdk-python

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file auggie_sdk-0.1.10-py3-none-any.whl.

File metadata

  • Download URL: auggie_sdk-0.1.10-py3-none-any.whl
  • Upload date:
  • Size: 65.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for auggie_sdk-0.1.10-py3-none-any.whl
Algorithm Hash digest
SHA256 119caa88623f7aa414c8e9db911639ae84d2fb3fef90bccd204e5609ad223abc
MD5 10c1d059c6ef0c4aa27932a815092396
BLAKE2b-256 c89d673e640ab10105e01001f5cc245ce86f20e87ec2dbddf441013dc5bfbc4f

See more details on using hashes here.

Provenance

The following attestation bundles were made for auggie_sdk-0.1.10-py3-none-any.whl:

Publisher: publish.yml on augmentcode/auggie-sdk-python

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page