import os import sys import time import logging import json import ast import contextlib from rich.console import Console from rich.panel import Panel logger = logging.getLogger(__name__) console = Console() def init_trace_file(debug, 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") if debug: logger.debug(f"Trace logging initialized: {filename}") return filename def save_agent_trace(filepath, messages, full_history=None): try: data_to_save = { "timestamp": time.time(), "context_window": messages } with open(filepath, 'w', encoding='utf-8') as f: json.dump(data_to_save, f, indent=2, ensure_ascii=False) except Exception as e: logger.error(f"Failed to save trace file: {e}") 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.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) 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 if debug: 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"[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: if debug: logger.error("Snippet repair failed to parse. Falling back to full repair.") pass 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(debug, client, messages, keep_last_pairs=2): keep_count = keep_last_pairs * 2 if len(messages) < (2 + 2 + keep_count): if debug: 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 ---" ) if debug: logger.debug(f"Compressing {len(to_compress)} messages...") 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:] if debug: logger.info(f"Compression complete. Reduced {len(messages)} msgs to {len(new_messages)}.") return new_messages def generate_final_report(debug, 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). """ if debug: logger.debug("Generating natural language report...") return client.completion([ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ]) def load_file(filepath): try: with open(filepath, 'r', encoding='utf-8') as f: return f.read() except FileNotFoundError: logger.error(f"File not found: {filepath}") sys.exit(1)