New web interface, move llm engine to separate process

sia/llm_engine.py

@@ -0,0 +1,212 @@
+from typing import Iterator
+from .util import pretty_print_element
+import subprocess
+import sys
+import xml.etree.ElementTree as ET
+
+class LlmEngine:
+    """
+    LlmEngine manages communication with LLM engine subprocesses.
+    Each LLM type runs in its own subprocess with a tailored environment.
+    """
+
+    EOT = '\x04'  # EOT character (ASCII 4) as bytes
+    
+    def __init__(self, executable_path: str, action_schema_path: str):
+        """
+        Initialize the LLM engine subprocess.
+        
+        Args:
+            executable_path (str): Path to the LLM engine executable
+            action_schema_path (str): Path to the XML action schema
+        """
+        self.executable_path = executable_path
+        self.action_schema_path = action_schema_path
+        self.action_schema = open(action_schema_path, 'r').read()
+        self.process = None
+        self.restart()
+    
+    def _read_until_eot(self) -> str:
+        """
+        Read from subprocess stdout until EOT character.
+            
+        Returns:
+            str: Complete response without the EOT character
+        """
+        response = []
+        
+        while True:
+            # Read available data
+            data = self.process.stdout.read(1024)  # Read up to 1024 bytes at a time
+            
+            if not data:  # process died
+                self.restart()
+                raise RuntimeError(f"LLM subprocess terminated unexpectedly")
+
+            data = data.decode('utf-8')
+            
+            # Check if EOT is in the data
+            if self.EOT in data:
+                eot_index = data.index(self.EOT)
+                response.append(data[:eot_index])  # Add data before EOT
+                break
+            else:
+                response.append(data)
+        
+        return  "".join(response)
+    
+    def token_limit(self) -> int:
+        """
+        Get the maximum token limit of the LLM.
+        
+        Returns:
+            int: Maximum token limit
+        """
+        self.process.stdin.write(b"<token_limit/>\n")
+        self.process.stdin.flush()
+        response = self._read_until_eot()
+        return int(response.strip())
+    
+    def token_count(self, system_prompt: str, main_context: ET.Element, prefix: str = "") -> int:
+        """
+        Count the number of tokens in the prompt.
+        
+        Args:
+            system_prompt (str): System prompt text
+            main_context (ET.Element): Main context as ElementTree
+            prefix (str): Optional prefix for continuing generation
+            
+        Returns:
+            int: Token count
+        """
+        # Create the XML document
+        root = ET.Element("token_count")
+        
+        # Add system prompt
+        system_prompt_elem = ET.SubElement(root, "system")
+        system_prompt_elem.text = self._append_action_schema(system_prompt)
+        
+        # Add context element
+        context_elem = ET.SubElement(root, "context")
+        context_elem.text = pretty_print_element(main_context)
+        
+        # Add prefix if provided
+        if prefix:
+            prefix_elem = ET.SubElement(root, "prefix")
+            prefix_elem.text = prefix
+        
+        # Send to subprocess - convert to bytes
+        xml_str = ET.tostring(root, encoding='utf-8')
+        self.process.stdin.write(xml_str + b"\n")
+        self.process.stdin.flush()
+        
+        # Read response
+        response = self._read_until_eot()
+        return int(response.strip())
+    
+    def infer(self, system_prompt: str, main_context: ET.Element, prefix: str = "") -> Iterator[str]:
+        """
+        Generate text from the LLM.
+        
+        Args:
+            system_prompt (str): System prompt text
+            main_context (ET.Element): Main context as ElementTree
+            prefix (str): Optional prefix for continuing generation
+            
+        Returns:
+            Iterator[str]: Generated text, yielded as it's produced
+        """
+        # Create the XML document
+        root = ET.Element("infer_xml")
+        
+        # Add action schema path
+        schema_path_elem = ET.SubElement(root, "schema")
+        schema_path_elem.text = self.action_schema_path
+        
+        # Add system prompt in CDATA
+        system_prompt_elem = ET.SubElement(root, "system")
+        system_prompt_elem.text = self._append_action_schema(system_prompt)
+        
+        # Add context element
+        context_elem = ET.SubElement(root, "context")
+        context_elem.text = pretty_print_element(main_context)
+        
+        # Add prefix if provided
+        if prefix:
+            prefix_elem = ET.SubElement(root, "prefix")
+            prefix_elem.text = prefix
+        
+        # Send to subprocess - convert to bytes
+        xml_str = ET.tostring(root, encoding='utf-8')
+        self.process.stdin.write(xml_str + b"\n")
+        self.process.stdin.flush()
+        
+        while True:
+            # Read available data
+            data = self.process.stdout.read(1024)  # Read up to 1024 bytes at a time
+            
+            if not data:  # Process died
+                self.restart()
+                raise RuntimeError("LLM subprocess terminated unexpectedly")
+            
+            data = data.decode('utf-8')
+            
+            if self.EOT in data:
+                eot_index = data.index(self.EOT)
+                if eot_index > 0:
+                    yield data[:eot_index]
+                break
+            else:
+                yield data
+    
+    def restart(self):
+        """Start the LLM engine subprocess."""
+        # Ensure any existing process is terminated
+        if self.process:
+            self._terminate_process()
+        
+        # Start the subprocess with pipes for stdin/stdout and direct stderr
+        self.process = subprocess.Popen(
+            ["/bin/bash", "-c", self.executable_path],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=sys.stderr,
+            text=False, # Use binary mode to avoid buffering and encoding issues
+            bufsize=0,
+        )
+        
+        # Check if the process started successfully
+        if self.process.poll() is not None:
+            raise RuntimeError(f"Failed to start LLM engine at {self.executable_path}")
+    
+    def _terminate_process(self):
+        """Terminate the LLM engine subprocess safely."""
+        if self.process:
+            try:
+                self.process.stdin.close()
+                self.process.stdout.close()
+                self.process.terminate()
+                try:
+                    self.process.wait(timeout=5)  # Wait for process to terminate
+                except subprocess.TimeoutExpired:
+                    # Force kill if termination takes too long
+                    self.process.kill()
+                    self.process.wait(timeout=2)
+            finally:
+                self.process = None
+    
+    def _append_action_schema(self, system_prompt: str) -> str:
+        """
+        Append the action schema to the system prompt.
+
+        Args:
+            system_prompt (str): Original system prompt
+
+        Returns:
+            str: Updated system prompt with action schema
+        """
+        return f"{system_prompt}\n\n--- Action Schema ---\n{self.action_schema}"
+    
+    def __del__(self):
+        """Cleanup when the object is destroyed."""
+        self._terminate_process()