**feat: refactor agent loop to state machine with dynamic trajectory checks**

- Replace hard-coded `MAX_REPL_STEPS` death-limit with a dynamic Checkpoint Interval.
- Introduce 4-phase state machine: `PLANNING`, `EXECUTING`, `TRAJECTORY_CHECK`, and `PIVOTING`.
- Add Trajectory Checks allowing the agent to self-assess progress against defined success criteria to extend its execution budget.
- Implement Hard (wipe strategy/kernel) and Soft (retry step) Pivoting logic.
- Establish persistent notebook-style kernel (flushed only on Hard Pivots).
- Upgrade sandbox safety with strict import whitelisting, dunder blocking, and taint tracking.

**bugfix:

- logging_config.py: Add markdown_it to suppressed logger list
- utils.py: Rewrite save_agent_trace for delta-based logging
- edge_rlm.py: Add last_trace_msg_count tracking and pass deltas at all 3 call sites
This commit is contained in:
2026-07-18 20:47:44 +01:00
parent 834183807a
commit f70c9b2054
6 changed files with 591 additions and 135 deletions
+391 -125
View File
@@ -5,7 +5,7 @@ import json
import argparse
import io
import logging
import types
import re
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
@@ -17,17 +17,35 @@ from logging_config import setup_logging
import utils
import prompts
def _load_dotenv(path=".env"):
try:
with open(path) as f:
for line in f:
line = line.strip()
if not line or line.startswith("#") or "=" not in line:
continue
key, _, val = line.partition("=")
key, val = key.strip(), val.strip()
if (val.startswith('"') and val.endswith('"')) or (val.startswith("'") and val.endswith("'")):
val = val[1:-1]
os.environ.setdefault(key, val)
except FileNotFoundError:
pass
_load_dotenv()
logger = logging.getLogger(__name__)
console = Console()
# Configuration
DEFAULT_AGENT_API = "http://localhost:8080/v1"
DEFAULT_REPL_API = "http://localhost:8090/v1"
# Configuration (from .env with fallback defaults)
AGENT_API = os.getenv("AGENT_API", "http://localhost:8080/v1")
REPL_API = os.getenv("REPL_API", "http://localhost:8090/v1")
DEFAULT_CONTEXT_FILE = "context.txt"
DEFAULT_TASK_FILE = "task.txt"
MAX_REPL_STEPS = 20
MAX_VIRTUAL_CONTEXT_RATIO = 0.85
CONTEXT_FILE = os.getenv("CONTEXT_FILE", "context.txt")
TASK_FILE = os.getenv("TASK_FILE", "task.txt")
MAX_REPL_STEPS = int(os.getenv("MAX_REPL_STEPS", "20"))
MAX_VIRTUAL_CONTEXT_RATIO = float(os.getenv("MAX_VIRTUAL_CONTEXT_RATIO", "0.85"))
class LlamaClient:
@@ -207,24 +225,8 @@ class AgentOutputBuffer:
self.global_truncated = False
return value
def run_agent(agent_client, repl_client, context_text, task_text):
tools = AgentTools(repl_client, context_text)
agent_schema = {
"type": "object",
"properties": {
"thought": {"type": "string", "description": "Reasoning about current state and what to do next."},
"action": {"type": "string", "enum": ["execute_python", "final_answer"]},
"content": {"type": "string", "description": "Python code or Final Answer text."}
},
"required": ["thought", "action", "content"]
}
out_buffer = AgentOutputBuffer()
trace_filepath = utils.init_trace_file()
exec_env = {
def _make_exec_env(tools, out_buffer):
return {
"RAW_CORPUS": tools.RAW_CORPUS,
"llm_query": tools.llm_query,
"re": __import__("re"),
@@ -236,150 +238,414 @@ def run_agent(agent_client, repl_client, context_text, task_text):
"datetime": __import__("datetime"),
"difflib": __import__("difflib"),
"string": __import__("string"),
"print": out_buffer.custom_print
"print": out_buffer.custom_print,
}
system_instruction = prompts.get_system_prompt()
def _compress_if_needed(agent_client, messages, inference_messages):
usage = agent_client.count_tokens(inference_messages)
logger.debug("Context Usage: %d / %d", usage, agent_client.max_input_tokens)
if usage > agent_client.max_input_tokens:
logger.warning("Context limit exceeded. Triggering History Compression.")
messages = utils.compress_history(agent_client, messages, keep_last_pairs=2)
rebuilt = messages + inference_messages[-1:]
new_usage = agent_client.count_tokens(rebuilt)
logger.debug("Context Usage after compression: %d", new_usage)
if new_usage > agent_client.max_input_tokens:
logger.error("Compression insufficient. Forcing hard truncation.")
messages.pop(2)
return messages
def run_agent(agent_client, repl_client, context_text, task_text):
tools = AgentTools(repl_client, context_text)
planning_schema = {
"type": "object",
"properties": {
"goal": {"type": "string", "description": "The high-level objective."},
"strategy": {
"type": "array",
"items": {"type": "string"},
"description": "Discrete, programmatic actions.",
},
"success_criteria": {
"type": "array",
"items": {"type": "string"},
"description": "Success criteria parallel to strategy.",
},
},
"required": ["goal", "strategy", "success_criteria"],
}
executing_schema = {
"type": "object",
"properties": {
"thought": {"type": "string", "description": "Reasoning about current state and what to do next."},
"action": {"type": "string", "enum": ["execute_python", "final_answer"]},
"content": {"type": "string", "description": "Python code or Final Answer text."},
"step_completed": {
"type": "boolean",
"description": "Set to true when the current strategy step is fully complete and you are ready to advance.",
},
},
"required": ["thought", "action", "content"],
}
trajectory_schema = {
"type": "object",
"properties": {
"progress": {"type": "string", "enum": ["YES", "NO"]},
"justification": {"type": "string", "description": "Brief explanation of the assessment."},
},
"required": ["progress", "justification"],
}
pivot_schema = {
"type": "object",
"properties": {
"diagnosis": {"type": "string", "description": "Why progress stalled."},
"new_strategy": {
"type": "array",
"items": {"type": "string"},
"description": "Revised strategy (omit for Soft Pivot).",
},
"new_success_criteria": {
"type": "array",
"items": {"type": "string"},
"description": "Revised success criteria (omit for Soft Pivot).",
},
},
"required": ["diagnosis"],
}
out_buffer = AgentOutputBuffer()
trace_filepath = utils.init_trace_file()
exec_env = _make_exec_env(tools, out_buffer)
system_instruction = prompts.get_system_prompt(max_repl_steps=MAX_REPL_STEPS)
messages = [
{"role": "system", "content": system_instruction},
{"role": "user", "content": f"USER TASK: {task_text}"}
{"role": "user", "content": f"USER TASK: {task_text}"},
]
step = 0
while step < MAX_REPL_STEPS:
step += 1
logger.debug("Step %d of %d", step, MAX_REPL_STEPS)
state = "PLANNING"
plan = None
turns_since_checkpoint = 0
total_turns = 0
last_diagnosis = ""
last_trace_msg_count = 0
modules = []
functions = []
variables = []
ACTIVE_VAR_SNIPPET_LEN = 100
while state != "FINISHED":
logger.debug("State: %s | Turns since checkpoint: %d / %d", state, turns_since_checkpoint, MAX_REPL_STEPS)
for name, val in exec_env.items():
if name.startswith("__"): continue
if name == "print": continue
if state == "PLANNING":
planning_prompt = prompts.get_planning_prompt(task_text)
if isinstance(val, types.ModuleType):
modules.append(name)
elif callable(val):
functions.append(name)
else:
type_name = type(val).__name__
s_val = str(val)
snippet = (s_val[:ACTIVE_VAR_SNIPPET_LEN] + '...') if len(s_val) > ACTIVE_VAR_SNIPPET_LEN else s_val
variables.append(f"{name} ({type_name}): {snippet}")
planning_messages = [
{"role": "system", "content": system_instruction},
{"role": "user", "content": planning_prompt},
]
dynamic_state_msg = (
f"[SYSTEM STATE REMINDER]\n"
f"Current Step: {step}/{MAX_REPL_STEPS}\n"
f"Available Libraries: {', '.join(modules)}\n"
f"Available Tools: {', '.join(functions)}\n"
f"Active Variables:\n" + ("\n".join([f" - {v}" for v in variables]) if variables else " (None)") + "\n---"
)
planning_messages = _compress_if_needed(agent_client, planning_messages, planning_messages)
inference_messages = messages.copy()
inference_messages.append({"role": "user", "content": dynamic_state_msg})
response_text = agent_client.completion(
planning_messages, schema=planning_schema, temperature=0.3
)
usage = agent_client.count_tokens(inference_messages)
logger.debug("Context Usage: %d / %d", usage, agent_client.max_input_tokens)
try:
plan = json.loads(response_text)
except json.JSONDecodeError:
logger.error("Planning: JSON parse error. Retrying.")
messages.append({"role": "user", "content": "System: Invalid PlanObject JSON. Please retry."})
continue
if usage > agent_client.max_input_tokens:
logger.warning("Context limit exceeded. Triggering History Compression.")
valid, err = utils.validate_plan_object(plan)
if not valid:
logger.error("Planning: Invalid PlanObject: %s", err)
messages.append({"role": "user", "content": f"System: Invalid PlanObject: {err}. Please retry."})
continue
messages = utils.compress_history(agent_client, messages, keep_last_pairs=2)
plan["current_step_index"] = 0
plan["state"] = "executing"
turns_since_checkpoint = 0
inference_messages = messages.copy()
inference_messages.append({"role": "user", "content": dynamic_state_msg})
messages.append({"role": "assistant", "content": json.dumps(plan, indent=2, ensure_ascii=False)})
messages.append({"role": "user", "content": f"System: Plan accepted. Starting execution of step 1/{len(plan['strategy'])}: {plan['strategy'][0]}"})
new_usage = agent_client.count_tokens(inference_messages)
logger.debug("Context Usage after compression: %d", new_usage)
logger.info("Plan accepted: goal=%s, steps=%d", plan["goal"], len(plan["strategy"]))
state = "EXECUTING"
if new_usage > agent_client.max_input_tokens:
logger.error("Compression insufficient. Forcing hard truncation.")
messages.pop(2)
inference_messages = messages.copy()
inference_messages.append({"role": "user", "content": dynamic_state_msg})
elif state == "EXECUTING":
dynamic_state_msg = prompts.get_execution_context(
plan, exec_env, turns_since_checkpoint, MAX_REPL_STEPS
)
response_text = agent_client.completion(inference_messages, schema=agent_schema, temperature=0.5)
inference_messages = messages.copy()
inference_messages.append({"role": "user", "content": dynamic_state_msg})
try:
response_json = json.loads(response_text)
except json.JSONDecodeError:
logger.error("JSON Parse Error")
messages.append({"role": "user", "content": "System: Invalid JSON returned. Please retry."})
continue
messages = _compress_if_needed(agent_client, messages, inference_messages)
thought = response_json.get("thought", "")
action = response_json.get("action", "")
content = response_json.get("content", "")
inference_messages = messages.copy()
inference_messages.append({"role": "user", "content": dynamic_state_msg})
if action == "execute_python" and content:
content = utils.safeguard_and_repair(agent_client.debug, agent_client, messages, agent_schema, content)
response_text = agent_client.completion(
inference_messages, schema=executing_schema, temperature=0.5
)
if agent_client.debug:
console.print(Panel(
try:
response_json = json.loads(response_text)
except json.JSONDecodeError:
logger.error("EXECUTING: JSON Parse Error")
messages.append({"role": "user", "content": "System: Invalid JSON returned. Please retry."})
continue
thought = response_json.get("thought", "")
action = response_json.get("action", "")
content = response_json.get("content", "")
step_completed = response_json.get("step_completed", False)
if action == "execute_python" and content:
content = utils.safeguard_and_repair(
agent_client.debug, agent_client, messages, executing_schema, content
)
if agent_client.debug:
console.print(Panel(
f"[italic]{thought}[/italic]",
title="Agent Thought",
title_align="left",
border_style="magenta"
))
messages.append({"role": "assistant", "content": json.dumps(response_json, indent=2, ensure_ascii=False)})
messages.append({
"role": "assistant",
"content": json.dumps(response_json, indent=2, ensure_ascii=False),
})
if action == "final_answer":
logger.debug("Raw Agent Output: %s", content[:200])
if action == "final_answer":
logger.debug("Raw Agent Output: %s", content[:200])
final_report = utils.generate_final_report(agent_client, task_text, content)
final_report_md = Markdown(final_report)
print("\n\n")
console.print(final_report_md)
print("\n")
state = "FINISHED"
final_report = utils.generate_final_report(agent_client, task_text, content)
elif action == "execute_python":
if agent_client.debug and content != response_json.get("content"):
console.print(Panel(content, title="Executing Code via Safeguard", title_align="left", border_style="cyan"))
elif agent_client.debug and content == response_json.get("content"):
console.print(Panel(content, title="Executing Code", title_align="left", border_style="yellow"))
final_report_md = Markdown(final_report)
print("\n\n")
console.print(final_report_md)
print("\n")
break
observation = ""
try:
out_buffer.read_and_clear()
exec(content, exec_env)
observation = out_buffer.read_and_clear()
if not observation:
observation = "Code executed successfully (no output)."
except Exception as e:
observation = f"Python Error: {e}"
logger.error("Code Execution Error: %s", e)
elif action == "execute_python":
if agent_client.debug and content != response_json.get("content"):
console.print(Panel(content, title="Executing Code via Safeguard", title_align="left", border_style="cyan"))
elif agent_client.debug and content == response_json.get("content"):
console.print(Panel(content, title="Executing Code", title_align="left", border_style="yellow"))
if agent_client.debug:
console.print(Panel(
f"{observation.strip()}",
title="Observation",
title_align="left",
border_style="dark_green"
))
messages.append({"role": "user", "content": f"Observation:\n{observation}"})
total_turns += 1
turns_since_checkpoint += 1
if step_completed:
plan["current_step_index"] += 1
if plan["current_step_index"] >= len(plan["strategy"]):
plan["current_step_index"] = len(plan["strategy"]) - 1
messages.append({
"role": "user",
"content": "System: All plan steps completed. Provide your final answer."
})
else:
next_idx = plan["current_step_index"]
messages.append({
"role": "user",
"content": (
f"System: Advancing to step {next_idx + 1}/{len(plan['strategy'])}: "
f"{plan['strategy'][next_idx]}. "
f"Success criteria: {plan['success_criteria'][next_idx]}"
)
})
turns_since_checkpoint = 0
if turns_since_checkpoint >= MAX_REPL_STEPS:
state = "TRAJECTORY_CHECK"
messages.append({
"role": "user",
"content": "System: Checkpoint reached. Pausing for trajectory review."
})
else:
messages.append({"role": "user", "content": f"System: Unknown action '{action}'."})
trace_delta = messages[last_trace_msg_count:] if last_trace_msg_count <= len(messages) else messages
utils.save_agent_trace(
trace_filepath, trace_delta, step=total_turns,
state=state, plan_step=plan["current_step_index"],
total_messages=len(messages),
)
last_trace_msg_count = len(messages)
elif state == "TRAJECTORY_CHECK":
step_idx = plan["current_step_index"]
criteria = plan["success_criteria"][step_idx]
check_prompt = prompts.get_trajectory_check_prompt(criteria)
check_messages = [
{"role": "system", "content": system_instruction},
{"role": "user", "content": check_prompt},
]
if len(messages) >= 4:
check_messages += messages[-4:]
else:
check_messages += messages
response_text = agent_client.completion(
check_messages, schema=trajectory_schema, temperature=0.2
)
observation = ""
try:
out_buffer.read_and_clear()
exec(content, exec_env)
observation = out_buffer.read_and_clear()
check_json = json.loads(response_text)
except json.JSONDecodeError:
logger.error("TRAJECTORY_CHECK: JSON parse error. Assuming NO.")
check_json = {"progress": "NO", "justification": "JSON parse failure."}
if not observation:
observation = "Code executed successfully (no output)."
except Exception as e:
observation = f"Python Error: {e}"
logger.error("Code Execution Error: %s", e)
progress = check_json.get("progress", "NO")
justification = check_json.get("justification", "")
if agent_client.debug:
console.print(Panel(
f"{observation.strip()}",
title="Observation",
title_align="left",
border_style="dark_green"
))
messages.append({"role": "user", "content": f"Observation:\n{observation}"})
messages.append({
"role": "assistant",
"content": json.dumps(check_json, indent=2, ensure_ascii=False),
})
else:
messages.append({"role": "user", "content": f"System: Unknown action '{action}'."})
if progress == "YES":
logger.info("Trajectory check PASSED: %s", justification[:100])
turns_since_checkpoint = 0
state = "EXECUTING"
messages.append({
"role": "user",
"content": (
f"System: Trajectory check passed. Resuming execution. "
f"Turns since checkpoint reset to 0."
)
})
else:
logger.warning("Trajectory check FAILED: %s", justification[:100])
last_diagnosis = justification
state = "PIVOTING"
messages.append({
"role": "user",
"content": "System: Trajectory check failed. Transitioning to pivot."
})
utils.save_agent_trace(trace_filepath, messages, step=step)
trace_delta = messages[last_trace_msg_count:] if last_trace_msg_count <= len(messages) else messages
utils.save_agent_trace(
trace_filepath, trace_delta, step=total_turns,
state=state, plan_step=plan["current_step_index"],
total_messages=len(messages),
)
last_trace_msg_count = len(messages)
elif state == "PIVOTING":
pivot_prompt = prompts.get_pivot_prompt(last_diagnosis, task_text)
pivot_messages = [
{"role": "system", "content": system_instruction},
{"role": "user", "content": pivot_prompt},
]
if len(messages) >= 4:
pivot_messages += messages[-4:]
else:
pivot_messages += messages
response_text = agent_client.completion(
pivot_messages, schema=pivot_schema, temperature=0.4
)
try:
pivot_json = json.loads(response_text)
except json.JSONDecodeError:
logger.error("PIVOTING: JSON parse error. Forcing Soft Pivot.")
pivot_json = {"diagnosis": "JSON parse error. Retrying current step."}
valid, err = utils.validate_pivot_object(pivot_json)
if not valid:
logger.error("PIVOTING: Invalid PivotObject: %s. Forcing Soft Pivot.", err)
pivot_json = {"diagnosis": pivot_json.get("diagnosis", "Validation error: " + err)}
has_new_strategy = "new_strategy" in pivot_json and "new_success_criteria" in pivot_json
messages.append({
"role": "assistant",
"content": json.dumps(pivot_json, indent=2, ensure_ascii=False),
})
if has_new_strategy:
logger.info("Hard Pivot: new strategy with %d steps.", len(pivot_json["new_strategy"]))
plan = {
"goal": plan["goal"],
"strategy": pivot_json["new_strategy"],
"success_criteria": pivot_json["new_success_criteria"],
"current_step_index": 0,
"state": "executing",
}
exec_env = _make_exec_env(tools, out_buffer)
messages.append({
"role": "user",
"content": (
f"System: Hard Pivot applied. Kernel reset. "
f"New plan has {len(plan['strategy'])} steps. "
f"Starting step 1: {plan['strategy'][0]}"
)
})
else:
logger.info("Soft Pivot: retrying current step with fresh budget.")
messages.append({
"role": "user",
"content": (
f"System: Soft Pivot applied. Kernel state preserved. "
f"Resuming step {plan['current_step_index'] + 1}/{len(plan['strategy'])}: "
f"{plan['strategy'][plan['current_step_index']]}"
)
})
turns_since_checkpoint = 0
state = "EXECUTING"
trace_delta = messages[last_trace_msg_count:] if last_trace_msg_count <= len(messages) else messages
utils.save_agent_trace(
trace_filepath, trace_delta, step=total_turns,
state=state, plan_step=plan["current_step_index"],
total_messages=len(messages),
)
last_trace_msg_count = len(messages)
logger.info("Agent finished. Total execution turns: %d", total_turns)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="""Edge Recursive Language Model
A sophisticated data extraction and analysis tool that mimics the process of a human data scientist, carefully exploring and structuring a large dataset before performing targeted queries.""")
parser.add_argument("--context", default=DEFAULT_CONTEXT_FILE, help="Path to text file to process")
parser.add_argument("--task", default=DEFAULT_TASK_FILE, help="Path to task instruction file")
parser.add_argument("--context", default=CONTEXT_FILE, help="Path to text file to process")
parser.add_argument("--task", default=TASK_FILE, help="Path to task instruction file")
parser.add_argument("--override_task", help="Direct string override for the task")
parser.add_argument("--agent_api", default=DEFAULT_AGENT_API, help="URL for the Main Agent LLM")
parser.add_argument("--repl_api", default=DEFAULT_REPL_API, help="URL for the Sub-call/REPL LLM")
parser.add_argument("--agent_api", default=AGENT_API, help="URL for the Main Agent LLM")
parser.add_argument("--repl_api", default=REPL_API, help="URL for the Sub-call/REPL LLM")
parser.add_argument("--debug", action="store_true", help="Enable verbose debug logging and JSON log file")
args = parser.parse_args()