409 lines
16 KiB
Python
409 lines
16 KiB
Python
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)
|