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

@@ -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)

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:
+    if level == logging.DEBUG:
         logging.getLogger(__name__).debug("[dim]Debug mode active.[/dim]")

requirements.txt

@@ -0,0 +1,2 @@
+requests>=2.28.0
+rich>=13.0.0

templates.py

@@ -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 = "<start_of_turn>"
-    IM_END = "<end_of_turn>"
-    ROLE_SYSTEM = "user" # Gemma has no system role
-    ROLE_USER = "user"
-    ROLE_ASSISTANT = "model"
-
-agent_template = TemplateQwen()
-repl_template = TemplateGemma()

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)