import os import time import requests import json import argparse import io import logging import types 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 logger = logging.getLogger(__name__) console = Console() # Configuration DEFAULT_AGENT_API = "http://localhost:8080/v1" DEFAULT_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 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 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 = { "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 } system_instruction = prompts.get_system_prompt() messages = [ {"role": "system", "content": system_instruction}, {"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) modules = [] functions = [] variables = [] ACTIVE_VAR_SNIPPET_LEN = 100 for name, val in exec_env.items(): if name.startswith("__"): continue if name == "print": continue 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}") 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---" ) inference_messages = messages.copy() inference_messages.append({"role": "user", "content": dynamic_state_msg}) 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) inference_messages = messages.copy() inference_messages.append({"role": "user", "content": dynamic_state_msg}) new_usage = agent_client.count_tokens(inference_messages) 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) inference_messages = messages.copy() inference_messages.append({"role": "user", "content": dynamic_state_msg}) response_text = agent_client.completion(inference_messages, schema=agent_schema, temperature=0.5) 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 thought = response_json.get("thought", "") action = response_json.get("action", "") content = response_json.get("content", "") if action == "execute_python" and content: content = utils.safeguard_and_repair(agent_client.debug, agent_client, messages, agent_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") break 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}"}) else: messages.append({"role": "user", "content": f"System: Unknown action '{action}'."}) utils.save_agent_trace(trace_filepath, messages, step=step) 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("--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("--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)