Files
SIA/sia/iteration_logger.py
2024-11-19 20:24:21 +01:00

49 lines
1.6 KiB
Python

from datetime import datetime
from pathlib import Path
import xml.etree.ElementTree as ET
import hashlib
class IterationLogger:
"""Logs agent iterations to XML files"""
def __init__(
self,
iterations_dir: Path,
system_prompt: str,
action_schema: str,
):
"""Initialize with directory for storing iteration files"""
self.iterations_dir = iterations_dir
self.iterations_dir.mkdir(parents=True, exist_ok=True)
self._system_prompt_hash = hashlib.sha256(system_prompt.encode()).hexdigest()
self._action_schema_hash = hashlib.sha256(action_schema.encode()).hexdigest()
def log_iteration(
self,
context: str,
response: str,
):
"""
Save an iteration to an XML file
Args:
context: The context as ElementTree
response: Raw response from LLM
"""
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
filename = f"iteration_{timestamp}.xml"
filepath = self.iterations_dir / filename
root = ET.Element("iteration")
root.set("system_prompt_hash", self._system_prompt_hash)
root.set("action_schema_hash", self._action_schema_hash)
context_elem = ET.SubElement(root, "context")
context_elem.text = context
response_elem = ET.SubElement(root, "response")
response_elem.text = response
tree = ET.ElementTree(root)
tree.write(filepath, encoding="utf-8", xml_declaration=True)