Added iteration_logger

sia/__main__.py

@@ -1,13 +1,10 @@
-from concurrent.futures import ThreadPoolExecutor
 from aiohttp import web
-from pathlib import Path
 import asyncio
 import mimetypes
-import time
 
 from .config import Config
 from .hf_llm_engine import HfLlmEngine
-from .llm_engine import LlmEngine
+from .iteration_logger import IterationLogger
 from .local_llm_engine import LocalLlmEngine
 from .mistral_llm_engine import MistralLlmEngine
 from .openai_llm_engine import OpenAILlmEngine
@@ -40,8 +37,7 @@ class Main:
                 self._llm = LocalLlmEngine(
                     self._config.model,
                     self._config.temperature,
-                    self._config.token_limit,
-                    self._config.api_token,
+                    self._config.token_limit
                 )
             case "hf":
                 self._llm = HfLlmEngine(
@@ -73,7 +69,8 @@ class Main:
             metrics=SystemMetrics(),
             llm=self._llm,
             validator=XMLValidator(self._action_schema),
-            parser=ResponseParser(self._io_buffer)
+            parser=ResponseParser(self._io_buffer),
+            iteration_logger=IterationLogger(self._config.iterations_dir, self._system_prompt, self._action_schema),
         )
         self._ws_manager = await WebSocketManager.create(
             self._agent,
@@ -132,4 +129,4 @@ if __name__ == "__main__":
     config = Config()
     main = loop.run_until_complete(Main.create(config))
     print(f"Web server started at http://localhost:{config.port}")
-    web.run_app(main.app, loop=loop, host=config.host, port=config.port)
+    web.run_app(main.app, loop=loop, host=config.host, port=config.port)

sia/base_agent.py

@@ -1,7 +1,5 @@
-from abc import ABC, abstractmethod
-from typing import List
+from abc import ABC
 import xml.etree.ElementTree as ET
-import time
 
 from .llm_engine import LlmEngine
 from .response_parser import ResponseParser
@@ -26,7 +24,7 @@ class BaseAgent(ABC):
             metrics: SystemMetrics,
             llm: LlmEngine,
             validator: XMLValidator,
-            parser: ResponseParser
+            parser: ResponseParser,
         ):
         """
         Initialize agent with required components.

sia/config.py

@@ -33,6 +33,12 @@ class Config:
             default=os.getenv('SIA_ACTION_SCHEMA', 'action_schema.xsd'),
             help='Path to the action schema file (default: action_schema.xsd, env: SIA_ACTION_SCHEMA)'
         )
+        parser.add_argument(
+            '--iterations-dir',
+            type=Path,
+            default=os.getenv('SIA_ITERATIONS_DIR', 'iterations'),
+            help='Path to the directory for storing iterations (default: iterations, env: SIA_ITERATIONS_DIR)'
+        )
         parser.add_argument(
             '--server',
             action='store_true',
@@ -112,6 +118,11 @@ class Config:
     def action_schema(self) -> Path:
         """Path to the action schema file."""
         return self.args.action_schema
+    
+    @property
+    def iterations_dir(self) -> Path:
+        """Path to the directory for storing iterations."""
+        return self.args.iterations_dir
 
     @property
     def server(self) -> bool:

sia/iteration_logger.py

@@ -0,0 +1,49 @@
+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)

sia/web_agent.py

@@ -5,6 +5,7 @@ from typing import Callable, List, Optional
 from .base_agent import BaseAgent
 from .command import Command
 from .command_result import CommandResult
+from .iteration_logger import IterationLogger
 from .llm_engine import LlmEngine
 from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
@@ -41,7 +42,8 @@ class WebAgent(BaseAgent):
             metrics: SystemMetrics,
             llm: LlmEngine,
             validator: XMLValidator,
-            parser: ResponseParser
+            parser: ResponseParser,
+            iteration_logger: IterationLogger,
         ):
         """
         Initialize web agent with required components.
@@ -63,6 +65,7 @@ class WebAgent(BaseAgent):
         self._command_result: Optional[CommandResult] = None
         self._context = self._compile_context()
         self._state_lock = Lock()
+        self._iteration_logger = iteration_logger
 
     @property
     def state(self) -> WebAgentState:
@@ -166,6 +169,7 @@ class WebAgent(BaseAgent):
         print()
     
     def _approve_response_thread(self) -> None:
+        self._iteration_logger.log_iteration(self._context, self._response)
         parse_result = self._parser.parse(self._response)
         if isinstance(parse_result, Command):
             result = parse_result.execute(self._working_memory)