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)