From 17a70e08d5dd4c3ae7f43c11e0715975ddaa4a5b Mon Sep 17 00:00:00 2001
From: geens <niels.geens@gmail.com>
Date: Fri, 1 Nov 2024 20:38:44 +0100
Subject: [PATCH] Updated main file

---
 sia/__main__.py         | 62 ++++++++++++++++++++++++++++++++++-------
 test/base_agent_test.py | 34 ++++++++++++++++------
 2 files changed, 77 insertions(+), 19 deletions(-)

diff --git a/sia/__main__.py b/sia/__main__.py
index 7d686fa..3625312 100644
--- a/sia/__main__.py
+++ b/sia/__main__.py
@@ -1,17 +1,59 @@
+import asyncio
+from pathlib import Path
+
 from .llm_engine import LlmEngine
+from .system_metrics import SystemMetrics
+from .web_agent import WebAgent
+from .web_io_buffer import WebIOBuffer
+from .web_server import WebServer
+from .working_memory import WorkingMemory
+from .xml_validator import XMLValidator
+from .response_parser import ResponseParser
+
+async def run_web_server():
+    """Run the web server with default configuration."""
+    # Load system prompt and action schema
+    base_dir = Path(__file__).parent.parent
+    with open(base_dir / "system_prompt.txt", "r") as f:
+        system_prompt = f.read()
+    with open(base_dir / "action_schema.xsd", "r") as f:
+        action_schema = f.read()
+
+    # Initialize core components
+    llm = LlmEngine("/root/model")
+    working_memory = WorkingMemory()
+    system_metrics = SystemMetrics()
+    validator = XMLValidator(action_schema)
+    io_buffer = WebIOBuffer()
+    parser = ResponseParser(io_buffer)
+
+    # Create web agent
+    agent = WebAgent(
+        system_prompt=system_prompt,
+        action_schema=action_schema,
+        working_memory=working_memory,
+        system_metrics=system_metrics,
+        llm=llm,
+        validator=validator,
+        parser=parser
+    )
+
+    # Create and start web server
+    server = WebServer(agent, io_buffer)
+    await server.start()
+
+    try:
+        # Keep running until interrupted
+        while True:
+            await asyncio.sleep(1)
+    except KeyboardInterrupt:
+        print("\nShutting down...")
+    finally:
+        await server.stop()
 
 def main():
     """Main entry point for the SIA application."""
-    system_prompt_path = "system_prompt.txt"
-    action_schema_path = "action_schema.xsd"
-    model_path = "/root/model"
-
-    with open(system_prompt_path, 'r') as f:
-        system_prompt = f.read()
-    with open(action_schema_path, 'r') as f:
-        action_schema = f.read()
-
-    print("todo")
+    asyncio.run(run_web_server())
 
 if __name__ == "__main__":
     main()
\ No newline at end of file
diff --git a/test/base_agent_test.py b/test/base_agent_test.py
index 12cd5e3..2418031 100644
--- a/test/base_agent_test.py
+++ b/test/base_agent_test.py
@@ -18,10 +18,17 @@ from sia.parse_error_entry import ParseErrorEntry
 from sia.delete_command import DeleteCommand
 from sia.entry import Entry
 
-class TestBaseAgent(BaseAgent):
-    """Concrete implementation of BaseAgent for testing."""
-    def run(self) -> None:
-        pass
+class TestLogHandler(logging.Handler):
+    """Custom logging handler that stores records for test verification."""
+    def __init__(self):
+        super().__init__()
+        self.records = []
+
+    def emit(self, record):
+        self.records.append(record)
+
+    def clear(self):
+        self.records = []
 
 class MockEntry(Entry):
     """Mock entry class that properly extends Entry."""
@@ -40,6 +47,11 @@ class MockEntry(Entry):
         elem.text = "<![CDATA[mock content]]>"
         return elem
 
+class TestBaseAgent(BaseAgent):
+    """Concrete implementation of BaseAgent for testing."""
+    def run(self) -> None:
+        pass
+
 class BaseAgentTest(unittest.TestCase):
     def setUp(self):
         """Set up test cases with mocked components."""
@@ -82,10 +94,10 @@ class BaseAgentTest(unittest.TestCase):
         # Set timestamp for entries
         self.test_timestamp = datetime(2024, 10, 31, 12, 0, 0)
         
-        # Capture logging output
-        self.log_records = []
-        logging.getLogger().handlers = []
-        logging.getLogger().addHandler(self.log_handler)
+        # Setup logging
+        self.log_handler = TestLogHandler()
+        logger = logging.getLogger()
+        logger.addHandler(self.log_handler)
         
     def log_handler(self, record):
         """Capture log records for testing."""
@@ -93,6 +105,10 @@ class BaseAgentTest(unittest.TestCase):
     
     def tearDown(self):
         """Clean up resources."""
+        # Remove log handler
+        logger = logging.getLogger()
+        logger.removeHandler(self.log_handler)
+        
         if hasattr(self, 'agent'):
             del self.agent
             
@@ -170,4 +186,4 @@ class BaseAgentTest(unittest.TestCase):
         self.assertEqual(len(reasoning_elems), 3)
         for i, elem in enumerate(reasoning_elems):
             self.assertEqual(elem.get("id"), f"id-{i}")
-            self.assertEqual(elem.text, f"<![CDATA[test {i}]]>")
\ No newline at end of file
+            self.assertEqual(elem.text, f"<![CDATA[test {i}]]>")