New web interface, move llm engine to separate process

lib/llm_engine_utils/pyproject.toml

@@ -0,0 +1,13 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "llm_engine_utils"
+version = "0.1.0"
+requires-python = ">=3.8"
+
+[project.optional-dependencies]
+dataset = [
+    "torch>=4.0.0",
+]

lib/llm_engine_utils/src/llm_engine_utils/__init__.py

@@ -0,0 +1,7 @@
+try:
+    from . import dataset
+except ImportError:
+    pass
+from . import iterators
+from . import protocol
+from .llm_engine import LlmEngine

lib/llm_engine_utils/src/llm_engine_utils/dataset.py

@@ -0,0 +1,144 @@
+from datasets import Dataset as TransformersDataset
+from transformers import PreTrainedTokenizer
+from pathlib import Path
+from typing import Dict, List, Iterator
+import hashlib
+import torch
+import xml.etree.ElementTree as ET
+import yaml
+
+class Dataset(torch.utils.data.Dataset):
+    """Training dataset from XML iteration files"""
+    
+    def __init__(self, config_filename: str):
+        with open(config_filename) as f:
+            config_data = yaml.safe_load(f)
+        
+        data_paths = [Path(p) for p in config_data['data']]
+        self.files = self._find_xml_files(data_paths)
+        
+        self.system_prompt_file = Path(config_data['model']['system_prompt_path'])
+        self.action_schema_file = Path(config_data['model']['action_schema'])
+        
+        self.system_prompt = self.system_prompt_file.read_text()
+        self.system_prompt_hash = self._calculate_hash(self.system_prompt)
+        
+        self.action_schema = self.action_schema_file.read_text()
+        self.action_schema_hash = self._calculate_hash(self.action_schema)
+        
+    def _find_xml_files(self, data_paths: List[Path]) -> List[Path]:
+        """Find all XML files in the given data paths"""
+        xml_files = list()
+        for path in data_paths:
+            if not path.exists():
+                raise Exception(f"Data path not found: {path}")
+            xml_files.extend(path.rglob('*.xml'))
+        return xml_files
+        
+    def _calculate_hash(self, content: str) -> str:
+        """Calculate SHA-256 hash of content"""
+        return hashlib.sha256(content.encode()).hexdigest()
+    
+    def _parse_iteration_file(self, file_path: Path) -> Dict:
+        """Parse a single iteration XML file into a training example"""
+        tree = ET.parse(file_path)
+        root = tree.getroot()
+        
+        context_elem = root.find('context')
+        response_elem = root.find('response')
+        
+        context = context_elem.text
+        response = response_elem.text
+        
+        return {
+            "messages": [
+                {
+                    "role": "system",
+                    "content": self.system_prompt + "\n" + self.action_schema
+                },
+                {
+                    "role": "user",
+                    "content": context
+                },
+                {
+                    "role": "assistant", 
+                    "content": response
+                }
+            ]
+        }
+    
+    def __len__(self) -> int:
+        """Return the number of samples in the dataset"""
+        return len(self.files)
+    
+    def __getitem__(self, idx: int) -> Dict:
+        """Indexing for a single sample"""
+        if idx < 0 or idx >= len(self):
+            raise IndexError(f"Index {idx} out of range for dataset with {len(self)} samples")
+        file_path = self.files[idx]
+        return self._parse_iteration_file(file_path)
+        
+    def __iter__(self) -> Iterator[Dict]:
+        """Allow iteration over samples"""
+        for i in range(len(self)):
+            yield self[i]
+    
+    def to_list(self) -> List[Dict]:
+        """Convert dataset to a list"""
+        results = []
+        for i in range(len(self)):
+            results.append(self[i])
+        return results
+    
+    def to_transformers_dataset(self, tokenizer: PreTrainedTokenizer) -> TransformersDataset:
+        def generator():
+            for item in self:
+                messages = item["messages"]
+                formatted_text = tokenizer.apply_chat_template(
+                    messages,
+                    tokenize=False,
+                    add_generation_prompt=False
+                )
+                yield {"messages": formatted_text}
+        return TransformersDataset.from_generator(generator)
+    
+    def validate(self) -> None:
+        """Validate XML files"""
+        print(f"Validating {len(self.files)} XML files...")
+        
+        for i in range(len(self.files)):
+            self.validate_sample(i)
+
+        print(f"Validation complete. Found {len(self.files)} valid files.")
+    
+    def validate_sample(self, index: int) -> None:
+        file = self.files[index]
+        print("file:", file)
+        tree = ET.parse(file)
+        root = tree.getroot()
+        
+        # Check system prompt hash
+        file_system_hash = root.get('system_prompt_hash')
+        if file_system_hash != self.system_prompt_hash:
+            print(f"WARNING: System prompt hash mismatch in {file}")
+        
+        # Check action schema hash
+        file_schema_hash = root.get('action_schema_hash')
+        if file_schema_hash != self.action_schema_hash:
+            print(f"WARNING: Action schema hash mismatch in {file}")
+        
+        # Check for required elements
+        context_elem = root.find('context')
+        response_elem = root.find('response')
+        
+        if context_elem is None:
+            raise Exception(f"Missing context element")
+            
+        if response_elem is None:
+            raise Exception(f"Missing response element")
+            
+        if not context_elem.text:
+            raise Exception(f"Empty context")
+            
+        if not response_elem.text:
+            raise Exception(f"Empty response")

