From 834183807abf269e97204734121a4308726d8c1b Mon Sep 17 00:00:00 2001 From: Morpheus Sandmann Date: Fri, 17 Jul 2026 21:26:24 +0100 Subject: [PATCH] Improved token counting for OpenAI endpoints (llama.cpp token count -> cached prompt_tokens from last output -> naive string estimate). Logger refactoring with JSON format debugging to log file. Various bugfixes. Removed unused templates.py --- edge_rlm.py | 94 ++++++++++++++++++++++++++++++----------------- logging_config.py | 59 ++++++++++++++++++++++------- requirements.txt | 2 + templates.py | 18 --------- utils.py | 34 ++++++++--------- 5 files changed, 123 insertions(+), 84 deletions(-) create mode 100644 requirements.txt delete mode 100644 templates.py diff --git a/edge_rlm.py b/edge_rlm.py index ab71023..b86ca39 100644 --- a/edge_rlm.py +++ b/edge_rlm.py @@ -1,3 +1,4 @@ +import os import time import requests import json @@ -14,7 +15,7 @@ from rich.json import JSON # Local imports from logging_config import setup_logging import utils -import prompts as prompts +import prompts logger = logging.getLogger(__name__) console = Console() @@ -30,15 +31,17 @@ MAX_VIRTUAL_CONTEXT_RATIO = 0.85 class LlamaClient: - def __init__(self, base_url, name="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 debug: logger.debug(f"Connected to {name} ({base_url}). Model: {self.model}. Context: {self.n_ctx}. Max Input: {self.max_input_tokens}") + 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: @@ -62,7 +65,7 @@ class LlamaClient: else: self.model = "default" except Exception as e: - logger.error(f"[{self.name}] Failed to get model info: {e}. Defaulting.") + logger.error("[%s] Failed to get model info: %s. Defaulting.", self.name, e) self.model = "default" def count_tokens(self, messages): @@ -75,7 +78,12 @@ class LlamaClient: if resp.status_code == 200: return resp.json().get("input_tokens", 0) except Exception as e: - logger.debug(f"[{self.name}] Token count failed: {e}. Using estimate.") + 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): @@ -92,7 +100,7 @@ class LlamaClient: "type": "json_schema", "json_schema": {"name": "response", "schema": schema} } - if debug: + 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, @@ -101,10 +109,13 @@ class LlamaClient: border_style=self.color )) try: - resp = requests.post(f"{self.base_url}/chat/completions", json=payload) + resp = requests.post(f"{self.base_url}/chat/completions", json=payload, timeout=120.0) resp.raise_for_status() - content = resp.json()["choices"][0]["message"]["content"].strip() - if debug: + 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", @@ -113,7 +124,7 @@ class LlamaClient: )) return content except Exception as e: - logger.error(f"[{self.name}] Error calling LLM: {e}") + logger.error("[%s] Error calling LLM: %s", self.name, e) return f"Error: {e}" class AgentTools: @@ -133,7 +144,7 @@ class AgentTools: query_tokens = self.client.count_text_tokens(query) total = chunk_tokens + query_tokens + 150 - if debug: logger.debug(f"[Sub-LLM] Processing Query with {total} tokens.") + 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." @@ -151,7 +162,7 @@ class AgentTools: ] results = self.client.completion(sub_messages) result_tokens = self.client.count_text_tokens(results) - if debug: logger.debug(f"[Sub-LLM] Responded with {result_tokens} tokens.") + logger.debug("[Sub-LLM] Responded with %d tokens.", result_tokens) return results class AgentOutputBuffer: @@ -211,7 +222,7 @@ def run_agent(agent_client, repl_client, context_text, task_text): out_buffer = AgentOutputBuffer() - trace_filepath = utils.init_trace_file(debug) + trace_filepath = utils.init_trace_file() exec_env = { "RAW_CORPUS": tools.RAW_CORPUS, @@ -239,7 +250,7 @@ def run_agent(agent_client, repl_client, context_text, task_text): step = 0 while step < MAX_REPL_STEPS: step += 1 - if debug: logger.debug(f"Step {step} of {MAX_REPL_STEPS}") + logger.debug("Step %d of %d", step, MAX_REPL_STEPS) modules = [] functions = [] @@ -272,19 +283,24 @@ def run_agent(agent_client, repl_client, context_text, task_text): inference_messages.append({"role": "user", "content": dynamic_state_msg}) usage = agent_client.count_tokens(inference_messages) - if debug: logger.debug(f"Context Usage: {usage} / {agent_client.max_input_tokens}") + logger.debug("Context Usage: %d / %d", usage, agent_client.max_input_tokens) if usage > agent_client.max_input_tokens: - if debug: logger.warning("Context limit exceeded. Triggering History Compression.") + logger.warning("Context limit exceeded. Triggering History Compression.") - messages = utils.compress_history(debug, agent_client, messages, keep_last_pairs=2) + 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) - if debug: logger.debug(f"Context Usage after compression: {new_usage}") + 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) @@ -300,9 +316,9 @@ def run_agent(agent_client, repl_client, context_text, task_text): content = response_json.get("content", "") if action == "execute_python" and content: - content = utils.safeguard_and_repair(debug, agent_client, messages, agent_schema, content) + content = utils.safeguard_and_repair(agent_client.debug, agent_client, messages, agent_schema, content) - if debug: + if agent_client.debug: console.print(Panel( f"[italic]{thought}[/italic]", title="Agent Thought", @@ -313,9 +329,9 @@ def run_agent(agent_client, repl_client, context_text, task_text): messages.append({"role": "assistant", "content": json.dumps(response_json, indent=2, ensure_ascii=False)}) if action == "final_answer": - if debug: logger.debug(f"Raw Agent Output: {content}") + logger.debug("Raw Agent Output: %s", content[:200]) - final_report = utils.generate_final_report(debug, agent_client, task_text, content) + final_report = utils.generate_final_report(agent_client, task_text, content) final_report_md = Markdown(final_report) print("\n\n") @@ -324,9 +340,9 @@ def run_agent(agent_client, repl_client, context_text, task_text): break elif action == "execute_python": - if debug and content != response_json.get("content"): + 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 debug and content == response_json.get("content"): + elif agent_client.debug and content == response_json.get("content"): console.print(Panel(content, title="Executing Code", title_align="left", border_style="yellow")) observation = "" @@ -339,9 +355,9 @@ def run_agent(agent_client, repl_client, context_text, task_text): observation = "Code executed successfully (no output)." except Exception as e: observation = f"Python Error: {e}" - logger.error(f"Code Execution Error: {e}") + logger.error("Code Execution Error: %s", e) - if debug: + if agent_client.debug: console.print(Panel( f"{observation.strip()}", title="Observation", @@ -353,7 +369,7 @@ def run_agent(agent_client, repl_client, context_text, task_text): else: messages.append({"role": "user", "content": f"System: Unknown action '{action}'."}) - utils.save_agent_trace(trace_filepath, messages) + utils.save_agent_trace(trace_filepath, messages, step=step) if __name__ == "__main__": parser = argparse.ArgumentParser(description="""Edge Recursive Language Model @@ -364,19 +380,29 @@ if __name__ == "__main__": 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") + 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 - setup_logging(level=log_level, debug=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...") - if debug: logger.info("Starting EdgeRLM...") context_content = utils.load_file(args.context) - if debug: logger.debug(f"Loaded Context: {len(context_content)} characters.") + 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") - repl_client = LlamaClient(args.repl_api, "REPL") + 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) diff --git a/logging_config.py b/logging_config.py index a3cf378..7d9b571 100644 --- a/logging_config.py +++ b/logging_config.py @@ -1,25 +1,56 @@ -import sys +import json import logging +import logging.handlers +from datetime import datetime, timezone from rich.logging import RichHandler -from rich.console import Console -def setup_logging(level=logging.INFO, debug=False): - - # silence noisy libraries - for lib_name in ("urllib3","requests","http.client","markdown","Markdown"): + +class JSONFormatter(logging.Formatter): + def format(self, record): + log_entry = { + "timestamp": datetime.now(timezone.utc).isoformat(), + "level": record.levelname, + "logger": record.name, + "message": record.getMessage(), + "module": record.module, + "function": record.funcName, + "line": record.lineno, + } + standard_keys = set(logging.LogRecord("", 0, "", 0, "", (), None).__dict__.keys()) + for key, value in record.__dict__.items(): + if key not in standard_keys and key not in ("message", "msg", "args"): + log_entry[key] = value + if record.exc_info and record.exc_info[0]: + log_entry["exception"] = self.formatException(record.exc_info) + return json.dumps(log_entry, default=str, ensure_ascii=False) + + +def setup_logging(level=logging.INFO, log_file=None): + for lib_name in ("urllib3", "requests", "http.client", "markdown", "Markdown"): logging.getLogger(lib_name).setLevel(logging.WARNING) + handlers = [ + RichHandler( + rich_tracebacks=True, + show_path=False, + log_time_format="[%H:%M:%S]", + markup=True, + ) + ] + + if log_file: + file_handler = logging.handlers.RotatingFileHandler( + log_file, maxBytes=10_000_000, backupCount=3, encoding="utf-8" + ) + file_handler.setFormatter(JSONFormatter()) + handlers.append(file_handler) + logging.basicConfig( level=level, format="%(message)s", datefmt="[%X]", - handlers=[RichHandler( - rich_tracebacks=True, - show_path=False, - log_time_format="[%H:%M:%S]", - markup=True - )], + handlers=handlers, ) - if debug: - logging.getLogger(__name__).debug("[dim]Debug mode active.[/dim]") \ No newline at end of file + if level == logging.DEBUG: + logging.getLogger(__name__).debug("[dim]Debug mode active.[/dim]") diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..3f937bd --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +requests>=2.28.0 +rich>=13.0.0 diff --git a/templates.py b/templates.py deleted file mode 100644 index 931c04b..0000000 --- a/templates.py +++ /dev/null @@ -1,18 +0,0 @@ -class TemplateQwen(): - # --- Prompt Template Configuration (ChatML) --- - IM_START = "<|im_start|>" - IM_END = "<|im_end|>" - ROLE_SYSTEM = "system" - ROLE_USER = "user" - ROLE_ASSISTANT = "assistant" - -class TemplateGemma(): - # --- Prompt Template Configuration (Gemma3) --- - IM_START = "" - IM_END = "" - ROLE_SYSTEM = "user" # Gemma has no system role - ROLE_USER = "user" - ROLE_ASSISTANT = "model" - -agent_template = TemplateQwen() -repl_template = TemplateGemma() \ No newline at end of file diff --git a/utils.py b/utils.py index 0b20b0c..fb40cf4 100644 --- a/utils.py +++ b/utils.py @@ -4,7 +4,6 @@ import time import logging import json import ast -import contextlib from rich.console import Console from rich.panel import Panel @@ -12,27 +11,28 @@ logger = logging.getLogger(__name__) console = Console() -def init_trace_file(debug, log_dir="logs"): +def init_trace_file(log_dir="logs"): if not os.path.exists(log_dir): os.makedirs(log_dir) timestamp = time.strftime("%Y%m%d-%H%M%S") filename = os.path.join(log_dir, f"trace_{timestamp}.json") - if debug: logger.debug(f"Trace logging initialized: {filename}") + logger.debug("Trace logging initialized: %s", filename) return filename -def save_agent_trace(filepath, messages, full_history=None): +def save_agent_trace(filepath, messages, step=0): try: data_to_save = { "timestamp": time.time(), + "step": step, "context_window": messages } - with open(filepath, 'w', encoding='utf-8') as f: - json.dump(data_to_save, f, indent=2, ensure_ascii=False) + with open(filepath, 'a', encoding='utf-8') as f: + f.write(json.dumps(data_to_save, ensure_ascii=False) + "\n") except Exception as e: - logger.error(f"Failed to save trace file: {e}") + logger.error("Failed to save trace file: %s", e) def _analyze_code_safety(code_str): @@ -101,10 +101,9 @@ def safeguard_and_repair(debug, client, messages, schema, original_code): if is_safe: return original_code + logger.warning("Safeguard triggered: %s (Line: %s)", error_msg, line_no) if debug: - logger.warning(f"Safeguard triggered: {error_msg} (Line: {line_no})") console.print(Panel(f"{error_msg}", title="Safeguard Interrupt", style="bold red")) - console.print(Panel( f"[italic]{error_msg}[/italic]", title="Unsafe Code", @@ -146,8 +145,7 @@ def safeguard_and_repair(debug, client, messages, schema, original_code): return full_fixed_code except json.JSONDecodeError: - if debug: logger.error("Snippet repair failed to parse. Falling back to full repair.") - pass + logger.error("Snippet repair failed to parse. Falling back to full repair.") repair_messages = messages + [ {"role": "assistant", "content": json.dumps({ @@ -171,11 +169,11 @@ def safeguard_and_repair(debug, client, messages, schema, original_code): except json.JSONDecodeError: return "" -def compress_history(debug, client, messages, keep_last_pairs=2): +def compress_history(client, messages, keep_last_pairs=2): keep_count = keep_last_pairs * 2 if len(messages) < (2 + 2 + keep_count): - if debug: logger.warning("History too short to compress, but context is full. Crashing safely.") + logger.warning("History too short to compress, but context is full. Crashing safely.") return messages to_compress = messages[2:-keep_count] @@ -194,7 +192,7 @@ def compress_history(debug, client, messages, keep_last_pairs=2): f"--- HISTORY START ---\n{history_text}\n--- HISTORY END ---" ) - if debug: logger.debug(f"Compressing {len(to_compress)} messages...") + logger.debug("Compressing %d messages...", len(to_compress)) summary_text = client.completion([{"role": "user", "content": summary_prompt}]) @@ -205,10 +203,10 @@ def compress_history(debug, client, messages, keep_last_pairs=2): new_messages = [messages[0], messages[1]] + [summary_message] + messages[-keep_count:] - if debug: logger.info(f"Compression complete. Reduced {len(messages)} msgs to {len(new_messages)}.") + logger.info("Compression complete. Reduced %d msgs to %d.", len(messages), len(new_messages)) return new_messages -def generate_final_report(debug, client, task_text, raw_answer): +def generate_final_report(client, task_text, raw_answer): system_prompt = ( "You are a professional report writer. " "Your goal is to convert the provided Raw Data into a clear, concise, " @@ -226,7 +224,7 @@ def generate_final_report(debug, client, task_text, raw_answer): Write the final response in natural language (Markdown). """ - if debug: logger.debug("Generating natural language report...") + logger.debug("Generating natural language report...") return client.completion([ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} @@ -237,5 +235,5 @@ def load_file(filepath): with open(filepath, 'r', encoding='utf-8') as f: return f.read() except FileNotFoundError: - logger.error(f"File not found: {filepath}") + logger.error("File not found: %s", filepath) sys.exit(1)