morpheus/erlm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor API llama.cpp native -> OpenAI compatible

Browse Source
This commit is contained in:
morpheus
2026-07-16 08:59:01 +01:00
parent 5d94e292b8
commit 07ee36da05
2 changed files with 163 additions and 266 deletions
Download Patch File Download Diff File Expand all files Collapse all files
edge_rlm.py
+115 -147
View File
@@ -15,14 +15,13 @@ from rich.json import JSON
from logging_config import setup_logging from logging_config import setup_logging
import utils import utils
import prompts as prompts import prompts as prompts
from templates import agent_template, repl_template
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
console = Console() console = Console()
# Configuration # Configuration
DEFAULT_AGENT_API = "http://localhost:8080" DEFAULT_AGENT_API = "http://localhost:8080/v1"
DEFAULT_REPL_API = "http://localhost:8090" DEFAULT_REPL_API = "http://localhost:8090/v1"
DEFAULT_CONTEXT_FILE = "context.txt" DEFAULT_CONTEXT_FILE = "context.txt"
DEFAULT_TASK_FILE = "task.txt" DEFAULT_TASK_FILE = "task.txt"
@@ -32,161 +31,160 @@ MAX_VIRTUAL_CONTEXT_RATIO = 0.85
class LlamaClient: class LlamaClient:
def __init__(self, base_url, name="LlamaClient"): def __init__(self, base_url, name="LlamaClient"):
self.base_url = base_url self.base_url = base_url.rstrip("/")
self.name = name self.name = name
self.n_ctx = self._get_context_size() self.model = None
self.n_ctx = 4096
self._get_model_info()
self.max_input_tokens = int(self.n_ctx * MAX_VIRTUAL_CONTEXT_RATIO) self.max_input_tokens = int(self.n_ctx * MAX_VIRTUAL_CONTEXT_RATIO)
self.color = self._determine_color() # Add this line self.color = self._determine_color()
if debug: logger.debug(f"Connected to {name} ({base_url}). Model Context: {self.n_ctx}. Max Input Safe Limit: {self.max_input_tokens}. Color: {self.color}") if debug: logger.debug(f"Connected to {name} ({base_url}). Model: {self.model}. Context: {self.n_ctx}. Max Input: {self.max_input_tokens}")
def _determine_color(self): def _determine_color(self):
if self.base_url == DEFAULT_AGENT_API: # Assuming args.agent_api is a string if "8080" in self.base_url:
return "dodger_blue1" return "dodger_blue1"
elif self.base_url == DEFAULT_REPL_API: elif "8090" in self.base_url:
return "dodger_blue3" return "dodger_blue3"
else: else:
return "cyan1" # Default color if base_url is unknown return "cyan1"
def _get_context_size(self): def _get_model_info(self):
try: try:
resp = requests.get(f"{self.base_url}/props") resp = requests.get(f"{self.base_url}/models")
resp.raise_for_status() resp.raise_for_status()
data = resp.json() data = resp.json()
model_data = data.get("data", [])
if 'n_ctx' in data: return data['n_ctx'] if model_data:
if 'default_n_ctx' in data: return data['default_n_ctx'] model = model_data[0]
if 'default_generation_settings' in data: self.model = model.get("id", "default")
settings = data['default_generation_settings'] meta = model.get("meta", {})
if 'n_ctx' in settings: return settings['n_ctx'] self.n_ctx = meta.get("n_ctx", 4096)
else:
return 4096 self.model = "default"
except Exception as e: except Exception as e:
logger.error(f"[{self.name}] Failed to get props: {e}. Defaulting to 4096.") logger.error(f"[{self.name}] Failed to get model info: {e}. Defaulting.")
return 4096 self.model = "default"
def tokenize(self, text): def count_tokens(self, messages):
try: try:
resp = requests.post(f"{self.base_url}/tokenize", json={"content": text}) resp = requests.post(
resp.raise_for_status() f"{self.base_url}/chat/completions/input_tokens",
return len(resp.json().get('tokens', [])) json={"model": self.model, "messages": messages},
except Exception: timeout=30.0,
return len(text) // 4 )
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.")
return sum(len(json.dumps(m)) // 4 + 4 for m in messages)
def completion(self, prompt, schema=None, temperature=0.1): def count_text_tokens(self, text):
return self.count_tokens([{"role": "user", "content": text}])
def completion(self, messages, schema=None, temperature=0.1):
payload = { payload = {
"prompt": prompt, "model": self.model,
"n_predict": -1, "messages": messages,
"temperature": temperature, "temperature": temperature,
"cache_prompt": True
} }
if schema: if schema:
payload["json_schema"] = schema payload["response_format"] = {
else: "type": "json_schema",
payload["stop"] = ["<|eot_id|>", "<|im_end|>", "Observation:", "User:"] "json_schema": {"name": "response", "schema": schema}
}
if debug: if debug:
last_content = messages[-1].get("content", "") if messages else ""
console.print(Panel( console.print(Panel(
prompt[500:], last_content[-500:] if len(last_content) > 500 else last_content,
title=f"Last 500 Characters of {self.name} Call", title=f"Last message to {self.name}",
title_align="left", title_align="left",
border_style=self.color border_style=self.color
)) ))
try: try:
resp = requests.post(f"{self.base_url}/completion", json=payload) resp = requests.post(f"{self.base_url}/chat/completions", json=payload)
resp.raise_for_status()
content = resp.json()["choices"][0]["message"]["content"].strip()
if debug: if debug:
console.print(Panel( console.print(Panel(
JSON.from_data(resp.json().get('content', '').strip()), JSON.from_data(content),
title=f"{self.name} Response", title=f"{self.name} Response",
title_align="left", title_align="left",
border_style=self.color border_style=self.color
)) ))
resp.raise_for_status() return content
return resp.json().get('content', '').strip()
except Exception as e: except Exception as e:
logger.error(f"[{self.name}] Error calling LLM: {e}") logger.error(f"[{self.name}] Error calling LLM: {e}")
return f"Error: {e}" return f"Error: {e}"
class AgentTools: class AgentTools:
def __init__(self, repl_client: LlamaClient, data_content: str): def __init__(self, repl_client: LlamaClient, data_content: str):
self.client = repl_client self.client = repl_client
self.RAW_CORPUS = data_content self.RAW_CORPUS = data_content
def llm_query(self, content_chunk, query): def llm_query(self, content_chunk, query):
if content_chunk == "RAW_CORPUS": 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, ...)`)." 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, ...)`)."
# --- OPTIMIZATION FIX: Heuristic check before network call ---
# Assume approx 4 chars per token. If it's wildly larger than context,
# fail fast to prevent network timeout on the /tokenize call.
estimated_tokens = len(content_chunk) // 3 estimated_tokens = len(content_chunk) // 3
if estimated_tokens > (self.client.n_ctx * 2): if estimated_tokens > (self.client.n_ctx * 2):
return f"ERROR: Chunk is massively too large (approx {estimated_tokens} tokens). Slice strictly." return f"ERROR: Chunk is massively too large (approx {estimated_tokens} tokens). Slice strictly."
# 2. Precise Safety check chunk_tokens = self.client.count_text_tokens(content_chunk)
chunk_tokens = self.client.tokenize(content_chunk) query_tokens = self.client.count_text_tokens(query)
query_tokens = self.client.tokenize(query)
total = chunk_tokens + query_tokens + 150 total = chunk_tokens + query_tokens + 150
if debug: logger.debug(f"[Sub-LLM] Processing Query with {total} tokens.") if debug: logger.debug(f"[Sub-LLM] Processing Query with {total} tokens.")
if total > self.client.n_ctx: if total > self.client.n_ctx:
msg = f"ERROR: Chunk too large ({chunk_tokens} tokens). Limit is {self.client.n_ctx}. Slice smaller." msg = f"ERROR: Chunk too large ({chunk_tokens} tokens). Limit is {self.client.n_ctx}. Slice smaller."
logger.warning(msg) logger.warning(msg)
return msg return msg
# 3. Strict Grounding Prompt
sub_messages = [ sub_messages = [
{"role": repl_template.ROLE_SYSTEM, "content": ( {"role": "system", "content": (
"You are a strict reading assistant. " "You are a strict reading assistant. "
"Answer the question based ONLY on the provided Context. " "Answer the question based ONLY on the provided Context. "
"Do not use outside training data. " "Do not use outside training data. "
f"If the answer is not in the text, say 'NULL'." "If the answer is not in the text, say 'NULL'."
)}, )},
{"role": repl_template.ROLE_USER, "content": f"Context:\n{content_chunk}\n\nQuestion: {query}"} {"role": "user", "content": f"Context:\n{content_chunk}\n\nQuestion: {query}"}
] ]
results = self.client.completion(utils.build_chat_prompt(sub_messages)) results = self.client.completion(sub_messages)
result_tokens = self.client.tokenize(results) result_tokens = self.client.count_text_tokens(results)
if debug: logger.debug(f"[Sub-LLM] Responded with {result_tokens} tokens.") if debug: logger.debug(f"[Sub-LLM] Responded with {result_tokens} tokens.")
return results return results
class AgentOutputBuffer: class AgentOutputBuffer:
def __init__(self, max_total_chars=20000, max_len_per_print=1009): def __init__(self, max_total_chars=20000, max_len_per_print=1009):
self._io = io.StringIO() self._io = io.StringIO()
self.max_total_chars = max_total_chars # Hard cap for infinite loop protection self.max_total_chars = max_total_chars
self.max_len_per_print = max_len_per_print # Soft cap for raw data dumping protection self.max_len_per_print = max_len_per_print
self.current_chars = 0 self.current_chars = 0
self.global_truncated = False self.global_truncated = False
def custom_print(self, *args, **kwargs): def custom_print(self, *args, **kwargs):
# 1. Capture the content of THIS specific print call
temp_io = io.StringIO() temp_io = io.StringIO()
print(*args, file=temp_io, **kwargs) print(*args, file=temp_io, **kwargs)
text = temp_io.getvalue() text = temp_io.getvalue()
# 2. Check PER-PRINT limit (The "Density" Check)
# This prevents printing raw corpus data, but allows short summaries to pass through
if len(text) > self.max_len_per_print: if len(text) > self.max_len_per_print:
# Slice the text
truncated_text = text[:self.max_len_per_print] truncated_text = text[:self.max_len_per_print]
# Create a localized warning that doesn't stop the whole stream
text = ( text = (
f"{truncated_text}\n" f"{truncated_text}\n"
f"... [LINE TRUNCATED: Output exceeded {self.max_len_per_print-9} chars. " f"... [LINE TRUNCATED: Output exceeded {self.max_len_per_print-9} chars. "
f"Use slicing or llm_query() to inspect data.] ...\n" f"Use slicing or llm_query() to inspect data.] ...\n"
) )
# 3. Check GLOBAL limit (The "Sanity" Check)
# This prevents infinite loops (while True: print('a')) from crashing memory
if self.current_chars + len(text) > self.max_total_chars: if self.current_chars + len(text) > self.max_total_chars:
remaining = self.max_total_chars - self.current_chars remaining = self.max_total_chars - self.current_chars
if remaining > 0: if remaining > 0:
self._io.write(text[:remaining]) self._io.write(text[:remaining])
if not self.global_truncated: if not self.global_truncated:
self._io.write(f"\n... [SYSTEM HALT: Total output limit ({self.max_total_chars}) reached] ...\n") self._io.write(f"\n... [SYSTEM HALT: Total output limit ({self.max_total_chars}) reached] ...\n")
self.global_truncated = True self.global_truncated = True
self.current_chars += len(text) self.current_chars += len(text)
else: else:
self._io.write(text) self._io.write(text)
self.current_chars += len(text) self.current_chars += len(text)
@@ -209,19 +207,15 @@ def run_agent(agent_client, repl_client, context_text, task_text):
"content": {"type": "string", "description": "Python code or Final Answer text."} "content": {"type": "string", "description": "Python code or Final Answer text."}
}, },
"required": ["thought", "action", "content"] "required": ["thought", "action", "content"]
} }
# 1. Instantiate the buffer
out_buffer = AgentOutputBuffer() out_buffer = AgentOutputBuffer()
trace_filepath = utils.init_trace_file(debug) trace_filepath = utils.init_trace_file(debug)
# 2. Add it to the environment
exec_env = { exec_env = {
"RAW_CORPUS": tools.RAW_CORPUS, "RAW_CORPUS": tools.RAW_CORPUS,
"llm_query": tools.llm_query, "llm_query": tools.llm_query,
# Standard Libs
"re": __import__("re"), "re": __import__("re"),
"math": __import__("math"), "math": __import__("math"),
"json": __import__("json"), "json": __import__("json"),
@@ -231,23 +225,22 @@ def run_agent(agent_client, repl_client, context_text, task_text):
"datetime": __import__("datetime"), "datetime": __import__("datetime"),
"difflib": __import__("difflib"), "difflib": __import__("difflib"),
"string": __import__("string"), "string": __import__("string"),
# Overrides "print": out_buffer.custom_print
"print": out_buffer.custom_print
} }
system_instruction = prompts.get_system_prompt() system_instruction = prompts.get_system_prompt()
messages = [ messages = [
{"role": agent_template.ROLE_SYSTEM, "content": system_instruction}, {"role": "system", "content": system_instruction},
{"role": agent_template.ROLE_USER, "content": f"USER TASK: {task_text}"} {"role": "user", "content": f"USER TASK: {task_text}"}
] ]
step = 0 step = 0
while step < MAX_REPL_STEPS: while step < MAX_REPL_STEPS:
step += 1 step += 1
if debug: logger.debug(f"Step {step} of {MAX_REPL_STEPS}") if debug: logger.debug(f"Step {step} of {MAX_REPL_STEPS}")
modules = [] modules = []
functions = [] functions = []
variables = [] variables = []
@@ -255,21 +248,18 @@ def run_agent(agent_client, repl_client, context_text, task_text):
for name, val in exec_env.items(): for name, val in exec_env.items():
if name.startswith("__"): continue if name.startswith("__"): continue
if name == "print": continue # Hide print, it's implied if name == "print": continue
if isinstance(val, types.ModuleType): if isinstance(val, types.ModuleType):
modules.append(name) modules.append(name)
elif callable(val): elif callable(val):
functions.append(name) functions.append(name)
else: else:
# For variables, provide a type and a short preview
type_name = type(val).__name__ type_name = type(val).__name__
s_val = str(val) s_val = str(val)
# Truncate long values for display (e.g. RAW_CORPUS)
snippet = (s_val[:ACTIVE_VAR_SNIPPET_LEN] + '...') if len(s_val) > ACTIVE_VAR_SNIPPET_LEN else s_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}") variables.append(f"{name} ({type_name}): {snippet}")
# 2. Create the status message
dynamic_state_msg = ( dynamic_state_msg = (
f"[SYSTEM STATE REMINDER]\n" f"[SYSTEM STATE REMINDER]\n"
f"Current Step: {step}/{MAX_REPL_STEPS}\n" f"Current Step: {step}/{MAX_REPL_STEPS}\n"
@@ -278,41 +268,31 @@ def run_agent(agent_client, repl_client, context_text, task_text):
f"Active Variables:\n" + ("\n".join([f" - {v}" for v in variables]) if variables else " (None)") + "\n---" f"Active Variables:\n" + ("\n".join([f" - {v}" for v in variables]) if variables else " (None)") + "\n---"
) )
# 3. Create a temporary message list for this specific inference
# We append the state to the very end so it has high 'recency' bias
inference_messages = messages.copy() inference_messages = messages.copy()
inference_messages.append({"role": agent_template.ROLE_USER, "content": dynamic_state_msg}) inference_messages.append({"role": "user", "content": dynamic_state_msg})
# 4. Build prompt using the INFERENCE messages (not the permanent history)
full_prompt = utils.build_chat_prompt(inference_messages)
usage = agent_client.tokenize(full_prompt) usage = agent_client.count_tokens(inference_messages)
if debug: logger.debug(f"Context Usage: {usage} / {agent_client.max_input_tokens}") if debug: logger.debug(f"Context Usage: {usage} / {agent_client.max_input_tokens}")
# Check context use and attempt compression
if usage > agent_client.max_input_tokens: if usage > agent_client.max_input_tokens:
if debug: logger.warning("Context limit exceeded. Triggering History Compression.") if debug: logger.warning("Context limit exceeded. Triggering History Compression.")
messages = utils.compress_history(debug, agent_client, messages, keep_last_pairs=2) messages = utils.compress_history(debug, agent_client, messages, keep_last_pairs=2)
# Re-check usage after compression new_usage = agent_client.count_tokens(inference_messages)
full_prompt = utils.build_chat_prompt(messages)
new_usage = agent_client.tokenize(full_prompt)
if debug: logger.debug(f"Context Usage after compression: {new_usage}") if debug: logger.debug(f"Context Usage after compression: {new_usage}")
# Panic mode: If it's STILL too big (unlikely), truncate the summary
if new_usage > agent_client.max_input_tokens: if new_usage > agent_client.max_input_tokens:
logger.error("Compression insufficient. Forcing hard truncation.") logger.error("Compression insufficient. Forcing hard truncation.")
messages.pop(2) messages.pop(2)
response_text = agent_client.completion(inference_messages, schema=agent_schema, temperature=0.5)
# Agent Completion
response_text = agent_client.completion(full_prompt, schema=agent_schema, temperature=0.5)
try: try:
response_json = json.loads(response_text) response_json = json.loads(response_text)
except json.JSONDecodeError: except json.JSONDecodeError:
logger.error("JSON Parse Error") logger.error("JSON Parse Error")
messages.append({"role": agent_template.ROLE_USER, "content": "System: Invalid JSON returned. Please retry."}) messages.append({"role": "user", "content": "System: Invalid JSON returned. Please retry."})
continue continue
thought = response_json.get("thought", "") thought = response_json.get("thought", "")
@@ -320,59 +300,47 @@ def run_agent(agent_client, repl_client, context_text, task_text):
content = response_json.get("content", "") content = response_json.get("content", "")
if action == "execute_python" and content: if action == "execute_python" and content:
# Run the safeguard. If the code is bad, 'content' gets replaced
content = utils.safeguard_and_repair(debug, agent_client, messages, agent_schema, content) content = utils.safeguard_and_repair(debug, agent_client, messages, agent_schema, content)
if debug: if debug:
console.print(Panel( console.print(Panel(
f"[italic]{thought}[/italic]", f"[italic]{thought}[/italic]",
title="🧠 Agent Thought", title="Agent Thought",
title_align="left", title_align="left",
border_style="magenta" border_style="magenta"
)) ))
messages.append({"role": agent_template.ROLE_ASSISTANT, "content": json.dumps(response_json, indent=2, ensure_ascii=False)}) messages.append({"role": "assistant", "content": json.dumps(response_json, indent=2, ensure_ascii=False)})
# 3. Execution
if action == "final_answer": if action == "final_answer":
# 1. Capture the raw result (keep this for logs/debugging)
if debug: logger.debug(f"Raw Agent Output: {content}") if debug: logger.debug(f"Raw Agent Output: {content}")
# Check if content looks like JSON/Structure, if so, summarize it.
# Even if it's already text, a quick polish pass ensures consistent tone.
final_report = utils.generate_final_report(debug, agent_client, task_text, content) final_report = utils.generate_final_report(debug, agent_client, task_text, content)
# 3. Print the pretty version
final_report_md = Markdown(final_report) final_report_md = Markdown(final_report)
print("\n\n") print("\n\n")
console.print(final_report_md) console.print(final_report_md)
print("\n") print("\n")
break break
elif action == "execute_python": elif action == "execute_python":
# Update the thought/log to reflect potential changes for the human observer
if debug and content != response_json.get("content"): if debug and content != response_json.get("content"):
console.print(Panel(content, title="Executing Code via Safeguard", title_align="left", border_style="cyan")) console.print(Panel(content, title="Executing Code via Safeguard", title_align="left", border_style="cyan"))
elif debug and content == response_json.get("content"): elif debug and content == response_json.get("content"):
console.print(Panel(content, title="Executing Code", title_align="left", border_style="yellow")) console.print(Panel(content, title="Executing Code", title_align="left", border_style="yellow"))
observation = "" observation = ""
try: try:
# 1. Clear any leftover junk from previous steps (safety) out_buffer.read_and_clear()
out_buffer.read_and_clear()
# 2. Execute. The Agent calls 'print', which goes to out_buffer
exec(content, exec_env) exec(content, exec_env)
# 3. Extract the text
observation = out_buffer.read_and_clear() observation = out_buffer.read_and_clear()
if not observation: if not observation:
observation = "Code executed successfully (no output)." observation = "Code executed successfully (no output)."
except Exception as e: except Exception as e:
observation = f"Python Error: {e}" observation = f"Python Error: {e}"
logger.error(f"Code Execution Error: {e}") logger.error(f"Code Execution Error: {e}")
if debug: if debug:
console.print(Panel( console.print(Panel(
f"{observation.strip()}", f"{observation.strip()}",
@@ -380,16 +348,16 @@ def run_agent(agent_client, repl_client, context_text, task_text):
title_align="left", title_align="left",
border_style="dark_green" border_style="dark_green"
)) ))
messages.append({"role": agent_template.ROLE_USER, "content": f"Observation:\n{observation}"}) messages.append({"role": "user", "content": f"Observation:\n{observation}"})
else: else:
messages.append({"role": agent_template.ROLE_USER, "content": f"System: Unknown action '{action}'."}) messages.append({"role": "user", "content": f"System: Unknown action '{action}'."})
utils.save_agent_trace(trace_filepath, messages) utils.save_agent_trace(trace_filepath, messages)
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser(description="""Edge Recursive Language Model 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.""") 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("--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("--task", default=DEFAULT_TASK_FILE, help="Path to task instruction file")
@@ -397,7 +365,7 @@ if __name__ == "__main__":
parser.add_argument("--agent_api", default=DEFAULT_AGENT_API, help="URL for the Main Agent LLM") 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("--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")
args = parser.parse_args() args = parser.parse_args()
debug = args.debug debug = args.debug
log_level=logging.DEBUG if debug else logging.INFO log_level=logging.DEBUG if debug else logging.INFO
@@ -406,9 +374,9 @@ if __name__ == "__main__":
if debug: logger.info("Starting EdgeRLM...") if debug: logger.info("Starting EdgeRLM...")
context_content = utils.load_file(args.context) context_content = utils.load_file(args.context)
if debug: logger.debug(f"Loaded Context: {len(context_content)} characters.") if debug: logger.debug(f"Loaded Context: {len(context_content)} characters.")
task_content = args.override_task if args.override_task else load_file(args.task) task_content = args.override_task if args.override_task else utils.load_file(args.task)
agent_client = LlamaClient(args.agent_api, "Agent") agent_client = LlamaClient(args.agent_api, "Agent")
repl_client = LlamaClient(args.repl_api, "REPL") repl_client = LlamaClient(args.repl_api, "REPL")
run_agent(agent_client, repl_client, context_content, task_content) run_agent(agent_client, repl_client, context_content, task_content)
utils.py
+48 -119
View File
@@ -1,4 +1,5 @@
import os import os
import sys
import time import time
import logging import logging
import json import json
@@ -7,85 +8,56 @@ import contextlib
from rich.console import Console from rich.console import Console
from rich.panel import Panel from rich.panel import Panel
from templates import agent_template
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
console = Console() console = Console()
def init_trace_file(debug, log_dir="logs"): def init_trace_file(debug, log_dir="logs"):
"""
Creates the log directory and returns a unique filepath
based on the current timestamp.
"""
if not os.path.exists(log_dir): if not os.path.exists(log_dir):
os.makedirs(log_dir) os.makedirs(log_dir)
timestamp = time.strftime("%Y%m%d-%H%M%S") timestamp = time.strftime("%Y%m%d-%H%M%S")
filename = os.path.join(log_dir, f"trace_{timestamp}.json") filename = os.path.join(log_dir, f"trace_{timestamp}.json")
if debug: logger.debug(f"Trace logging initialized: {filename}") if debug: logger.debug(f"Trace logging initialized: {filename}")
return filename return filename
def save_agent_trace(filepath, messages, full_history=None): def save_agent_trace(filepath, messages, full_history=None):
"""
Dumps the current state of the conversation to a JSON file.
Overwrites the file each step so the last write is always the complete history.
"""
try: try:
data_to_save = { data_to_save = {
"timestamp": time.time(), "timestamp": time.time(),
# If you are using history compression, 'messages' might get cut. "context_window": messages
# If you want the RAW full history, pass full_history.
# Otherwise, we log what the agent currently 'sees'.
"context_window": messages
} }
with open(filepath, 'w', encoding='utf-8') as f: with open(filepath, 'w', encoding='utf-8') as f:
json.dump(data_to_save, f, indent=2, ensure_ascii=False) json.dump(data_to_save, f, indent=2, ensure_ascii=False)
except Exception as e: except Exception as e:
logger.error(f"Failed to save trace file: {e}") logger.error(f"Failed to save trace file: {e}")
def build_chat_prompt(messages):
prompt = ""
for msg in messages:
role = msg.get("role")
content = msg.get("content")
prompt += f"{agent_template.IM_START}{role}\n{content}{agent_template.IM_END}\n"
prompt += f"{agent_template.IM_START}{agent_template.ROLE_ASSISTANT}" # Removed trailing newline
return prompt
def _analyze_code_safety(code_str): def _analyze_code_safety(code_str):
"""
Returns: (is_safe: bool, error_msg: str, line_number: int | None)
"""
try: try:
tree = ast.parse(code_str) tree = ast.parse(code_str)
except SyntaxError as e: except SyntaxError as e:
# e.lineno is the line where the parser failed
return False, f"SyntaxError: {e.msg}", e.lineno return False, f"SyntaxError: {e.msg}", e.lineno
tainted_vars = {"RAW_CORPUS"} tainted_vars = {"RAW_CORPUS"}
has_print = False has_print = False
for node in ast.walk(tree): for node in ast.walk(tree):
# 1. Track assignments
if isinstance(node, ast.Assign): if isinstance(node, ast.Assign):
if isinstance(node.value, ast.Name) and node.value.id in tainted_vars: if isinstance(node.value, ast.Name) and node.value.id in tainted_vars:
for target in node.targets: for target in node.targets:
if isinstance(target, ast.Name): if isinstance(target, ast.Name):
tainted_vars.add(target.id) tainted_vars.add(target.id)
# 2. Check Call nodes
if isinstance(node, ast.Call): if isinstance(node, ast.Call):
if isinstance(node.func, ast.Name) and node.func.id == 'print': if isinstance(node.func, ast.Name) and node.func.id == 'print':
has_print = True has_print = True
for arg in node.args: for arg in node.args:
if isinstance(arg, ast.Name) and arg.id in tainted_vars: if isinstance(arg, ast.Name) and arg.id in tainted_vars:
return False, f"Safety Violation: Printing '{arg.id}' (RAW_CORPUS). Use slicing.", node.lineno return False, f"Safety Violation: Printing '{arg.id}' (RAW_CORPUS). Use slicing.", node.lineno
# Check re.compile arguments
is_re_compile = False is_re_compile = False
if isinstance(node.func, ast.Attribute) and node.func.attr == 'compile': if isinstance(node.func, ast.Attribute) and node.func.attr == 'compile':
is_re_compile = True is_re_compile = True
@@ -95,69 +67,55 @@ def _analyze_code_safety(code_str):
if is_re_compile and len(node.args) > 2: if is_re_compile and len(node.args) > 2:
return False, "Library Usage Error: `re.compile` accepts max 2 args.", node.lineno return False, "Library Usage Error: `re.compile` accepts max 2 args.", node.lineno
# 3. Global Check (No specific line number)
if not has_print: if not has_print:
return False, "Observability Error: No `print()` statements found.", None return False, "Observability Error: No `print()` statements found.", None
return True, None, None return True, None, None
def _extract_context_block(code_str, target_lineno): def _extract_context_block(code_str, target_lineno):
"""
Extracts lines surrounding target_lineno bounded by empty lines.
Returns: (start_index, end_index, snippet_str)
"""
lines = code_str.split('\n') lines = code_str.split('\n')
# target_lineno is 1-based, list is 0-based
idx = target_lineno - 1 idx = target_lineno - 1
# Clamp index just in case
if idx < 0: idx = 0 if idx < 0: idx = 0
if idx >= len(lines): idx = len(lines) - 1 if idx >= len(lines): idx = len(lines) - 1
start_idx = idx start_idx = idx
end_idx = idx end_idx = idx
# Scan Up
while start_idx > 0: while start_idx > 0:
if lines[start_idx - 1].strip() == "": if lines[start_idx - 1].strip() == "":
break break
start_idx -= 1 start_idx -= 1
# Scan Down
while end_idx < len(lines) - 1: while end_idx < len(lines) - 1:
if lines[end_idx + 1].strip() == "": if lines[end_idx + 1].strip() == "":
break break
end_idx += 1 end_idx += 1
# Extract the block including the found boundaries (or lack thereof)
snippet_lines = lines[start_idx : end_idx + 1] snippet_lines = lines[start_idx : end_idx + 1]
return start_idx, end_idx, "\n".join(snippet_lines) return start_idx, end_idx, "\n".join(snippet_lines)
def safeguard_and_repair(debug, client, messages, schema, original_code): def safeguard_and_repair(debug, client, messages, schema, original_code):
is_safe, error_msg, line_no = _analyze_code_safety(original_code) is_safe, error_msg, line_no = _analyze_code_safety(original_code)
if is_safe: if is_safe:
return original_code return original_code
if debug: if debug:
logger.warning(f"Safeguard triggered: {error_msg} (Line: {line_no})") 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"{error_msg}", title="Safeguard Interrupt", style="bold red"))
console.print(Panel( console.print(Panel(
f"[italic]{thought}[/italic]", f"[italic]{error_msg}[/italic]",
title="Unsafe Code", title="Unsafe Code",
title_align="left", title_align="left",
border_style="hot_pink2" border_style="hot_pink2"
)) ))
# STRATEGY 1: SNIPPET REPAIR (Optimization)
# If we have a specific line number, we only send that block.
if line_no is not None: if line_no is not None:
start_idx, end_idx, snippet = _extract_context_block(original_code, line_no) start_idx, end_idx, snippet = _extract_context_block(original_code, line_no)
# We create a temporary "micro-agent" prompt just for fixing the snippet repair_messages = [
# We reuse the schema to ensure we get a clean content block back
repair_prompt = [
{"role": "system", "content": "You are a code repair assistant. Output only the fixed code snippet in the JSON content field."}, {"role": "system", "content": "You are a code repair assistant. Output only the fixed code snippet in the JSON content field."},
{"role": "user", "content": ( {"role": "user", "content": (
f"The following Python code snippet failed validation.\n" f"The following Python code snippet failed validation.\n"
@@ -169,43 +127,35 @@ def safeguard_and_repair(debug, client, messages, schema, original_code):
] ]
if debug: if debug:
console.print(Panel(f"{snippet}", title="Attempting Snippet Repair", style="light_goldenrod1")) console.print(Panel(f"{snippet}", title="Attempting Snippet Repair", style="light_goldenrod1"))
response_text = client.completion(repair_prompt, schema=schema, temperature=0.0) response_text = client.completion(repair_messages, schema=schema, temperature=0.0)
try: try:
response_json = json.loads(response_text) response_json = json.loads(response_text)
fixed_snippet = response_json.get("content", "") fixed_snippet = response_json.get("content", "")
if debug: if debug:
console.print(Panel(f"{fixed_snippet}", title="Repaired Snippet", style="yellow1")) console.print(Panel(f"{fixed_snippet}", title="Repaired Snippet", style="yellow1"))
# Stitch the code back together
all_lines = original_code.split('\n') all_lines = original_code.split('\n')
# We replace the range we extracted with the new snippet
# Note: fixed_snippet might have different line count, that's fine.
pre_block = all_lines[:start_idx] pre_block = all_lines[:start_idx]
post_block = all_lines[end_idx + 1:] post_block = all_lines[end_idx + 1:]
# Reassemble
full_fixed_code = "\n".join(pre_block + [fixed_snippet] + post_block) full_fixed_code = "\n".join(pre_block + [fixed_snippet] + post_block)
return full_fixed_code return full_fixed_code
except json.JSONDecodeError: except json.JSONDecodeError:
# If the snippet repair fails to parse, fall through to full repair
if debug: logger.error("Snippet repair failed to parse. Falling back to full repair.") if debug: logger.error("Snippet repair failed to parse. Falling back to full repair.")
pass pass
# STRATEGY 2: FULL REPAIR (Fallback)
# Used for global errors (missing prints) or if snippet repair crashed
repair_messages = messages + [ repair_messages = messages + [
{"role": agent_template.ROLE_ASSISTANT, "content": json.dumps({ {"role": "assistant", "content": json.dumps({
"thought": "Drafting code...", "thought": "Drafting code...",
"action": "execute_python", "action": "execute_python",
"content": original_code "content": original_code
})}, })},
{"role": agent_template.ROLE_USER, "content": ( {"role": "user", "content": (
f"SYSTEM INTERRUPT: Your code failed pre-flight safety checks.\n" f"SYSTEM INTERRUPT: Your code failed pre-flight safety checks.\n"
f"Error: {error_msg}\n\n" f"Error: {error_msg}\n\n"
f"Generate the JSON response again with CORRECTED Python code.\n" f"Generate the JSON response again with CORRECTED Python code.\n"
@@ -213,7 +163,7 @@ def safeguard_and_repair(debug, client, messages, schema, original_code):
)} )}
] ]
response_text = client.completion(build_chat_prompt(repair_messages), schema=schema, temperature=0.0) response_text = client.completion(repair_messages, schema=schema, temperature=0.0)
try: try:
response_json = json.loads(response_text) response_json = json.loads(response_text)
@@ -222,31 +172,20 @@ def safeguard_and_repair(debug, client, messages, schema, original_code):
return "" return ""
def compress_history(debug, client, messages, keep_last_pairs=2): def compress_history(debug, client, messages, keep_last_pairs=2):
"""
Compresses the middle of the conversation history.
Preserves: System Prompt (0), User Task (1), and the last N pairs of interaction.
"""
# Calculate how many messages to keep at the end (pairs * 2)
keep_count = keep_last_pairs * 2 keep_count = keep_last_pairs * 2
# Check if we actually have enough history to compress
# We need: System + Task + (At least 2 messages to compress) + Keep_Count
if len(messages) < (2 + 2 + keep_count): if len(messages) < (2 + 2 + keep_count):
if debug: logger.warning("History too short to compress, but context is full. Crashing safely.") if debug: logger.warning("History too short to compress, but context is full. Crashing safely.")
return messages # Nothing we can do, let it fail or truncate manually return messages
# Define the slice to compress
# Start at 2 (after Task), End at -keep_count
to_compress = messages[2:-keep_count] to_compress = messages[2:-keep_count]
# 1. format the text for the summarizer
history_text = "" history_text = ""
for msg in to_compress: for msg in to_compress:
role = msg['role'].upper() role = msg['role'].upper()
content = msg['content'] content = msg['content']
history_text += f"[{role}]: {content}\n" history_text += f"[{role}]: {content}\n"
# 2. Build the summarization prompt
summary_prompt = ( summary_prompt = (
"You are a technical documentation assistant. " "You are a technical documentation assistant. "
"Summarize the following interaction history between an AI Agent and a System. " "Summarize the following interaction history between an AI Agent and a System. "
@@ -254,39 +193,29 @@ def compress_history(debug, client, messages, keep_last_pairs=2):
"Be concise. Do not chat.\n\n" "Be concise. Do not chat.\n\n"
f"--- HISTORY START ---\n{history_text}\n--- HISTORY END ---" f"--- HISTORY START ---\n{history_text}\n--- HISTORY END ---"
) )
if debug: logger.debug(f"Compressing {len(to_compress)} messages...") if debug: logger.debug(f"Compressing {len(to_compress)} messages...")
# 3. Call the LLM (We use the Agent Client for high-quality summaries) summary_text = client.completion([{"role": "user", "content": summary_prompt}])
# We use a simple generation call here.
summary_text = client.completion(
build_chat_prompt([{"role": "user", "content": summary_prompt}])
)
# 4. Create the new compressed message
summary_message = { summary_message = {
"role": "user", "role": "user",
"content": f"[SYSTEM SUMMARY OF PREVIOUS ACTIONS]\n{summary_text}" "content": f"[SYSTEM SUMMARY OF PREVIOUS ACTIONS]\n{summary_text}"
} }
# 5. Reconstruct the list
new_messages = [messages[0], messages[1]] + [summary_message] + messages[-keep_count:] 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)}.") if debug: logger.info(f"Compression complete. Reduced {len(messages)} msgs to {len(new_messages)}.")
return new_messages return new_messages
def generate_final_report(debug, client, task_text, raw_answer): def generate_final_report(debug, client, task_text, raw_answer):
"""
Converts the Agent's raw (likely structured/technical) answer into
a natural language response for the user.
"""
system_prompt = ( system_prompt = (
"You are a professional report writer. " "You are a professional report writer. "
"Your goal is to convert the provided Raw Data into a clear, concise, " "Your goal is to convert the provided Raw Data into a clear, concise, "
"and well-formatted response to the User's original request. " "and well-formatted response to the User's original request. "
"Do not add new facts. Just format and explain the existing data." "Do not add new facts. Just format and explain the existing data."
) )
user_prompt = f"""### USER REQUEST user_prompt = f"""### USER REQUEST
{task_text} {task_text}
@@ -296,12 +225,12 @@ def generate_final_report(debug, client, task_text, raw_answer):
### INSTRUCTION ### INSTRUCTION
Write the final response in natural language (Markdown). Write the final response in natural language (Markdown).
""" """
if debug: logger.debug("Generating natural language report...") if debug: logger.debug("Generating natural language report...")
return client.completion(build_chat_prompt([ return client.completion([
{"role": agent_template.ROLE_SYSTEM, "content": system_prompt}, {"role": "system", "content": system_prompt},
{"role": agent_template.ROLE_USER, "content": user_prompt} {"role": "user", "content": user_prompt}
])) ])
def load_file(filepath): def load_file(filepath):
try: try:
@@ -309,4 +238,4 @@ def load_file(filepath):
return f.read() return f.read()
except FileNotFoundError: except FileNotFoundError:
logger.error(f"File not found: {filepath}") logger.error(f"File not found: {filepath}")
sys.exit(1) sys.exit(1)