erlm/utils.py

import os
import sys
import time
import logging
import json
import ast
from rich.console import Console
from rich.panel import Panel

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

WHITELISTED_IMPORTS = {
    "json", "math", "datetime", "re", "collections",
    "statistics", "random", "difflib", "string",
}

DANGEROUS_BUILTINS = {"eval", "exec", "compile"}


def init_trace_file(log_dir="logs"):
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)

    timestamp = time.strftime("%Y%m%d-%H%M%S")
    filename = os.path.join(log_dir, f"trace_{timestamp}.json")
    logger.debug("Trace logging initialized: %s", filename)
    return filename

def save_agent_trace(filepath, new_messages, step=0, state=None, plan_step=0, total_messages=0):
    try:
        data_to_save = {
            "timestamp": time.time(),
            "step": step,
            "state": state,
            "plan_step": plan_step,
            "total_messages": total_messages,
            "new_messages": new_messages,
        }

        with open(filepath, 'a', encoding='utf-8') as f:
            f.write(json.dumps(data_to_save, ensure_ascii=False) + "\n")

    except Exception as e:
        logger.error("Failed to save trace file: %s", e)


def _check_dunder_access(node):
    if isinstance(node, ast.Attribute) and node.attr.startswith("__") and node.attr.endswith("__"):
        return True
    if isinstance(node, ast.Subscript):
        if isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, str):
            if node.slice.value.startswith("__") and node.slice.value.endswith("__"):
                return True
    return False

def _analyze_code_safety(code_str):
    try:
        tree = ast.parse(code_str)
    except SyntaxError as e:
        return False, f"SyntaxError: {e.msg}", e.lineno

    tainted_vars = {"RAW_CORPUS"}
    has_print = False

    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                if alias.name not in WHITELISTED_IMPORTS:
                    return False, f"Import Violation: '{alias.name}' is not whitelisted. Allowed: {', '.join(sorted(WHITELISTED_IMPORTS))}.", node.lineno

        if isinstance(node, ast.ImportFrom):
            if node.module not in WHITELISTED_IMPORTS:
                return False, f"Import Violation: '{node.module}' is not whitelisted. Allowed: {', '.join(sorted(WHITELISTED_IMPORTS))}.", node.lineno

        if _check_dunder_access(node):
            name = None
            if isinstance(node, ast.Attribute):
                name = node.attr
            elif isinstance(node, ast.Subscript):
                name = getattr(node.slice, 'value', getattr(node.slice, 's', None))
            return False, f"Dunder Access Violation: Accessing '{name}' is forbidden.", getattr(node, 'lineno', None)

        if isinstance(node, ast.Assign):
            if isinstance(node.value, ast.Name) and node.value.id in tainted_vars:
                for target in node.targets:
                    if isinstance(target, ast.Name):
                        tainted_vars.add(target.id)

        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name):
                if node.func.id in DANGEROUS_BUILTINS:
                    return False, f"Builtin Violation: '{node.func.id}()' is forbidden.", node.lineno

            if isinstance(node.func, ast.Name) and node.func.id == 'print':
                has_print = True
                for arg in node.args:
                    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

            is_re_compile = False
            if isinstance(node.func, ast.Attribute) and node.func.attr == 'compile':
                is_re_compile = True
            elif isinstance(node.func, ast.Name) and node.func.id == 'compile':
                is_re_compile = True

            if is_re_compile and len(node.args) > 2:
                return False, "Library Usage Error: `re.compile` accepts max 2 args.", node.lineno

    if not has_print:
        return False, "Observability Error: No `print()` statements found.", None

    return True, None, None

def _extract_context_block(code_str, target_lineno):
    lines = code_str.split('\n')
    idx = target_lineno - 1

    if idx < 0: idx = 0
    if idx >= len(lines): idx = len(lines) - 1

    start_idx = idx
    end_idx = idx

    while start_idx > 0:
        if lines[start_idx - 1].strip() == "":
            break
        start_idx -= 1

    while end_idx < len(lines) - 1:
        if lines[end_idx + 1].strip() == "":
            break
        end_idx += 1

    snippet_lines = lines[start_idx : end_idx + 1]
    return start_idx, end_idx, "\n".join(snippet_lines)

def safeguard_and_repair(debug, client, messages, schema, original_code):
    is_safe, error_msg, line_no = _analyze_code_safety(original_code)

    if is_safe:
        return original_code

    logger.warning("Safeguard triggered: %s (Line: %s)", error_msg, line_no)
    if debug:
        console.print(Panel(f"{error_msg}", title="Safeguard Interrupt", style="bold red"))
        console.print(Panel(
                f"[italic]{error_msg}[/italic]",
                title="Unsafe Code",
                title_align="left",
                border_style="hot_pink2"
            ))

    if line_no is not None:
        start_idx, end_idx, snippet = _extract_context_block(original_code, line_no)

        repair_messages = [
            {"role": "system", "content": "You are a code repair assistant. Output only the fixed code snippet in the JSON content field."},
            {"role": "user", "content": (
                f"The following Python code snippet failed validation.\n"
                f"Error: {error_msg} (occurred around line {line_no})\n\n"
                f"```python\n{snippet}\n```\n\n"
                f"Return the JSON with the fixed snippet. "
                f"Maintain original indentation. Add a comment (# FIXED) to changed lines."
            )}
        ]
        if debug:
            console.print(Panel(f"{snippet}", title="Attempting Snippet Repair", style="light_goldenrod1"))

        response_text = client.completion(repair_messages, schema=schema, temperature=0.0)

        try:
            response_json = json.loads(response_text)
            fixed_snippet = response_json.get("content", "")

            if debug:
                console.print(Panel(f"{fixed_snippet}", title="Repaired Snippet", style="yellow1"))

            all_lines = original_code.split('\n')
            pre_block = all_lines[:start_idx]
            post_block = all_lines[end_idx + 1:]

            full_fixed_code = "\n".join(pre_block + [fixed_snippet] + post_block)

            return full_fixed_code

        except json.JSONDecodeError:
            logger.error("Snippet repair failed to parse. Falling back to full repair.")

    repair_messages = messages + [
        {"role": "assistant", "content": json.dumps({
            "thought": "Drafting code...",
            "action": "execute_python",
            "content": original_code
        })},
        {"role": "user", "content": (
            f"SYSTEM INTERRUPT: Your code failed pre-flight safety checks.\n"
            f"Error: {error_msg}\n\n"
            f"Generate the JSON response again with CORRECTED Python code.\n"
            f"IMPORTANT: You must add a comment (# FIXED: ...) to the corrected line."
        )}
    ]

    response_text = client.completion(repair_messages, schema=schema, temperature=0.0)

    try:
        response_json = json.loads(response_text)
        return response_json.get("content", "")
    except json.JSONDecodeError:
        return ""

