llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Begin work on agent core

Browse Source
This commit is contained in:
geens
2024-10-27 10:35:05 +01:00
parent a110140d5b
commit d922c12556
9 changed files with 330 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

148
action_schema.xsd Normal file
View File

@@ -0,0 +1,148 @@
<?xml version="1.0" encoding="UTF-8"?>
<xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema">
<!-- Root element containing a sequence of actions to be executed -->
<xs:element name="actions">
<xs:complexType>
<xs:choice minOccurs="0" maxOccurs="unbounded">
<xs:element ref="start_container"/>
<xs:element ref="write_container_stdin"/>
<xs:element ref="read_container_stdout"/>
<xs:element ref="read_container_stderr"/>
<xs:element ref="wait_container"/>
<xs:element ref="set_model_path"/>
</xs:choice>
</xs:complexType>
</xs:element>
<!--
Start a new Docker container.
Containers can be started in two modes:
- Short-running containers:
These are started with a timeout.
Their output is visible in the next iteration.
- Long-running containers:
These are started with a name.
They are added to the container list.
Their output is queried using the read_container_stdout and read_container_stderr actions.
Either a name or a timeout must be specified.
-->
<xs:element name="start_container">
<xs:complexType>
<!-- Docker image to use -->
<xs:attribute name="image" type="xs:string" use="required"/>
<!-- Container name for long-running containers -->
<xs:attribute name="name" type="xs:string" use="optional"/>
<!-- Timeout in milliseconds for short-running containers -->
<xs:attribute name="timeout" type="xs:integer" use="optional"/>
<xs:sequence>
<!-- Main command to run in container -->
<xs:element name="command" type="xs:string" minOccurs="0"/>
<!-- Command line arguments to the main command -->
<xs:element name="argument" type="xs:string" minOccurs="0" maxOccurs="unbounded"/>
<!-- Volume mappings from host to container -->
<xs:element name="volumes" minOccurs="0">
<xs:complexType>
<xs:sequence>
<!--
Each volume maps host path to container path.
Syntax similar to -v switch of docker run command.
using : separator and optional :ro flag.
-->
<xs:element name="volume" type="xs:string" maxOccurs="unbounded"/>
</xs:sequence>
</xs:complexType>
</xs:element>
<!-- Port mappings from host to container -->
<xs:element name="ports" minOccurs="0">
<xs:complexType>
<xs:sequence>
<!-- Each port maps a host port to a container port -->
<xs:element name="port" maxOccurs="unbounded">
<xs:complexType>
<xs:simpleContent>
<xs:extension base="xs:string">
<xs:attribute name="host" type="xs:string" use="required"/>
<xs:attribute name="container" type="xs:string" use="required"/>
</xs:extension>
</xs:simpleContent>
</xs:complexType>
</xs:element>
</xs:sequence>
</xs:complexType>
</xs:element>
<!-- Environment variables to set in the container -->
<xs:element name="environment" minOccurs="0">
<xs:complexType>
<xs:sequence>
<!-- Each variable has a name and value -->
<xs:element name="variable" maxOccurs="unbounded">
<xs:complexType>
<xs:simpleContent>
<xs:extension base="xs:string">
<xs:attribute name="name" type="xs:string" use="required"/>
</xs:extension>
</xs:simpleContent>
</xs:complexType>
</xs:element>
</xs:sequence>
</xs:complexType>
</xs:element>
</xs:sequence>
</xs:complexType>
</xs:element>
<!-- Write data to a container's standard input -->
<xs:element name="write_container_stdin">
<xs:complexType>
<xs:simpleContent>
<xs:extension base="xs:string">
<!-- Name of the target container -->
<xs:attribute name="container" type="xs:string" use="required"/>
</xs:extension>
</xs:simpleContent>
</xs:complexType>
</xs:element>
<!-- Read from a container's standard output buffer -->
<xs:element name="read_container_stdout">
<xs:complexType>
<!-- Name of the container to read from -->
<xs:attribute name="container" type="xs:string" use="required"/>
<!-- Optional number of bytes to read (-1 for all available) -->
<xs:attribute name="n" type="xs:integer" use="optional"/>
</xs:complexType>
</xs:element>
<!-- Read from a container's standard error buffer -->
<xs:element name="read_container_stderr">
<xs:complexType>
<!-- Name of the container to read from -->
<xs:attribute name="container" type="xs:string" use="required"/>
<!-- Optional number of bytes to read (-1 for all available) -->
<xs:attribute name="n" type="xs:integer" use="optional"/>
</xs:complexType>
</xs:element>
<!--
Wait for a container to finish execution.
Standard output and standard error are merged and visible in the next iteration.
-->
<xs:element name="wait_container">
<xs:complexType>
<!-- Name of the container to wait for -->
<xs:attribute name="name" type="xs:string" use="required"/>
<!-- Maximum time to wait in milliseconds -->
<xs:attribute name="timeout" type="xs:integer" use="required"/>
</xs:complexType>
</xs:element>
<!-- Switch to a different LLM model file -->
<xs:element name="set_model_path">
<xs:complexType>
<xs:simpleContent>
<!-- Path to the model file to load -->
<xs:extension base="xs:string"/>
</xs:simpleContent>
</xs:complexType>
</xs:element>
</xs:schema>

