agent.py

"""Append-only prompt buffer tracking and ReAct state management."""

S0 = (
    "User Query: Find if Batman has a hidden connection to Spider-Man through an intermediate contact. "
    "Available Tools: Graph_Lookup(character_name) which returns immediate neighbors."
)

T0_A0 = 'I need to check who Batman is directly connected to first. Graph_Lookup("Batman")'
O1 = 'Neighbors of Batman = [Superman]'

T1_A1 = (
    'Batman connects to Superman. Now I need to see if Superman connects closer to Spider-Man or to someone else who does. '
    'Graph_Lookup("Superman")'
)
O2 = 'Neighbors of Superman = [Iron Man]'

T2_A2 = (
    'Superman connects to Iron Man. Let me check Iron Man\'s connections. '
    'Graph_Lookup("Iron Man")'
)
O3 = 'Neighbors of Iron Man = [Spider-Man]'

T3 = (
    'Iron Man connects directly to Spider-Man. Path found: Batman -> Superman -> Iron Man -> Spider-Man.'
)
R = (
    'Yes, Batman has a hidden connection to Spider-Man. '
    'The path is exactly: Batman -> Superman -> Iron Man -> Spider-Man.'
)

BUNDLE_STEP_5 = f"{T2_A2} {O3} {T3} {R}"


def count_words(text: str) -> int:
    """Count words by whitespace separation while ignoring punctuation tokens.

    This preserves tool call identifiers like Graph_Lookup(character_name) as a single token.
    """
    cleaned = text.replace("->", " ")
    for char in "\n\t\r.,:;":
        cleaned = cleaned.replace(char, " ")
    for char in "[]()\"'":
        cleaned = cleaned.replace(char, "")
    return len([token for token in cleaned.split() if token])


def build_context_sequence() -> list[str]:
    """Return the append-only text sequence for the ReAct execution trace."""
    return [S0, T0_A0, O1, T1_A1, O2, BUNDLE_STEP_5]


def compute_step_counts(sequence: list[str]) -> tuple[list[int], list[int]]:
    """Return per-step word counts and cumulative word volumes."""
    step_counts: list[int] = []
    cumulative_volumes: list[int] = []
    total = 0
    for item in sequence:
        current = count_words(item)
        step_counts.append(current)
        total += current
        cumulative_volumes.append(total)
    return step_counts, cumulative_volumes