def compress_history(client, messages, keep_last_pairs=2):
    keep_count = keep_last_pairs * 2

    if len(messages) < (2 + 2 + keep_count):
        logger.warning("History too short to compress, but context is full. Crashing safely.")
        return messages

    to_compress = messages[2:-keep_count]

    history_text = ""
    for msg in to_compress:
        role = msg['role'].upper()
        content = msg['content']
        history_text += f"[{role}]: {content}\n"

    summary_prompt = (
        "You are a technical documentation assistant. "
        "Summarize the following interaction history between an AI Agent and a System. "
        "Focus on: 1. Code executed, 2. Errors encountered, 3. Specific data/variables discovered. "
        "Be concise. Do not chat.\n\n"
        f"--- HISTORY START ---\n{history_text}\n--- HISTORY END ---"
    )

    logger.debug("Compressing %d messages...", len(to_compress))

    summary_text = client.completion([{"role": "user", "content": summary_prompt}])

    summary_message = {
        "role": "user",
        "content": f"[SYSTEM SUMMARY OF PREVIOUS ACTIONS]\n{summary_text}"
    }

    new_messages = [messages[0], messages[1]] + [summary_message] + messages[-keep_count:]

    logger.info("Compression complete. Reduced %d msgs to %d.", len(messages), len(new_messages))
    return new_messages

def generate_final_report(client, task_text, raw_answer):
    system_prompt = (
        "You are a professional report writer. "
        "Your goal is to convert the provided Raw Data into a clear, concise, "
        "and well-formatted response to the User's original request. "
        "Do not add new facts. Just format and explain the existing data."
    )

    user_prompt = f"""### USER REQUEST
{task_text}

### RAW DATA COLLECTED
{raw_answer}

### INSTRUCTION
Write the final response in natural language (Markdown).
    """

    logger.debug("Generating natural language report...")
    return client.completion([
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ])

def validate_plan_object(obj):
    if not isinstance(obj, dict):
        return False, "PlanObject must be a JSON object."
    for key in ("goal", "strategy", "success_criteria"):
        if key not in obj:
            return False, f"PlanObject missing required key '{key}'."
    if not isinstance(obj["goal"], str) or not obj["goal"].strip():
        return False, "PlanObject 'goal' must be a non-empty string."
    if not isinstance(obj["strategy"], list) or len(obj["strategy"]) == 0:
        return False, "PlanObject 'strategy' must be a non-empty list of strings."
    if not isinstance(obj["success_criteria"], list) or len(obj["success_criteria"]) == 0:
        return False, "PlanObject 'success_criteria' must be a non-empty list of strings."
    if len(obj["strategy"]) != len(obj["success_criteria"]):
        return False, "PlanObject 'strategy' and 'success_criteria' must have the same length."
    for i, s in enumerate(obj["strategy"]):
        if not isinstance(s, str) or not s.strip():
            return False, f"PlanObject 'strategy[{i}]' must be a non-empty string."
    for i, s in enumerate(obj["success_criteria"]):
        if not isinstance(s, str) or not s.strip():
            return False, f"PlanObject 'success_criteria[{i}]' must be a non-empty string."
    return True, None


def validate_pivot_object(obj):
    if not isinstance(obj, dict):
        return False, "PivotObject must be a JSON object."
    if "diagnosis" not in obj:
        return False, "PivotObject missing required key 'diagnosis'."
    if not isinstance(obj["diagnosis"], str) or not obj["diagnosis"].strip():
        return False, "PivotObject 'diagnosis' must be a non-empty string."
    if "new_strategy" in obj or "new_success_criteria" in obj:
        if "new_strategy" not in obj or "new_success_criteria" not in obj:
            return False, "PivotObject: 'new_strategy' and 'new_success_criteria' must both be present or both omitted."
        if not isinstance(obj["new_strategy"], list) or len(obj["new_strategy"]) == 0:
            return False, "PivotObject 'new_strategy' must be a non-empty list of strings."
        if not isinstance(obj["new_success_criteria"], list) or len(obj["new_success_criteria"]) == 0:
            return False, "PivotObject 'new_success_criteria' must be a non-empty list of strings."
        if len(obj["new_strategy"]) != len(obj["new_success_criteria"]):
            return False, "PivotObject 'new_strategy' and 'new_success_criteria' must have the same length."
    return True, None


def load_file(filepath):
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            return f.read()
    except FileNotFoundError:
        logger.error("File not found: %s", filepath)
        sys.exit(1)