System propt and action schema

sia/__main__.py

@@ -12,40 +12,33 @@ 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()
+
+    # Load system prompt and action schema
+    system_prompt = (base_dir / "system_prompt.txt").read_text()
+    action_schema = (base_dir / "action_schema.xsd").read_text()
 
     # 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,
+        working_memory=WorkingMemory(),
+        system_metrics=SystemMetrics(),
         llm=llm,
-        validator=validator,
-        parser=parser
+        validator=XMLValidator(action_schema),
+        parser=ResponseParser(io_buffer)
     )
 
-    # Create and start web server
+    # Start the web server
     server = WebServer(agent, io_buffer, "0.0.0.0", 8080)
     await server.start()
 
     try:
         while True:
             if server._response_tasks:
-                done, pending = await asyncio.wait(
+                done, _ = await asyncio.wait(
                     server._response_tasks,
                     timeout=0.1,
                     return_when=asyncio.FIRST_COMPLETED
@@ -53,9 +46,11 @@ async def run_web_server():
                 for task in done:
                     try:
                         await task
+                    except asyncio.CancelledError:
+                        print("Task was cancelled")
                     except Exception as e:
                         print(f"Task error: {e}")
-                    server._response_tasks.remove(task)
+                    server._response_tasks.discard(task)  # Use discard to prevent KeyError
             await asyncio.sleep(0.1)
     except KeyboardInterrupt:
         print("\nShutting down...")
@@ -67,4 +62,4 @@ def main():
     asyncio.run(run_web_server())
 
 if __name__ == "__main__":
-    main()
+    main()