lib/llm_engine_utils/src/llm_engine_utils/iterators.py

@@ -0,0 +1,61 @@
+from typing import Iterator
+
+def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]:
+    """
+    Creates an iterator that yields values from the input iterator
+    until it encounters the stop_value (exclusive).
+    
+    Args:
+        iterator: The source iterator
+        stop_value: The value to stop before
+        
+    Yields:
+        Values from the iterator until stop_value is encountered
+        If stop_value is part of an item, yields the part before stop_value
+    """
+    for item in iterator:
+        if stop_value in item:
+            split_point = item.index(stop_value)
+            if split_point > 0:
+                yield item[:split_point]
+            break
+        yield item
+
+def skip_prefix(iterator: Iterator[str], prefix: str) -> Iterator[str]:
+    """
+    Creates an iterator that skips a prefix from the input iterator
+    and yields only the content after the prefix.
+    
+    Args:
+        iterator: The source iterator
+        prefix: The prefix to skip
+        
+    Yields:
+        Values from the iterator after the prefix has been fully skipped
+    """
+    if not prefix:
+        # If no prefix to skip, yield everything
+        yield from iterator
+        return
+    
+    prefix_remaining = prefix
+    
+    for item in iterator:
+        if prefix_remaining:
+            # If the item starts with the remaining prefix
+            if prefix_remaining.startswith(item):
+                # Skip this item entirely
+                prefix_remaining = prefix_remaining[len(item):]
+                continue
+            elif item.startswith(prefix_remaining):
+                # Yield only the part after the prefix
+                yield item[len(prefix_remaining):]
+                prefix_remaining = ""
+            else:
+                # Item doesn't match prefix pattern, yield everything
+                # This is unexpected but we handle it gracefully
+                yield item
+                prefix_remaining = ""
+        else:
+            # No prefix remaining, yield all content
+            yield item

lib/llm_engine_utils/src/llm_engine_utils/llm_engine.py

@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Iterator
+
+class LlmEngine(ABC):
+    @abstractmethod
+    def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]:
+        pass
+
+    @abstractmethod
+    def token_count(self, system: str, context: str) -> int:
+        pass
+
+    @abstractmethod
+    def token_limit(self) -> int:
+        pass

lib/llm_engine_utils/src/llm_engine_utils/protocol.py

@@ -0,0 +1,87 @@
+"""
+Protocol handler for SIA LLM engine subprocess communication.
+
+This module provides tools to parse and handle the XML-based communication
+protocol between the SIA core and LLM engine subprocesses as described
+in the SIA README.
+"""
+
+from .llm_engine import LlmEngine
+from io import StringIO
+from pathlib import Path
+from typing import Optional
+import sys
+import xml.etree.ElementTree as ET
+
+# ASCII End of Transmission character used to signal the end of a response
+EOT = '\u0004'
+
+def read_command() -> Optional[ET.Element]:
+    """
+    Read an XML command from stdin.
+    
+    Returns:
+        ET.Element: The parsed XML element or None if EOF was reached.
+    """
+    buffer = StringIO()
+    
+    while True:
+        char = sys.stdin.read(1)
+        if not char:  # EOF
+            return None
+        
+        buffer.write(char)
+        content = buffer.getvalue()
+        
+        # Try to parse once we have a potentially complete XML document
+        if content.endswith('>'):
+            try:
+                return ET.fromstring(content)
+            except ET.ParseError:
+                # Not a complete XML document yet, continue reading
+                pass
+        
+def process(engine: LlmEngine) -> None:
+    """
+    Run the adapter loop, processing commands from stdin and writing responses to stdout.
+    """
+    while True:
+        print("reading command", file=sys.stderr)
+        command = read_command()
+        if command is None:
+            print("EOF", file=sys.stderr)
+            # End of input stream, exit
+            return
+
+        if command.tag == 'token_limit':
+            print("token_limit", file=sys.stderr)
+            limit = engine.token_limit()
+            sys.stdout.write(str(limit))
+            sys.stdout.write(EOT)
+            sys.stdout.flush()
+            
+        elif command.tag == 'token_count':
+            print("token_count", file=sys.stderr)
+            system = command.find('system').text
+            context = command.find('context').text
+            count = engine.token_count(system, context)
+            sys.stdout.write(str(count))
+            sys.stdout.write(EOT)
+            sys.stdout.flush()
+            
+        elif command.tag == 'infer_xml':
+            print("infer_xml", file=sys.stderr)
+            try:
+                schema = Path(command.find('schema').text)
+                system = command.find('system').text
+                context = command.find('context').text
+                prefix = command.find('prefix').text if command.find('prefix') is not None else None
+                for chunk in engine.infer_xml(schema, system, context, prefix):
+                    sys.stdout.write(chunk)
+                    sys.stdout.flush()
+                sys.stdout.write(EOT)
+                sys.stdout.flush()
+            except Exception as e:
+                print(f"Error: {e}", file=sys.stderr)
+                sys.stdout.write(EOT)
+                sys.stdout.flush()