erlm/edge_rlm.py

import os
import time
import requests
import json
import argparse
import io
import logging
import re
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


def _load_dotenv(path=".env"):
    try:
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#") or "=" not in line:
                    continue
                key, _, val = line.partition("=")
                key, val = key.strip(), val.strip()
                if (val.startswith('"') and val.endswith('"')) or (val.startswith("'") and val.endswith("'")):
                    val = val[1:-1]
                os.environ.setdefault(key, val)
    except FileNotFoundError:
        pass

_load_dotenv()

logger = logging.getLogger(__name__)
console = Console()

# Configuration (from .env with fallback defaults)
AGENT_API = os.getenv("AGENT_API", "http://localhost:8080/v1")
REPL_API = os.getenv("REPL_API", "http://localhost:8090/v1")

CONTEXT_FILE = os.getenv("CONTEXT_FILE", "context.txt")
TASK_FILE = os.getenv("TASK_FILE", "task.txt")
MAX_REPL_STEPS = int(os.getenv("MAX_REPL_STEPS", "20"))
MAX_VIRTUAL_CONTEXT_RATIO = float(os.getenv("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 _make_exec_env(tools, out_buffer):
    return {
        "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,
    }


def _compress_if_needed(agent_client, messages, inference_messages):
    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)
        rebuilt = messages + inference_messages[-1:]
        new_usage = agent_client.count_tokens(rebuilt)
        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)
    return messages


