Begin work on agent core

sia/__main__.py

@@ -1,5 +1,25 @@
+from .agent_core import AgentCore
+from .llm_engine import LlmEngine
+from .docker_module import DockerModule
+
 def main():
-    print("Hello, World! --sia")
+    """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()
+
+    agent_core = AgentCore(
+        system_prompt=f"{system_prompt}{action_schema}",
+        action_schema=action_schema,
+        docker_module=DockerModule(),
+        llm_engine=LlmEngine(model_path=model_path)
+    )
+    agent_core.run_iteration()
 
 if __name__ == "__main__":
     main()

sia/agent_core.py

@@ -0,0 +1,51 @@
+from itertools import tee
+from typing import Optional, Dict, List
+
+from .docker_module import DockerModule
+from .llm_engine import LlmEngine
+from .context_template import generate_context
+from .util import get_valid_root_elements, split_response
+
+class AgentCore:
+    """
+    Core orchestration class for SIA that manages the interaction between
+    different modules and runs the main agent loop.
+    """
+    
+    def __init__(
+        self, 
+        system_prompt: str,
+        action_schema: str,
+        docker_module: DockerModule,
+        llm_engine: LlmEngine
+    ):
+        """
+        Initialize the AgentCore with required components.
+        
+        Args:
+            system_prompt: System prompt to use for the LLM
+            action_schema_path: Path to the XML schema defining valid actions
+            docker_module: DockerModule instance
+            llm_engine: LLmEngine instance
+        """
+        self.system_prompt = system_prompt
+        self.action_schema = action_schema
+        self.docker_module = docker_module
+        self.llm_engine = llm_engine
+        
+        self.valid_elements = get_valid_root_elements(self.action_schema)
+    
+    def run_iteration(self) -> None:
+        """Run a single iteration of the main agent loop."""
+        containers = self.docker_module.get_all_container_statuses()
+        context = generate_context(containers)
+        tokens = self.llm_engine.infer(
+            self.system_prompt,
+            context
+        )
+        print_tokens, response_tokens = tee(tokens)
+        for token in print_tokens:
+            print(token, end="", flush=True)
+        response_tokens = ''.join(response_tokens)
+        result = split_response(response_tokens, self.valid_elements)
+        print(f"result: {result}")

sia/docker_module.py

@@ -80,7 +80,7 @@ class DockerModule:
         volumes: Optional[Dict[str, str]] = None,
         ports: Optional[Dict[str, str]] = None,
         environment: Optional[Dict[str, str]] = None,
-    ) -> str:
+    ) -> Optional[str]:
         """Start a new Docker container with the specified configuration.
 
         Args:
@@ -95,11 +95,7 @@ class DockerModule:
 
         Returns:
             For short-lived containers (with timeout): Container output
-            For long-running containers: Container name
-
-        Raises:
-            docker.errors.ImageNotFound: If specified image doesn't exist
-            docker.errors.APIError: If container creation/start fails
+            For long-running containers: None
         """
         # Validate inputs
         if name is None and timeout < 0:
@@ -143,7 +139,6 @@ class DockerModule:
         else:
             # For long-running containers
             self.containers[name] = (container, socket)
-            return name
 
     def write_container_stdin(self, name: str, data: str) -> None:
         """Write data to a container's standard input.
@@ -225,11 +220,6 @@ class DockerModule:
 
         Returns:
             Tuple of (exit_code, output)
-
-        Raises:
-            KeyError: If container name not found
-            docker.errors.APIError: If wait operation fails
-            TimeoutError: If container doesn't finish within timeout
         """
         if name not in self.containers:
             raise KeyError(f"Container {name} not found")
@@ -277,4 +267,5 @@ class DockerModule:
         """
         statuses = []
         for name, (container, socket) in self.containers.items():
-            statuses.append(self.get_container_status(name))
+            statuses.append(self.get_container_status(name))
+        return statuses

sia/llm_engine.py

@@ -1,4 +1,5 @@
 from threading import Thread
+from typing import Iterator
 
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
 import torch
@@ -44,7 +45,7 @@ class LlmEngine:
             return_full_text=False,
         )
 
-    def infer(self, system_prompt: str, main_context: str) -> TextIteratorStreamer:
+    def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
         """
         Run inference using the system prompt and main context, while validating actions against the provided XML schema.
 
@@ -53,7 +54,7 @@ class LlmEngine:
             main_context: The main context string after templating
 
         Returns:
-            TextIteratorStreamer: An iterator that yields the generated text.
+            Iterator[str]: An iterator that yields the generated text.
         """
         messages = [
             {"role": "system", "content": system_prompt},

sia/util.py

@@ -1,4 +1,4 @@
-from typing import Iterator, TypeVar
+from typing import Iterator
 import re
 import xml.etree.ElementTree as ET