Merge branch 'master' of https://git.nielsgeens.be/llm/SIA

sia/__main__.py

@@ -5,11 +5,11 @@ import asyncio
 import mimetypes
 import time
 
-
 from .config import Config
 from .hf_llm_engine import HfLlmEngine
 from .llm_engine import LlmEngine
 from .local_llm_engine import LocalLlmEngine
+from .openai_llm_engine import OpenAILlmEngine
 from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
 from .web_agent import WebAgent
@@ -51,6 +51,11 @@ class Main:
                     api_token=self._config.api_token,
                     temperature=self._config.temperature
                 )
+            case "openai":
+                self._llm = OpenAILlmEngine(
+                    model=self._config.model,
+                    api_key=self._config.api_token
+                )
             case "test":
                 self._llm = TestLLM()
             case _:

sia/openai_llm_engine.py

@@ -0,0 +1,57 @@
+from typing import Iterator, Optional
+import openai
+import json
+
+from .llm_engine import LlmEngine
+
+class OpenAILlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using OpenAI's API.
+    Supports streaming responses from chat completion models.
+    """
+    
+    def __init__(
+        self,
+        model: str,
+        api_key: str,
+    ):
+        """
+        Initialize the OpenAI LLM Engine.
+
+        Args:
+            model: OpenAI model to use (default: gpt-4)
+            api_key: OpenAI API key. If None, will try to read from OPENAI_API_KEY env var
+        """
+        self._model = model
+        
+        # Initialize OpenAI client
+        self.client = openai.Client(
+            api_key=api_key,
+        )
+                        
+    def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
+        """
+        Run inference using the system prompt and main context.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text in chunks.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        
+        stream = self.client.chat.completions.create(
+            model=self._model,
+            messages=messages,
+            stream=True,
+            temperature=0.3,
+        )
+
+        for chunk in stream:
+            if content := chunk.choices[0].delta.content:
+                yield content

sia/response_parser.py

@@ -36,46 +36,6 @@ class ResponseParser:
             io_buffer: Buffer to use for IO operations
         """
         self.io_buffer = io_buffer
-
-    def _extract_content(self, xml: str, root: ET.Element) -> str:
-        """
-        Extract content from XML, handling CDATA sections and preserving special characters.
-        
-        Args:
-            xml: Original XML string
-            root: Parsed XML element
-            
-        Returns:
-            str: Extracted content with proper handling of CDATA and special characters
-        """
-        if root.text is None:
-            return ""
-            
-        # Find opening tag end
-        tag_name = root.tag
-        tag_start = xml.find('<' + tag_name)
-        tag_end = xml.find('>', tag_start)
-        if tag_end == -1:
-            return ""
-            
-        # Check for self-closing tag
-        if xml[tag_end-1] == '/':
-            return ""
-            
-        # Find closing tag start
-        close_tag = '</' + tag_name + '>'
-        close_start = xml.rfind(close_tag)
-        if close_start == -1:
-            return ""
-            
-        # Extract content
-        content = xml[tag_end+1:close_start]
-        
-        # Handle CDATA sections
-        if content.startswith('<![CDATA[') and content.endswith(']]>'):
-            content = content[9:-3]  # Remove CDATA markers
-            
-        return content
         
     def parse(self, xml: str) -> Union[Command, Entry]:
         """
@@ -91,74 +51,50 @@ class ResponseParser:
         entry_id = str(uuid.uuid4())
         timestamp = datetime.now()
         
-        # First try to parse the XML
         try:
             root = ET.fromstring(xml.strip())
         except ET.ParseError as e:
             return ParseErrorEntry(xml, f"Invalid XML: {str(e)}", entry_id, timestamp)
-            
-        # Extract content preserving special characters and handling CDATA
-        content = self._extract_content(xml, root)
         
-        # Convert to appropriate type based on root tag
         try:
             if root.tag == 'delete':
-                # Extract target ID and create delete command
                 target_id = root.get('id')
                 if not target_id:
                     return ParseErrorEntry(xml, "Delete command missing required 'id' attribute", entry_id, timestamp)
                 if len(root.attrib) > 1:
                     return ParseErrorEntry(xml, "Delete command should only have 'id' attribute", entry_id, timestamp)
                 return DeleteCommand(target_id)
+
             elif root.tag == 'stop':
-                # Create stop command
                 return StopCommand()
                 
             elif root.tag == 'background':
-                # Create background script entry
-                if not content:
-                    return ParseErrorEntry(xml, "Background entry missing script content", entry_id, timestamp)
-                if root.attrib:
-                    return ParseErrorEntry(xml, "Background entry shouldn't have attributes", entry_id, timestamp)
-                return BackgroundEntry(content, entry_id, timestamp)
+                if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
+                    return ParseErrorEntry(xml, "Background entry requires (only) script content", entry_id, timestamp)
+                return BackgroundEntry(root.text, entry_id, timestamp)
                 
             elif root.tag == 'repeat':
-                # Create repeat script entry
-                if not content:
-                    return ParseErrorEntry(xml, "Repeat entry missing script content", entry_id, timestamp)
-                if root.attrib:
-                    return ParseErrorEntry(xml, "Repeat entry shouldn't have attributes", entry_id, timestamp)
-                return RepeatEntry(content, entry_id, timestamp)
+                if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
+                    return ParseErrorEntry(xml, "Repeat entry requires (only) script content", entry_id, timestamp)
+                return RepeatEntry(root.text, entry_id, timestamp)
                 
             elif root.tag == 'single':
-                # Create single shot script entry
-                if not content:
-                    return ParseErrorEntry(xml, "Single entry missing script content", entry_id, timestamp)
-                if root.attrib:
-                    return ParseErrorEntry(xml, "Single entry shouldn't have attributes", entry_id, timestamp)
-                return SingleEntry(content, entry_id, timestamp)
+                if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
+                    return ParseErrorEntry(xml, "Single entry requires (only) script content", entry_id, timestamp)
+                return SingleEntry(root.text, entry_id, timestamp)
                 
             elif root.tag == 'reasoning':
-                # Create reasoning entry
-                if not content:
-                    return ParseErrorEntry(xml, "Reasoning entry missing content", entry_id, timestamp)
-                if root.attrib:
-                    return ParseErrorEntry(xml, "Reasoning entry shouldn't have attributes", entry_id, timestamp)
-                return ReasoningEntry(content, entry_id, timestamp)
+                if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
+                    return ParseErrorEntry(xml, "Reasoning entry requires (only) text content", entry_id, timestamp)
+                return ReasoningEntry(root.text, entry_id, timestamp)
                 
             elif root.tag == 'read_stdin':
-                # Create read entry - no content required
-                if root.attrib:
-                    return ParseErrorEntry(xml, "Read stdin entry shouldn't have attributes", entry_id, timestamp)
                 return ReadEntry(self.io_buffer, entry_id, timestamp)
                 
             elif root.tag == 'write_stdout':
-                # Create write entry
-                if not content:
-                    return ParseErrorEntry(xml, "Write entry missing content", entry_id, timestamp)
-                if root.attrib:
-                    return ParseErrorEntry(xml, "Read stdin entry shouldn't have attributes", entry_id, timestamp)
-                return WriteEntry(content, self.io_buffer, entry_id, timestamp)
+                if len(root) != 0 or root.attrib or root.text is None or root.text.strip() == '':
+                    return ParseErrorEntry(xml, "Write stdout entry requires (only) text content", entry_id, timestamp)
+                return WriteEntry(root.text, self.io_buffer, entry_id, timestamp)
                 
             else:
                 return ParseErrorEntry(xml, f"Unknown root element: {root.tag}", entry_id, timestamp)