Skip to main content

​
Overview

The ContextManager maintains mutable execution state that flows between steps within a single execution unit. It’s the working memory of a running AXON program. 
from axon.runtime import ContextManager

ctx = ContextManager(
    system_prompt="You are a legal expert.",
    tracer=tracer
)

# Store step results
ctx.set_step_result("extract", {"clauses": [1, 2, 3]})

# Retrieve for downstream steps
result = ctx.get_step_result("extract")
print(result)  # {"clauses": [1, 2, 3]}

# Track conversation
ctx.append_message("user", "Analyze this contract")
ctx.append_message("assistant", "I found 3 clauses...")

print(f"Message history: {ctx.message_count} messages")

​
Class: ContextManager

​
Constructor

ContextManager(
    system_prompt: str = "",
    tracer: Tracer | None = None
)
​
system_prompt
str
default:""
The compiled system prompt for this execution unit
​
tracer
Tracer
default:"None"
Optional tracer for recording context mutations
Example: 
from axon.runtime import ContextManager
from axon.runtime.tracer import Tracer

tracer = Tracer(program_name="MyProgram", backend_name="anthropic")
ctx = ContextManager(
    system_prompt="You are an AI assistant.",
    tracer=tracer
)

​
Properties

​
system_prompt: str

The compiled system prompt for this execution unit (read-only). 
print(ctx.system_prompt)
# "You are LegalExpert. Tone: analytical. [CONSTRAINT: NoHallucination]..."

​
current_step: str

The name of the step currently being executed. 
ctx.current_step = "extract_clauses"
print(ctx.current_step)  # "extract_clauses"

​
completed_steps: list[str]

Names of all steps that have recorded results, in insertion order. 
ctx.set_step_result("step1", "result1")
ctx.set_step_result("step2", "result2")

print(ctx.completed_steps)  # ["step1", "step2"]

​
message_count: int

The number of messages in the conversation history. 
ctx.append_message("user", "Hello")
ctx.append_message("assistant", "Hi there")

print(ctx.message_count)  # 2

​
Step Result Management

​
set_step_result(step_name: str, result: Any) -> None

Record the output of a completed step.
​
step_name
str
required
The name of the step that produced the result
​
result
Any
required
The step’s output value (any type)
Raises: ValueError if step_name is empty 
# Store string result
ctx.set_step_result("extract", "Contract terms...")

# Store structured result
ctx.set_step_result("analyze", {
    "risk_score": 0.7,
    "clauses": [1, 2, 3]
})

# Store list result
ctx.set_step_result("classify", ["NDA", "Employment"])

​
get_step_result(step_name: str) -> Any

Retrieve the output of a previously completed step.
​
step_name
str
required
The name of the step whose result is needed
Returns: The step’s output value Raises: KeyError if the step has not completed yet 
try:
    result = ctx.get_step_result("extract")
    print(result)
except KeyError as e:
    print(f"Step not found: {e}")

​
has_step_result(step_name: str) -> bool

Check whether a step has a recorded result. 
if ctx.has_step_result("extract"):
    result = ctx.get_step_result("extract")
else:
    print("Step 'extract' has not completed yet")

​
Variable Bindings

​
set_variable(name: str, value: Any) -> None

Bind a named variable in the execution context.
​
name
str
required
The variable name
​
value
Any
required
The variable’s value
Raises: ValueError if name is empty 
# Store flow parameters
ctx.set_variable("document", contract_text)
ctx.set_variable("max_risk", 0.8)

# Store intermediate values
ctx.set_variable("temp_clauses", extracted_clauses)

​
get_variable(name: str) -> Any

Retrieve a named variable from the execution context.
​
name
str
required
The variable name
Returns: The variable’s value Raises: KeyError if the variable has not been set 
try:
    document = ctx.get_variable("document")
    print(f"Document length: {len(document)}")
except KeyError:
    print("Variable 'document' not found")

​
has_variable(name: str) -> bool

Check whether a named variable exists. 
if ctx.has_variable("document"):
    doc = ctx.get_variable("document")

​
get_variables() -> dict[str, Any]

Return a shallow copy of all variable bindings. 
vars = ctx.get_variables()
print(f"Bound variables: {list(vars.keys())}")

​
Message History

​
append_message(role: str, content: str) -> None

Add a message to the conversation history.
​
role
str
required
One of “system”, “user”, “assistant”
​
content
str
required
The message content
Raises:
  • ValueError if role is not valid
  • ValueError if content is empty
ctx.append_message("user", "Analyze this contract")
ctx.append_message("assistant", "I found 5 key clauses...")
ctx.append_message("user", "What are the risks?")
ctx.append_message("assistant", "The main risks are...")

​
get_message_history() -> list[dict[str, str]]

Return a copy of the full message history. 
history = ctx.get_message_history()
for msg in history:
    print(f"{msg['role']}: {msg['content'][:50]}...")

# Output:
# user: Analyze this contract...
# assistant: I found 5 key clauses...
# user: What are the risks?...
# assistant: The main risks are...

​
clear_messages() -> None

Clear the entire message history. 
ctx.clear_messages()
print(ctx.message_count)  # 0

​
State Snapshots

​
snapshot() -> ContextSnapshot

Capture an immutable snapshot of the current execution state. Returns: ContextSnapshot representing the current state 
snapshot = ctx.snapshot()

print(f"Completed steps: {list(snapshot.step_results.keys())}")
print(f"Message count: {snapshot.message_count}")
print(f"Variables: {list(snapshot.variables.keys())}")
print(f"Current step: {snapshot.current_step}")

​
ContextSnapshot

An immutable point-in-time capture of execution state. 
class ContextSnapshot:
    step_results: dict[str, Any]
    message_count: int
    variables: dict[str, Any]
    current_step: str
    
    def to_dict(self) -> dict[str, Any]:
        ...
Example: 
import json

snapshot = ctx.snapshot()

# Serialize to JSON
with open("context_snapshot.json", "w") as f:
    json.dump(snapshot.to_dict(), f, indent=2)

​
Reset

​
reset() -> None

Clear all state, returning the context to its initial condition. The system prompt is preserved; everything else is cleared. 
ctx.reset()

assert ctx.message_count == 0
assert ctx.completed_steps == []
assert ctx.get_variables() == {}
print(ctx.system_prompt)  # Still present

​
Usage Patterns

​
Step Dependencies

# In AXON:
# flow Analyze(doc: Document) -> Report {
#     step Extract { ... }
#     step Classify { ... }
#     step Synthesize { given: [Extract.output, Classify.output] }
# }

# In the executor:
ctx = ContextManager(system_prompt=unit.system_prompt)

# Step 1: Extract
ctx.current_step = "Extract"
extract_result = await call_model(...)
ctx.set_step_result("Extract", extract_result)

# Step 2: Classify
ctx.current_step = "Classify"
classify_result = await call_model(...)
ctx.set_step_result("Classify", classify_result)

# Step 3: Synthesize (depends on Extract and Classify)
ctx.current_step = "Synthesize"
extract = ctx.get_step_result("Extract")
classify = ctx.get_step_result("Classify")
prompt = f"Given: Extract={extract}, Classify={classify}, create a report"
synthesize_result = await call_model(prompt)

​
Template Substitution

def build_prompt(template: str, ctx: ContextManager) -> str:
    """Replace {{step_name}} with actual step results."""
    prompt = template
    
    for step_name in ctx.completed_steps:
        result = ctx.get_step_result(step_name)
        placeholder = f"{{{{{step_name}}}}}"
        if placeholder in prompt:
            prompt = prompt.replace(placeholder, str(result))
    
    return prompt

# Usage
template = "Given the extracted data: {{Extract}}, classify it"
prompt = build_prompt(template, ctx)
print(prompt)
# "Given the extracted data: [extracted clauses], classify it"

​
Multi-turn Conversation

ctx = ContextManager(system_prompt="You are an expert analyst.")

# Turn 1
ctx.append_message("user", "What are the key clauses?")
response1 = await model.call(
    system_prompt=ctx.system_prompt,
    user_prompt="What are the key clauses?"
)
ctx.append_message("assistant", response1.content)

# Turn 2 (with history)
ctx.append_message("user", "Are there any risks?")
response2 = await model.call(
    system_prompt=ctx.system_prompt,
    user_prompt="Are there any risks?",
    conversation_history=ctx.get_message_history()
)
ctx.append_message("assistant", response2.content)

​
Debugging with Snapshots

import json

ctx = ContextManager(system_prompt="...")

# Capture state at key points
snapshots = []

for step in flow.steps:
    ctx.current_step = step.name
    result = await execute_step(step)
    ctx.set_step_result(step.name, result)
    
    # Capture snapshot after each step
    snapshot = ctx.snapshot()
    snapshots.append(snapshot)

# Save all snapshots for debugging
with open("debug_trace.json", "w") as f:
    json.dump([s.to_dict() for s in snapshots], f, indent=2)

​
Complete Example

from axon.runtime import ContextManager, Executor, ModelClient, ModelResponse
from axon.runtime.tracer import Tracer

class SimpleClient(ModelClient):
    async def call(self, system_prompt: str, user_prompt: str, **kwargs):
        return ModelResponse(content=f"Response to: {user_prompt}")

async def run_flow():
    # Setup
    tracer = Tracer(program_name="Demo", backend_name="test")
    ctx = ContextManager(
        system_prompt="You are a helpful assistant.",
        tracer=tracer
    )
    client = SimpleClient()
    
    # Simulate flow execution
    ctx.set_variable("input_doc", "Contract text...")
    
    # Step 1: Extract
    ctx.current_step = "Extract"
    ctx.append_message("user", "Extract key clauses")
    response1 = await client.call(
        system_prompt=ctx.system_prompt,
        user_prompt="Extract key clauses from contract"
    )
    ctx.append_message("assistant", response1.content)
    ctx.set_step_result("Extract", {"clauses": [1, 2, 3]})
    
    # Step 2: Analyze (depends on Extract)
    ctx.current_step = "Analyze"
    extract_result = ctx.get_step_result("Extract")
    prompt2 = f"Analyze these clauses: {extract_result}"
    ctx.append_message("user", prompt2)
    response2 = await client.call(
        system_prompt=ctx.system_prompt,
        user_prompt=prompt2
    )
    ctx.append_message("assistant", response2.content)
    ctx.set_step_result("Analyze", {"risk_score": 0.7})
    
    # Final snapshot
    snapshot = ctx.snapshot()
    print(f"Completed steps: {snapshot.step_results.keys()}")
    print(f"Message count: {snapshot.message_count}")
    print(f"Variables: {snapshot.variables.keys()}")

await run_flow()

​
Next Steps

Executor API

Orchestrate complete program execution

Memory API

Store and retrieve semantic values

Build docs developers (and LLMs) love

Get started for freeTalk to us