SIA/sia/llm_engine.py

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()