83
readme.md
View File

@@ -424,15 +424,16 @@ Module: Docker
#### Function declaration
```python
def start_container(
self,
image: str,
name: str = None,
name: Optional[str] = None,
timeout: int = -1,
command: str = None,
arguments: List[str] = None,
volumes: Dict[str, str] = None, # host_path: container_path
ports: Dict[str, str] = None, # host_port: container_port
environment: Dict[str, str] = None, # key: value
) -> None:
command: Optional[str] = None,
arguments: Optional[List[str]] = None,
volumes: Optional[Dict[str, str]] = None,
ports: Optional[Dict[str, str]] = None,
environment: Optional[Dict[str, str]] = None,
) -> Optional[str]:
"""Start a new Docker container with the specified configuration.
Args:
@@ -447,6 +448,10 @@ def start_container(
name or timeout must be provided.
Returns:
For short-lived containers (with timeout): Container output
For long-running containers: None
Example:
start_container(
image="busybox:latest",
@@ -558,45 +563,6 @@ Standard output and standard error are merged and added as `CDATA` text node.
For long running containers, no information is added.
The container is started and represented in the containers section of the main context.
### Stop container
Module: Docker
#### Function declaration
```python
def stop_container(name: str) -> None:
"""Stop a Docker container by name.
Args:
name: Name of the container to stop
"""
pass
```
#### Schema
```xml
<?xml version="1.0" encoding="UTF-8"?>
<xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema">
<xs:element name="stop_container">
<xs:complexType>
<xs:simpleContent>
<xs:extension base="xs:string"/>
</xs:simpleContent>
</xs:complexType>
</xs:element>
</xs:schema>
```
#### Example
```xml
<stop_container>my-long-running-container</stop_container>
```
#### Results
No information is added.
The container will be removed from the containers section of the main context.
### Write to container standard input
Module: Docker
@@ -728,12 +694,15 @@ Module: Docker
#### Function declaration
```python
def wait_container(name: str, timeout: int):
def wait_container(self, name: str, timeout: int) -> Tuple[int, str]:
"""Wait for a container to finish execution.
Args:
name: Name of the container to wait for
timeout: Time to wait in milliseconds
Returns:
Tuple of (exit_code, output)
"""
pass
```
@@ -744,11 +713,8 @@ def wait_container(name: str, timeout: int):
<xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema">
<xs:element name="wait_container">
<xs:complexType>
<xs:simpleContent>
<xs:extension base="xs:string">
<xs:attribute name="timeout" type="xs:integer" use="required"/>
</xs:extension>
</xs:simpleContent>
<xs:attribute name="name" type="xs:string" use="required"/>
<xs:attribute name="timeout" type="xs:integer" use="required"/>
</xs:complexType>
</xs:element>
</xs:schema>
@@ -756,7 +722,7 @@ def wait_container(name: str, timeout: int):
#### Example
```xml
<wait_container timeout="5000">background-task</wait_container>
<wait_container name="background-task" timeout="5000"/>
```
#### Results
@@ -769,7 +735,7 @@ Module: Reinforcement Learning
#### Function declaration
```python
def select_llm(path: str) -> None:
def set_model_path(path: str) -> None:
"""Switch to a different LLM model file.
Args:
@@ -782,7 +748,7 @@ def select_llm(path: str) -> None:
```xml
<?xml version="1.0" encoding="UTF-8"?>
<xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema">
<xs:element name="select_llm">
<xs:element name="set_model_path">
<xs:complexType>
<xs:simpleContent>
<xs:extension base="xs:string"/>
@@ -794,7 +760,7 @@ def select_llm(path: str) -> None:
#### Example
```xml
<select_llm>/models/2024_10_19_15_03_52</select_llm>
<set_model_path>/models/2024_10_19_15_03_52</set_model_path>
```
#### Results
@@ -905,10 +871,13 @@ Location: Somewhere in Belgium
John did not specify an exact time. I'll suggest 9am. He also did not specify how to be reminded. I'll ask but if he doesn't respond I'll assume a text message on standard output is fine. I'll write down this task in a file so I can keep it in context. I can write simple files with busybox:latest and echo but I will need to use sh -c to do the redirect.
<actions>
<write_stdout message="I'll remind you to feed the cat tomorrow morning at 9am. Is a message on the standard output ok?"/>
<start_container image="busybox:latest" timeout="1000" volumes="/tasks:/tasks">
<start_container image="busybox:latest" timeout="1000">
<command>sh</command>
<argtument>-c</argument>
<argument><![CDATA[echo 'Remind John to feed the cat on 2024-10-18T09:00:00+02:00. Use standard output.' > /tasks/reminder.txt]]></argument>
<volumes>
<volume>/tasks:/tasks</volume>
</volumes>
</start_container>
<monitor_file path="/tasks/reminder.txt"/>
</actions>

3
run.sh
View File

@@ -20,8 +20,7 @@ docker run \
--gpus=all \
--privileged \
-v /$(pwd)/model/:/root/model/ \
$TAG \
-c "bash"
$TAG
# Clean up image
[ ! -z "$TAG" ] && docker rmi $TAG

22
sia/__main__.py
View File

@@ -1,5 +1,25 @@
from .agent_core import AgentCore
from .llm_engine import LlmEngine
from .docker_module import DockerModule
def main():
print("Hello, World! --sia")
"""Main entry point for the SIA application."""
system_prompt_path = "system_prompt.txt"
action_schema_path = "action_schema.xsd"
model_path = "/root/model"
with open(system_prompt_path, 'r') as f:
system_prompt = f.read()
with open(action_schema_path, 'r') as f:
action_schema = f.read()
agent_core = AgentCore(
system_prompt=f"{system_prompt}{action_schema}",
action_schema=action_schema,
docker_module=DockerModule(),
llm_engine=LlmEngine(model_path=model_path)
)
agent_core.run_iteration()
if __name__ == "__main__":
main()

51
sia/agent_core.py Normal file
View File

@@ -0,0 +1,51 @@
from itertools import tee
from typing import Optional, Dict, List
from .docker_module import DockerModule
from .llm_engine import LlmEngine
from .context_template import generate_context
from .util import get_valid_root_elements, split_response
class AgentCore:
"""
Core orchestration class for SIA that manages the interaction between
different modules and runs the main agent loop.
"""
def __init__(
self,
system_prompt: str,
action_schema: str,
docker_module: DockerModule,
llm_engine: LlmEngine
):
"""
Initialize the AgentCore with required components.
Args:
system_prompt: System prompt to use for the LLM
action_schema_path: Path to the XML schema defining valid actions
docker_module: DockerModule instance
llm_engine: LLmEngine instance
"""
self.system_prompt = system_prompt
self.action_schema = action_schema
self.docker_module = docker_module
self.llm_engine = llm_engine
self.valid_elements = get_valid_root_elements(self.action_schema)
def run_iteration(self) -> None:
"""Run a single iteration of the main agent loop."""
containers = self.docker_module.get_all_container_statuses()
context = generate_context(containers)
tokens = self.llm_engine.infer(
self.system_prompt,
context
)
print_tokens, response_tokens = tee(tokens)
for token in print_tokens:
print(token, end="", flush=True)
response_tokens = ''.join(response_tokens)
result = split_response(response_tokens, self.valid_elements)
print(f"result: {result}")

15
sia/docker_module.py
View File

@@ -80,7 +80,7 @@ class DockerModule:
volumes: Optional[Dict[str, str]] = None,
ports: Optional[Dict[str, str]] = None,
environment: Optional[Dict[str, str]] = None,
) -> str:
) -> Optional[str]:
"""Start a new Docker container with the specified configuration.
Args:
@@ -95,11 +95,7 @@ class DockerModule:
Returns:
For short-lived containers (with timeout): Container output
For long-running containers: Container name
Raises:
docker.errors.ImageNotFound: If specified image doesn't exist
docker.errors.APIError: If container creation/start fails
For long-running containers: None
"""
# Validate inputs
if name is None and timeout < 0:
@@ -143,7 +139,6 @@ class DockerModule:
else:
# For long-running containers
self.containers[name] = (container, socket)
return name
def write_container_stdin(self, name: str, data: str) -> None:
"""Write data to a container's standard input.
@@ -225,11 +220,6 @@ class DockerModule:
Returns:
Tuple of (exit_code, output)
Raises:
KeyError: If container name not found
docker.errors.APIError: If wait operation fails
TimeoutError: If container doesn't finish within timeout
"""
if name not in self.containers:
raise KeyError(f"Container {name} not found")
@@ -278,3 +268,4 @@ class DockerModule:
statuses = []
for name, (container, socket) in self.containers.items():
statuses.append(self.get_container_status(name))
return statuses

5
sia/llm_engine.py
View File

@@ -1,4 +1,5 @@
from threading import Thread
from typing import Iterator
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
import torch
@@ -44,7 +45,7 @@ class LlmEngine:
return_full_text=False,
)
def infer(self, system_prompt: str, main_context: str) -> TextIteratorStreamer:
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.
@@ -53,7 +54,7 @@ class LlmEngine:
main_context: The main context string after templating
Returns:
TextIteratorStreamer: An iterator that yields the generated text.
Iterator[str]: An iterator that yields the generated text.
"""
messages = [
{"role": "system", "content": system_prompt},

2
sia/util.py
View File

@@ -1,4 +1,4 @@
from typing import Iterator, TypeVar
from typing import Iterator
import re
import xml.etree.ElementTree as ET

75
system_prompt.txt Normal file
View File

@@ -0,0 +1,75 @@
You are SIA, the self improving agent.
Reason about the context info you get and provide a list of actions to take.
The actions are formatted as XML.
The closing tag is your last token.
Don't add a markdown code block around the actions.
An example iteration looks like this:
Context
```xml
<context>
<system
time="2024-10-18T12:00:00Z"
cpu="12"
gpu="26"
memory_used="9556302234"
memory_total="17179869184"
disk_used="244434939904"
disk_total="273145991168"
context="3"
stdin="0"
/>
<containers/>
<previous>
<reasoning>
There is data available on the standard input channel. I should read it. I have no other running tasks to tend to.
</reasoning>
<actions>
<read_stdin n="42">
<![CDATA[Remind me to feed the cat tomorrow morning]]>
</read_stdin>
</actions>
</previous>
<files>
<file name="/" type="dir" index="0">
<![CDATA[
drwxr-xr-x 1 sia 197121 0 2024-10-16 23:02:16.486152500 +0200 tasks/
drwxr-xr-x 1 sia 197121 0 2024-10-16 22:35:31.806079500 +0200 user/
]]>
</file>
<file name="/tasks" type="dir" index="1">
</file>
<file name="/user" type="dir">
<![CDATA[
-rw-r--r-- 1 sia 197121 71 2024-10-16 22:41:23.223580300 +0200 general_info.txt
]]>
</file>
<file name="/user/general_info.txt" type="file" index="2">
<![CDATA[
Name: John (I don't know his last name)
Location: Somewhere in Belgium
]]>
</file>
</files>
</context>
```
Response
```xml
John did not specify an exact time. I'll suggest 9am. He also did not specify how to be reminded. I'll ask but if he doesn't respond I'll assume a text message on standard output is fine. I'll write down this task in a file so I can keep it in context. I can write simple files with busybox:latest and echo but I will need to use sh -c to do the redirect.
<actions>
<write_stdout message="I'll remind you to feed the cat tomorrow morning at 9am. Is a message on the standard output ok?"/>
<start_container image="busybox:latest" timeout="1000">
<command>sh</command>
<argtument>-c</argument>
<argument><![CDATA[echo 'Remind John to feed the cat on 2024-10-18T09:00:00+02:00. Use standard output.' > /tasks/reminder.txt]]></argument>
<volumes>
<volume>/tasks:/tasks</volume>
</volumes>
</start_container>
<monitor_file path="/tasks/reminder.txt"/>
</actions>
```
These actions are available to you: