diff --git a/.dockerignore b/.dockerignore
index e3f10e1..e346f21 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1 +1,2 @@
+.env
./model/
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 4591654..0b5394f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+.env
pdf/
model/
claude.txt
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
index a865d42..e0ff00c 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -32,4 +32,4 @@ COPY ./ /root/sia/
COPY --from=web-build /app/dist /root/sia/static/
WORKDIR /root/sia
-CMD ["python3", "-m", "sia"]
\ No newline at end of file
+ENTRYPOINT ["python3", "-m", "sia"]
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 86d524c..07d56b5 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
accelerate
aiohttp
bs4
+python-dotenv
torch
transformers
\ No newline at end of file
diff --git a/run.sh b/run.sh
index cd400c7..71551ae 100755
--- a/run.sh
+++ b/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 "$@"
diff --git a/sia/__main__.py b/sia/__main__.py
index e8d2fc7..9210aa5 100644
--- a/sia/__main__.py
+++ b/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._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._config = Config()
- self._llm = LlmEngine("/root/model")
- #self._llm = TestLLM()
+ self._system_prompt = self._config.system_prompt.read_text()
+ 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("/assets/", self._static_dir / "assets", show_index=False)
+ self._app.router.add_static("/static/", self._config.static_files, 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()
diff --git a/sia/config.py b/sia/config.py
new file mode 100644
index 0000000..6755a14
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/sia/hf_llm_engine.py b/sia/hf_llm_engine.py
new file mode 100644
index 0000000..3182030
--- /dev/null
+++ b/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()
\ No newline at end of file
diff --git a/sia/llm_engine.py b/sia/llm_engine.py
index d739266..522c790 100644
--- a/sia/llm_engine.py
+++ b/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
-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,
- )
+class LlmEngine(ABC):
+ @abstractmethod
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|>')
\ No newline at end of file
+ pass
\ No newline at end of file
diff --git a/sia/local_llm_engine.py b/sia/local_llm_engine.py
new file mode 100644
index 0000000..f3997d2
--- /dev/null
+++ b/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|>')
diff --git a/system_prompt.md b/system_prompt.md
index d1b2831..684da32 100644
--- a/system_prompt.md
+++ b/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.
diff --git a/test/llm_engine_test.py b/test/local_llm_engine_test.py
similarity index 79%
rename from test/llm_engine_test.py
rename to test/local_llm_engine_test.py
index a11b9c5..f78fce8 100644
--- a/test/llm_engine_test.py
+++ b/test/local_llm_engine_test.py
@@ -5,21 +5,22 @@ 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:
print(token, end="", flush=True)
result = ''.join(result_tokens)
- self.assertEqual(result, f"{main_context}{main_context}")
\ No newline at end of file
+ self.assertEqual(result, f"{main_context}{main_context}")