240 lines
8.0 KiB
Python
240 lines
8.0 KiB
Python
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()
|
|
|
|
|
|
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, messages, step=0):
|
|
try:
|
|
data_to_save = {
|
|
"timestamp": time.time(),
|
|
"step": step,
|
|
"context_window": 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 _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
|
|
|
|
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 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)
|