Aegis: Active Token Viability Governor for LLM Agents

A lightweight, dependency-free Python reference implementation of a Non-Equilibrium Active Viability Governor designed to prevent autonomous agent collapse (such as infinite loops, token runaway, and context-window overflows).

---

The Engineering Challenge

When building autonomous developers or multi-agent routing configurations (such as AutoGPT, CrewAI, or LangGraph), systems are prone to entering unproductive processing loops (e.g., repeatedly failing to parse tool output, calling duplicate command sequences, or self-correcting indefinitely).

• Standard Reactive Limiters implement static thresholds (such as hard maximum token bounds or spend caps). When a cap is hit, the agent is immediately terminated mid-flight (energetic collapse). This wastes the entire token budget, yields no useful output, and breaks user workflows.
• Aegis uses active boundary-distance monitoring. It models the remaining token budget as a thermodynamic resource ($R_t$) and the context window as a dynamic geometric constraint ($C_t$). When an agent begins to loop or runs low on resources, the governor automatically restricts max_tokens limits, downgrades calls to cheap models, or forces memory consolidation (compression) tasks to guarantee successful, graceful task completion.

---

Architecture Overview

1. Viability Potential ($\Pi$): A normalized bottleneck calculation across the token budget, context limits, and looping scores.
2. Urgency Activation ($\Gamma$): A non-linear urgency scaling curve that ramps up dynamically as safety limits are approached: $$\Gamma(\Pi) = \tanh\left(\frac{\kappa}{\Pi + \phi_0}\right)$$
3. Orthant Strategy Separation: Dynamic control decisions mapping either to forward-progress tasks (high completion lengths using premium models) or safety-preservation tasks (history truncation, summaries, cheaper model routing).

---

Getting Started

Prerequisites

• Python 3.8 or newer (Standard Library only, zero external packages required).

Execution

Clone this repository locally, navigate to the directory, and run the main interactive script:

python demo.py

File Layout

• token_governor/core.py: Mathematical calculation of the viability boundaries and urgency multipliers.
• token_governor/agent.py: Discrete-step stateful simulator representing cost tracking, loop tracking, and agent behaviors.
• demo.py: Interactive command-line testing suite and batch comparisons.

---

Transitioning from Simulation to Production

This repository is a simulated reference implementation utilizing decoupled, mock state vectors. To deploy this governor in an enterprise system with live APIs (e.g., OpenAI, Anthropic), you must implement three primary integrations:

1. High-Accuracy Token Counting

Mock token counting must be replaced by true byte-pair encoding (BPE) calculations to prevent unexpected 400 Context Limit errors.

Implementation Pattern (using tiktoken):

import tiktoken

def get_exact_context_tokens(messages: list, model: str = "gpt-4") -> int:
    try:
        encoding = tiktoken.encoding_for_model(model)
    except KeyError:
        encoding = tiktoken.get_encoding("cl100k_base")
        
    num_tokens = 0
    for message in messages:
        num_tokens += 4  # message metadata overhead
        for key, value in message.items():
            num_tokens += len(encoding.encode(value))
    num_tokens += 2  # priming tokens
    return num_tokens

2. Semantic Loop Detection

Replace the simple loop counter with semantic vector embeddings. Repetitive tool calls or redundant thoughts can be detected by monitoring cosine similarities across recent history.

Implementation Pattern (conceptual):

import numpy as np

def calculate_loop_score(recent_thoughts: list[str], embedding_client) -> float:
    if len(recent_thoughts) < 3:
        return 0.0
        
    # Retrieve embeddings for the last few steps
    response = embedding_client.embeddings.create(
        input=recent_thoughts,
        model="text-embedding-3-small"
    )
    vecs = [e.embedding for e in response.data]
    
    # Calculate cosine similarities between consecutive steps
    similarities = []
    for i in range(len(vecs) - 1):
        v1, v2 = np.array(vecs[i]), np.array(vecs[i+1])
        sim = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
        similarities.append(sim)
        
    # High similarity (> 0.9) suggests repetitive semantic states
    avg_similarity = sum(similarities) / len(similarities)
    return float(np.clip((avg_similarity - 0.7) / 0.25, 0.0, 1.0))

3. Asynchronous Concurrency and Rate Limits

Real-world systems run concurrently and are constrained by provider Rate Limits (RPM/TPM). The governor state must be treated as a thread-safe or async resource.

• Thread Safety: If multiple concurrent agent steps are executing, use asyncio.Lock or threading.Lock when modifying the shared budget variable (tokens_remaining).
• Rate-Limit Backoff: Integrate the tenacity library to wrap API calls, ensuring rate limit responses (HTTP 429) trigger non-blocking, exponential back-offs while preserving the viability clock.