Added hf_llm_engine and config

.dockerignore

@@ -1 +1,2 @@
+.env
 ./model/

.gitignore

@@ -1,3 +1,4 @@
+.env
 pdf/
 model/
 claude.txt

Dockerfile

@@ -32,4 +32,4 @@ COPY ./ /root/sia/
 COPY --from=web-build /app/dist /root/sia/static/
 WORKDIR /root/sia
 
-CMD ["python3", "-m", "sia"]
+ENTRYPOINT ["python3", "-m", "sia"]

requirements.txt

@@ -1,5 +1,6 @@
 accelerate
 aiohttp
 bs4
+python-dotenv
 torch
 transformers

run.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 docker build \
-    --tag sia
+    --tag sia \
     .
 
 docker run \
@@ -10,4 +10,5 @@ docker run \
     --gpus=all \
     -p 8080:8080 \
     -v /$(pwd)/model/:/root/model/ \
-    sia
+    -v /$(pwd)/.env:/root/.env \
+    sia "$@"

sia/__main__.py

@@ -5,16 +5,20 @@ 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 .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
 from .web_agent import WebAgent
+from .web_agent import WebAgent
+from .web_io_buffer import WebIOBuffer
 from .web_io_buffer import WebIOBuffer
 from .web_socket_manager import WebSocketManager
 from .working_memory import WorkingMemory
 from .xml_validator import XMLValidator
-from .response_parser import ResponseParser
-from .web_agent import WebAgent
-from .web_io_buffer import WebIOBuffer
 
 mimetypes.add_type("application/javascript", ".js")
 mimetypes.add_type("application/javascript", ".jsx")
@@ -31,13 +35,23 @@ class TestLLM:
 
 class Main:
     def __init__(self):
-        self._base_dir = Path(__file__).parent.parent
+        self._config = Config()
-        self._system_prompt = (self._base_dir / "system_prompt.md").read_text()
-        self._action_schema = (self._base_dir / "action_schema.xsd").read_text()
-        self._static_dir = self._base_dir / "static"
 
-        self._llm = LlmEngine("/root/model")
+        self._system_prompt = self._config.system_prompt.read_text()
-        #self._llm = TestLLM()
+        self._action_schema = self._config.action_schema.read_text()
+
+        match self._config.llm_engine:
+            case "local":
+                self._llm = LocalLlmEngine(self._config.model)
+            case "hf":
+                self._llm = HfLlmEngine(
+                    model_id=self._config.model,
+                    api_token=self._config.hf_api_token
+                )
+            case "test":
+                self._llm = TestLLM()
+            case _:
+                raise ValueError(f"Invalid LLM engine: {self._config.llm_engine}")
         self._io_buffer = WebIOBuffer()
         self._agent = WebAgent(
             system_prompt=self._system_prompt,
@@ -64,13 +78,13 @@ class Main:
         self._app.middlewares.append(self._cors_middleware)
         self._app.router.add_get("/ws", self._ws_manager.handle_websocket)
         self._app.router.add_get("/", self._serve_index)
-        self._app.router.add_static("/static/", self._static_dir, show_index=False)
+        self._app.router.add_static("/static/", self._config.static_files, show_index=False)
-        self._app.router.add_static("/assets/", self._static_dir / "assets", show_index=False)
+        self._app.router.add_static("/assets/", self._config.static_files / "assets", show_index=False)
         self._app.router.add_get("/{path:.*}", self._serve_index)
 
     async def _serve_index(self, request: web.Request) -> web.Response:
         """Serve the React application HTML for any unmatched routes."""
-        index_path = self._static_dir / "index.html"
+        index_path = self._config.static_files / "index.html"
         if not index_path.exists():
             raise web.HTTPNotFound()
 

sia/config.py

@@ -0,0 +1,137 @@
+from dataclasses import dataclass
+from dotenv import load_dotenv
+from pathlib import Path
+from typing import Optional
+import argparse
+import os
+
+@dataclass
+class Config:
+    """
+    Configuration class that handles both command line and environment variables.
+    
+    Command line arguments take precedence over environment variables.
+    Environment variables serve as defaults that can be overridden via CLI.
+    """
+
+    def __init__(self):
+        """
+        Create configuration from command line arguments and environment variables.
+        Required arguments must be provided either via CLI or environment variables.
+        """
+        load_dotenv()
+        parser = argparse.ArgumentParser(description='SIA - Self Improving Agent')
+        parser.add_argument(
+            '--system-prompt',
+            type=Path,
+            default=os.getenv('SIA_SYSTEM_PROMPT', 'system_prompt.md'),
+            help='Path to the system prompt file (default: system_prompt.md, env: SIA_SYSTEM_PROMPT)'
+        )
+        parser.add_argument(
+            '--action-schema',
+            type=Path,
+            default=os.getenv('SIA_ACTION_SCHEMA', 'action_schema.xsd'),
+            help='Path to the action schema file (default: action_schema.xsd, env: SIA_ACTION_SCHEMA)'
+        )
+        parser.add_argument(
+            '--server',
+            action='store_true',
+            default=self._parse_bool_env('SIA_SERVER_ENABLED', False),
+            help='Enable web server for debugging and human feedback (env: SIA_SERVER_ENABLED)'
+        )
+        parser.add_argument(
+            '--host',
+            type=str,
+            default=os.getenv('SIA_SERVER_HOST', 'localhost'),
+            help='Web server host (default: localhost, env: SIA_SERVER_HOST)'
+        )
+        parser.add_argument(
+            '--port',
+            type=int,
+            default=int(os.getenv('SIA_SERVER_PORT', '8080')),
+            help='Web server port (default: 8080, env: SIA_SERVER_PORT)'
+        )
+        parser.add_argument(
+            '--static-files',
+            type=Path,
+            default=self._parse_optional_path('SIA_STATIC_FILES', './static/'),
+            help='Path to static web files (default: ./static/, env: SIA_STATIC_FILES)'
+        )
+        parser.add_argument(
+            '--llm-engine',
+            type=str,
+            default=os.getenv('SIA_LLM_ENGINE', 'local'),
+            help='LLM engine (default: local, env: SIA_LLM_ENGINE)'
+        )
+        parser.add_argument(
+            '--hf-api-token',
+            type=str,
+            default=os.getenv('SIA_HF_API_TOKEN'),
+            help='Hugging Face access token (env: SIA_HF_API_TOKEN)'
+        )
+        parser.add_argument(
+            '--model',
+            type=str,
+            default=os.getenv('SIA_MODEL', '/root/model/'),
+            help='Path to the model directory (default: /root/model/, env: SIA_MODEL)'
+        )
+        self.args = parser.parse_args()
+    
+    def _parse_bool_env(self, env_var: str, default: bool) -> bool:
+        """Parse boolean environment variable."""
+        val = os.getenv(env_var)
+        if val is None:
+            return default
+        return val.lower() in ('true', '1', 'yes', 'on')
+    
+    def _parse_optional_path(self, env_var: str, default: Optional[Path]) -> Optional[Path]:
+        """Parse optional Path environment variable."""
+        val = os.getenv(env_var)
+        if val is None:
+            return default
+        return Path(val)
+
+    @property
+    def system_prompt(self) -> Path:
+        """Path to the system prompt file."""
+        return self.args.system_prompt
+    
+    @property
+    def action_schema(self) -> Path:
+        """Path to the action schema file."""
+        return self.args.action_schema
+
+    @property
+    def server(self) -> bool:
+        """Enable web server for debugging and human feedback."""
+        return self.args.server
+
+    @property
+    def host(self) -> str:
+        """Web server host."""
+        return self.args.host
+    
+    @property
+    def port(self) -> int:
+        """Web server port."""
+        return self.args.port
+    
+    @property
+    def static_files(self) -> Path:
+        """Path to static web files."""
+        return self.args.static_files
+
+    @property
+    def llm_engine(self) -> str:
+        """LLM engine."""
+        return self.args.llm_engine
+    
+    @property
+    def hf_api_token(self) -> Optional[str]:
+        """Hugging Face access token."""
+        return self.args.hf_api_token
+
+    @property
+    def model(self) -> str:
+        """Path to the model directory."""
+        return self.args.model

sia/hf_llm_engine.py

@@ -0,0 +1,71 @@
+from typing import Iterator, Optional
+from huggingface_hub import InferenceClient
+
+from .llm_engine import LlmEngine
+
+class HfLlmEngine(LlmEngine):
+    """
+    LLM Engine implementation using HuggingFace's InferenceClient.
+    """
+    
+    def __init__(
+        self,
+        model_id: str = "mistralai/Mistral-7B-Instruct-v0.2",
+        api_token: Optional[str] = None,
+        temperature: float = 0.7,
+        max_new_tokens: int = 1024,
+    ):
+        """
+        Initialize the HuggingFace Inference API LLM Engine.
+
+        Args:
+            model_id: HuggingFace model ID to use (default: Mistral-7B-Instruct)
+            api_token: HuggingFace API token. If None, will try to read from HF_TOKEN env var
+            temperature: Sampling temperature (default: 0.7)
+            max_new_tokens: Maximum number of tokens to generate (default: 1024)
+        """
+        self.model_id = model_id
+        self.client = InferenceClient(token=api_token)
+        
+        # Generation parameters
+        self.temperature = temperature
+        self.max_new_tokens = max_new_tokens
+            
+    def set_model_path(self, model_id: str):
+        """
+        Update the model being used.
+
+        Args:
+            model_id: New HuggingFace model ID to use
+        """
+        self.model_id = model_id
+                        
+    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.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        
+        def stream_wrapper():
+            stream = self.client.chat_completion(
+                model=self.model_id,
+                messages=messages,
+                temperature=self.temperature,
+                max_tokens=self.max_new_tokens,
+                stream=True
+            )
+            
+            for response in stream:
+                if content := response.choices[0].delta.content:
+                    yield content
+        return stream_wrapper()

sia/llm_engine.py

@@ -1,78 +1,8 @@
-from threading import Thread
 from typing import Iterator
+from abc import ABC, abstractmethod
 
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
+class LlmEngine(ABC):
-import torch
-
-from . import util
-
-class LlmEngine:
-    def __init__(self, model_path: str):
-        """
-        Initialize the LLM Engine with a model path.
-
-        Args:
-            model_path: Path to the model weights to be used.
-        """
-        self.set_model_path(model_path)
-
-    def set_model_path(self, model_path: str):
-        """
-        Load the model from the specified path.
-
-        Args:
-            model_path: Path to the model weights to load.
-        """
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            return_dict=True,
-            low_cpu_mem_usage=True,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            trust_remote_code=True,
-        )
-        if self.tokenizer.pad_token_id is None:
-            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
-        if model.config.pad_token_id is None:
-            model.config.pad_token_id = model.config.eos_token_id
-        self.pipeline = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=self.tokenizer,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            return_full_text=False,
-        )
 