def run_agent(agent_client, repl_client, context_text, task_text):
    tools = AgentTools(repl_client, context_text)

    planning_schema = {
        "type": "object",
        "properties": {
            "goal": {"type": "string", "description": "The high-level objective."},
            "strategy": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Discrete, programmatic actions.",
            },
            "success_criteria": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Success criteria parallel to strategy.",
            },
        },
        "required": ["goal", "strategy", "success_criteria"],
    }

    executing_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."},
            "step_completed": {
                "type": "boolean",
                "description": "Set to true when the current strategy step is fully complete and you are ready to advance.",
            },
        },
        "required": ["thought", "action", "content"],
    }

    trajectory_schema = {
        "type": "object",
        "properties": {
            "progress": {"type": "string", "enum": ["YES", "NO"]},
            "justification": {"type": "string", "description": "Brief explanation of the assessment."},
        },
        "required": ["progress", "justification"],
    }

    pivot_schema = {
        "type": "object",
        "properties": {
            "diagnosis": {"type": "string", "description": "Why progress stalled."},
            "new_strategy": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Revised strategy (omit for Soft Pivot).",
            },
            "new_success_criteria": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Revised success criteria (omit for Soft Pivot).",
            },
        },
        "required": ["diagnosis"],
    }

    out_buffer = AgentOutputBuffer()
    trace_filepath = utils.init_trace_file()
    exec_env = _make_exec_env(tools, out_buffer)

    system_instruction = prompts.get_system_prompt(max_repl_steps=MAX_REPL_STEPS)

    messages = [
        {"role": "system", "content": system_instruction},
        {"role": "user", "content": f"USER TASK: {task_text}"},
    ]

    state = "PLANNING"
    plan = None
    turns_since_checkpoint = 0
    total_turns = 0
    last_diagnosis = ""
    last_trace_msg_count = 0

    while state != "FINISHED":
        logger.debug("State: %s | Turns since checkpoint: %d / %d", state, turns_since_checkpoint, MAX_REPL_STEPS)

        if state == "PLANNING":
            planning_prompt = prompts.get_planning_prompt(task_text)

            planning_messages = [
                {"role": "system", "content": system_instruction},
                {"role": "user", "content": planning_prompt},
            ]

            planning_messages = _compress_if_needed(agent_client, planning_messages, planning_messages)

            response_text = agent_client.completion(
                planning_messages, schema=planning_schema, temperature=0.3
            )

            try:
                plan = json.loads(response_text)
            except json.JSONDecodeError:
                logger.error("Planning: JSON parse error. Retrying.")
                messages.append({"role": "user", "content": "System: Invalid PlanObject JSON. Please retry."})
                continue

            valid, err = utils.validate_plan_object(plan)
            if not valid:
                logger.error("Planning: Invalid PlanObject: %s", err)
                messages.append({"role": "user", "content": f"System: Invalid PlanObject: {err}. Please retry."})
                continue

            plan["current_step_index"] = 0
            plan["state"] = "executing"
            turns_since_checkpoint = 0

            messages.append({"role": "assistant", "content": json.dumps(plan, indent=2, ensure_ascii=False)})
            messages.append({"role": "user", "content": f"System: Plan accepted. Starting execution of step 1/{len(plan['strategy'])}: {plan['strategy'][0]}"})

            logger.info("Plan accepted: goal=%s, steps=%d", plan["goal"], len(plan["strategy"]))
            state = "EXECUTING"

        elif state == "EXECUTING":
            dynamic_state_msg = prompts.get_execution_context(
                plan, exec_env, turns_since_checkpoint, MAX_REPL_STEPS
            )

            inference_messages = messages.copy()
            inference_messages.append({"role": "user", "content": dynamic_state_msg})

            messages = _compress_if_needed(agent_client, messages, inference_messages)

            inference_messages = messages.copy()
            inference_messages.append({"role": "user", "content": dynamic_state_msg})

            response_text = agent_client.completion(
                inference_messages, schema=executing_schema, temperature=0.5
            )

            try:
                response_json = json.loads(response_text)
            except json.JSONDecodeError:
                logger.error("EXECUTING: 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", "")
            step_completed = response_json.get("step_completed", False)

            if action == "execute_python" and content:
                content = utils.safeguard_and_repair(
                    agent_client.debug, agent_client, messages, executing_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")
                state = "FINISHED"

            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}"})

                total_turns += 1
                turns_since_checkpoint += 1

                if step_completed:
                    plan["current_step_index"] += 1
                    if plan["current_step_index"] >= len(plan["strategy"]):
                        plan["current_step_index"] = len(plan["strategy"]) - 1
                        messages.append({
                            "role": "user",
                            "content": "System: All plan steps completed. Provide your final answer."
                        })
                    else:
                        next_idx = plan["current_step_index"]
                        messages.append({
                            "role": "user",
                            "content": (
                                f"System: Advancing to step {next_idx + 1}/{len(plan['strategy'])}: "
                                f"{plan['strategy'][next_idx]}. "
                                f"Success criteria: {plan['success_criteria'][next_idx]}"
                            )
                        })
                    turns_since_checkpoint = 0

                if turns_since_checkpoint >= MAX_REPL_STEPS:
                    state = "TRAJECTORY_CHECK"
                    messages.append({
                        "role": "user",
                        "content": "System: Checkpoint reached. Pausing for trajectory review."
                    })

            else:
                messages.append({"role": "user", "content": f"System: Unknown action '{action}'."})

            trace_delta = messages[last_trace_msg_count:] if last_trace_msg_count <= len(messages) else messages
            utils.save_agent_trace(
                trace_filepath, trace_delta, step=total_turns,
                state=state, plan_step=plan["current_step_index"],
                total_messages=len(messages),
            )
            last_trace_msg_count = len(messages)

        elif state == "TRAJECTORY_CHECK":
            step_idx = plan["current_step_index"]
            criteria = plan["success_criteria"][step_idx]
            check_prompt = prompts.get_trajectory_check_prompt(criteria)

            check_messages = [
                {"role": "system", "content": system_instruction},
                {"role": "user", "content": check_prompt},
            ]
            if len(messages) >= 4:
                check_messages += messages[-4:]
            else:
                check_messages += messages

            response_text = agent_client.completion(
                check_messages, schema=trajectory_schema, temperature=0.2
            )

            try:
                check_json = json.loads(response_text)
            except json.JSONDecodeError:
                logger.error("TRAJECTORY_CHECK: JSON parse error. Assuming NO.")
                check_json = {"progress": "NO", "justification": "JSON parse failure."}

            progress = check_json.get("progress", "NO")
            justification = check_json.get("justification", "")

            messages.append({
                "role": "assistant",
                "content": json.dumps(check_json, indent=2, ensure_ascii=False),
            })

            if progress == "YES":
                logger.info("Trajectory check PASSED: %s", justification[:100])
                turns_since_checkpoint = 0
                state = "EXECUTING"
                messages.append({
                    "role": "user",
                    "content": (
                        f"System: Trajectory check passed. Resuming execution. "
                        f"Turns since checkpoint reset to 0."
                    )
                })
            else:
                logger.warning("Trajectory check FAILED: %s", justification[:100])
                last_diagnosis = justification
                state = "PIVOTING"
                messages.append({
                    "role": "user",
                    "content": "System: Trajectory check failed. Transitioning to pivot."
                })

            trace_delta = messages[last_trace_msg_count:] if last_trace_msg_count <= len(messages) else messages
            utils.save_agent_trace(
                trace_filepath, trace_delta, step=total_turns,
                state=state, plan_step=plan["current_step_index"],
                total_messages=len(messages),
            )
            last_trace_msg_count = len(messages)

        elif state == "PIVOTING":
            pivot_prompt = prompts.get_pivot_prompt(last_diagnosis, task_text)

            pivot_messages = [
                {"role": "system", "content": system_instruction},
                {"role": "user", "content": pivot_prompt},
            ]
            if len(messages) >= 4:
                pivot_messages += messages[-4:]
            else:
                pivot_messages += messages

            response_text = agent_client.completion(
                pivot_messages, schema=pivot_schema, temperature=0.4
            )

            try:
                pivot_json = json.loads(response_text)
            except json.JSONDecodeError:
                logger.error("PIVOTING: JSON parse error. Forcing Soft Pivot.")
                pivot_json = {"diagnosis": "JSON parse error. Retrying current step."}

            valid, err = utils.validate_pivot_object(pivot_json)
            if not valid:
                logger.error("PIVOTING: Invalid PivotObject: %s. Forcing Soft Pivot.", err)
                pivot_json = {"diagnosis": pivot_json.get("diagnosis", "Validation error: " + err)}

            has_new_strategy = "new_strategy" in pivot_json and "new_success_criteria" in pivot_json

            messages.append({
                "role": "assistant",
                "content": json.dumps(pivot_json, indent=2, ensure_ascii=False),
            })

            if has_new_strategy:
                logger.info("Hard Pivot: new strategy with %d steps.", len(pivot_json["new_strategy"]))
                plan = {
                    "goal": plan["goal"],
                    "strategy": pivot_json["new_strategy"],
                    "success_criteria": pivot_json["new_success_criteria"],
                    "current_step_index": 0,
                    "state": "executing",
                }
                exec_env = _make_exec_env(tools, out_buffer)
                messages.append({
                    "role": "user",
                    "content": (
                        f"System: Hard Pivot applied. Kernel reset. "
                        f"New plan has {len(plan['strategy'])} steps. "
                        f"Starting step 1: {plan['strategy'][0]}"
                    )
                })
            else:
                logger.info("Soft Pivot: retrying current step with fresh budget.")
                messages.append({
                    "role": "user",
                    "content": (
                        f"System: Soft Pivot applied. Kernel state preserved. "
                        f"Resuming step {plan['current_step_index'] + 1}/{len(plan['strategy'])}: "
                        f"{plan['strategy'][plan['current_step_index']]}"
                    )
                })

            turns_since_checkpoint = 0
            state = "EXECUTING"

            trace_delta = messages[last_trace_msg_count:] if last_trace_msg_count <= len(messages) else messages
            utils.save_agent_trace(
                trace_filepath, trace_delta, step=total_turns,
                state=state, plan_step=plan["current_step_index"],
                total_messages=len(messages),
            )
            last_trace_msg_count = len(messages)

    logger.info("Agent finished. Total execution turns: %d", total_turns)

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=CONTEXT_FILE, help="Path to text file to process")
    parser.add_argument("--task", 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=AGENT_API, help="URL for the Main Agent LLM")
    parser.add_argument("--repl_api", 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)