import os import time import requests import json import argparse import io import logging import re from rich.console import Console from rich.panel import Panel from rich.markdown import Markdown from rich.json import JSON # Local imports 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 (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") 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: def __init__(self, base_url, name="LlamaClient", debug=False): self.base_url = base_url.rstrip("/") self.name = name self.debug = debug self.model = None self.n_ctx = 4096 self._last_prompt_tokens = 0 self._get_model_info() self.max_input_tokens = int(self.n_ctx * MAX_VIRTUAL_CONTEXT_RATIO) self.color = self._determine_color() if self.debug: logger.debug("Connected to %s (%s). Model: %s. Context: %s. Max Input: %s", name, base_url, self.model, self.n_ctx, self.max_input_tokens) def _determine_color(self): if "8080" in self.base_url: return "dodger_blue1" elif "8090" in self.base_url: return "dodger_blue3" else: return "cyan1" def _get_model_info(self): try: resp = requests.get(f"{self.base_url}/models") resp.raise_for_status() data = resp.json() model_data = data.get("data", []) if model_data: model = model_data[0] self.model = model.get("id", "default") meta = model.get("meta", {}) self.n_ctx = meta.get("n_ctx", 4096) else: self.model = "default" except Exception as e: logger.error("[%s] Failed to get model info: %s. Defaulting.", self.name, e) self.model = "default" def count_tokens(self, messages): try: resp = requests.post( f"{self.base_url}/chat/completions/input_tokens", json={"model": self.model, "messages": messages}, timeout=30.0, ) if resp.status_code == 200: return resp.json().get("input_tokens", 0) except Exception as e: logger.debug("[%s] Token count endpoint failed: %s", self.name, e) if self._last_prompt_tokens > 0: logger.debug("[%s] Using cached prompt_tokens estimate: %d", self.name, self._last_prompt_tokens) return self._last_prompt_tokens return sum(len(json.dumps(m)) // 4 + 4 for m in messages) def count_text_tokens(self, text): return self.count_tokens([{"role": "user", "content": text}]) def completion(self, messages, schema=None, temperature=0.1): payload = { "model": self.model, "messages": messages, "temperature": temperature, } if schema: payload["response_format"] = { "type": "json_schema", "json_schema": {"name": "response", "schema": schema} } if self.debug: last_content = messages[-1].get("content", "") if messages else "" console.print(Panel( last_content[-500:] if len(last_content) > 500 else last_content, title=f"Last message to {self.name}", title_align="left", border_style=self.color )) try: resp = requests.post(f"{self.base_url}/chat/completions", json=payload, timeout=120.0) resp.raise_for_status() resp_data = resp.json() content = resp_data["choices"][0]["message"]["content"].strip() usage = resp_data.get("usage", {}) self._last_prompt_tokens = usage.get("prompt_tokens", 0) if self.debug: console.print(Panel( JSON.from_data(content), title=f"{self.name} Response", title_align="left", border_style=self.color )) return content except Exception as e: logger.error("[%s] Error calling LLM: %s", self.name, e) return f"Error: {e}" class AgentTools: def __init__(self, repl_client: LlamaClient, data_content: str): self.client = repl_client self.RAW_CORPUS = data_content def llm_query(self, content_chunk, query): if content_chunk == "RAW_CORPUS": return "ERROR: You passed the string 'RAW_CORPUS' You must pass the CONTENT of the variable (e.g., `chunk = RAW_CORPUS[:1000]`, then `llm_query(chunk, ...)`)." estimated_tokens = len(content_chunk) // 3 if estimated_tokens > (self.client.n_ctx * 2): return f"ERROR: Chunk is massively too large (approx {estimated_tokens} tokens). Slice strictly." chunk_tokens = self.client.count_text_tokens(content_chunk) query_tokens = self.client.count_text_tokens(query) total = chunk_tokens + query_tokens + 150 logger.debug("[Sub-LLM] Processing Query with %d tokens.", total) if total > self.client.n_ctx: msg = f"ERROR: Chunk too large ({chunk_tokens} tokens). Limit is {self.client.n_ctx}. Slice smaller." logger.warning(msg) return msg sub_messages = [ {"role": "system", "content": ( "You are a strict reading assistant. " "Answer the question based ONLY on the provided Context. " "Do not use outside training data. " "If the answer is not in the text, say 'NULL'." )}, {"role": "user", "content": f"Context:\n{content_chunk}\n\nQuestion: {query}"} ] results = self.client.completion(sub_messages) result_tokens = self.client.count_text_tokens(results) logger.debug("[Sub-LLM] Responded with %d tokens.", result_tokens) return results class AgentOutputBuffer: def __init__(self, max_total_chars=20000, max_len_per_print=1009): self._io = io.StringIO() self.max_total_chars = max_total_chars self.max_len_per_print = max_len_per_print self.current_chars = 0 self.global_truncated = False def custom_print(self, *args, **kwargs): temp_io = io.StringIO() print(*args, file=temp_io, **kwargs) text = temp_io.getvalue() if len(text) > self.max_len_per_print: truncated_text = text[:self.max_len_per_print] text = ( f"{truncated_text}\n" f"... [LINE TRUNCATED: Output exceeded {self.max_len_per_print-9} chars. " f"Use slicing or llm_query() to inspect data.] ...\n" ) if self.current_chars + len(text) > self.max_total_chars: remaining = self.max_total_chars - self.current_chars if remaining > 0: self._io.write(text[:remaining]) if not self.global_truncated: self._io.write(f"\n... [SYSTEM HALT: Total output limit ({self.max_total_chars}) reached] ...\n") self.global_truncated = True self.current_chars += len(text) else: self._io.write(text) self.current_chars += len(text) def read_and_clear(self): value = self._io.getvalue() self._io = io.StringIO() self.current_chars = 0 self.global_truncated = False return value def _make_exec_env(tools, out_buffer): return { "RAW_CORPUS": tools.RAW_CORPUS, "llm_query": tools.llm_query, "re": __import__("re"), "math": __import__("math"), "json": __import__("json"), "collections": __import__("collections"), "statistics": __import__("statistics"), "random": __import__("random"), "datetime": __import__("datetime"), "difflib": __import__("difflib"), "string": __import__("string"), "print": out_buffer.custom_print, } 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}"}, ] state = "PLANNING" plan = None turns_since_checkpoint = 0 total_turns = 0 last_diagnosis = "" last_trace_msg_count = 0 while state != "FINISHED": logger.debug("State: %s | Turns since checkpoint: %d / %d", state, turns_since_checkpoint, MAX_REPL_STEPS) if state == "PLANNING": planning_prompt = prompts.get_planning_prompt(task_text) planning_messages = [ {"role": "system", "content": system_instruction}, {"role": "user", "content": planning_prompt}, ] planning_messages = _compress_if_needed(agent_client, planning_messages, planning_messages) response_text = agent_client.completion( planning_messages, schema=planning_schema, temperature=0.3 ) 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 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 plan["current_step_index"] = 0 plan["state"] = "executing" turns_since_checkpoint = 0 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]}"}) logger.info("Plan accepted: goal=%s, steps=%d", plan["goal"], len(plan["strategy"])) state = "EXECUTING" elif state == "EXECUTING": dynamic_state_msg = prompts.get_execution_context( plan, exec_env, turns_since_checkpoint, MAX_REPL_STEPS ) inference_messages = messages.copy() inference_messages.append({"role": "user", "content": dynamic_state_msg}) messages = _compress_if_needed(agent_client, messages, inference_messages) inference_messages = messages.copy() inference_messages.append({"role": "user", "content": dynamic_state_msg}) response_text = agent_client.completion( inference_messages, schema=executing_schema, temperature=0.5 ) 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), }) 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" 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")) 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) 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 ) try: 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."} progress = check_json.get("progress", "NO") justification = check_json.get("justification", "") messages.append({ "role": "assistant", "content": json.dumps(check_json, indent=2, ensure_ascii=False), }) 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." }) 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=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=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() debug = args.debug log_level = logging.DEBUG if debug else logging.INFO log_file = None if debug: os.makedirs("logs", exist_ok=True) timestamp = time.strftime("%Y%m%d-%H%M%S") log_file = os.path.join("logs", f"erlm_debug_{timestamp}.jsonl") setup_logging(level=log_level, log_file=log_file) if log_file: logger.info("JSON log file: %s", log_file) logger.info("Starting EdgeRLM...") context_content = utils.load_file(args.context) logger.debug("Loaded Context: %d characters.", len(context_content)) task_content = args.override_task if args.override_task else utils.load_file(args.task) agent_client = LlamaClient(args.agent_api, "Agent", debug=debug) repl_client = LlamaClient(args.repl_api, "REPL", debug=debug) run_agent(agent_client, repl_client, context_content, task_content)