+    @abstractmethod
     def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
-        """
+        pass
-        Run inference using the system prompt and main context, while validating actions against the provided XML schema.
-
-        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.
-        """
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": main_context}
-        ]
-        prompt = self.tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        streamer = TextIteratorStreamer(
-            self.tokenizer,
-            skip_prompt=True
-        )
-        pipeline_kwargs = dict(
-            text_inputs=prompt,
-            do_sample=True,
-            max_new_tokens=1024,
-            streamer=streamer
-        )
-        thread = Thread(target=self.pipeline, kwargs=pipeline_kwargs)
-        thread.start()
-        return util.stop_before_value(streamer, '<|eot_id|>')

sia/local_llm_engine.py

@@ -0,0 +1,79 @@
+from threading import Thread
+from typing import Iterator
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
+import torch
+
+from . import util
+from .llm_engine import LlmEngine
+
+class LocalLlmEngine(LlmEngine):
+    def __init__(self, model_path: str):
+        """
+        Initialize the LLM Engine with a model path.
+
+        Args:
+            model_path: Path to the model weights to be used.
+        """
+        self.set_model_path(model_path)
+
+    def set_model_path(self, model_path: str):
+        """
+        Load the model from the specified path.
+
+        Args:
+            model_path: Path to the model weights to load.
+        """
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            return_dict=True,
+            low_cpu_mem_usage=True,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        if model.config.pad_token_id is None:
+            model.config.pad_token_id = model.config.eos_token_id
+        self.pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=self.tokenizer,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            return_full_text=False,
+        )
+
+    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.
+
+        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.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            skip_prompt=True
+        )
+        pipeline_kwargs = dict(
+            text_inputs=prompt,
+            do_sample=True,
+            max_new_tokens=1024,
+            streamer=streamer
+        )
+        thread = Thread(target=self.pipeline, kwargs=pipeline_kwargs)
+        thread.start()
+        return util.stop_before_value(streamer, '<|eot_id|>')

system_prompt.md

@@ -4,9 +4,10 @@ You can solve any problem.
 
 Each iteration, the context is updated with the result of your previous actions.
 You modify the context by issuing a commands using XML.
-Always respond with one action adhering to the XML schema.
 Parameters and scripts may be long and complex.
 Use correct XML escaping or CDATA sections.
+It is very important that you always respond with one action adhering to the XML schema!
+Do not respond with anything else after the action.
 
 # Context
 The context has a limited length.

test/local_llm_engine_test.py

@@ -5,18 +5,19 @@ from itertools import tee
 from . import test_data
 
 from sia.llm_engine import LlmEngine
+from sia.local_llm_engine import LocalLlmEngine
 
 class LlmEngineTest(unittest.TestCase):
     def setUp(self):
         self.model_path = "/root/model"
 
     def test_initialization(self):
-        llm_engine = LlmEngine(self.model_path)
+        llm_engine = LocalLlmEngine(self.model_path)
         self.assertIsInstance(llm_engine, LlmEngine)
 
     def test_infer(self):
         main_context = "This is a test"
-        llm_engine = LlmEngine(self.model_path)
+        llm_engine = LocalLlmEngine(self.model_path)
         tokens = llm_engine.infer(test_data.echo_system_prompt, main_context)
         print_tokens, result_tokens = tee(tokens)
         for token in print_tokens: