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

New web interface, move llm engine to separate process

Browse Source
This commit is contained in:
geens
2025-05-20 09:43:17 +02:00
parent 895a533e01
commit d4a4902b94
137 changed files with 4850 additions and 3503 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
sia/__main__.py
View File

@@ -2,20 +2,16 @@ from aiohttp import web
import asyncio
from .auto_approver import AutoApprover
from .chat_io_buffer import ChatIOBuffer
from .config import Config
from .iteration_logger import IterationLogger
from .llm_engine.hf_llm_engine import HfLlmEngine
from .llm_engine.local_llm_engine import LocalLlmEngine
from .llm_engine.mistral_llm_engine import MistralLlmEngine
from .llm_engine.openai_llm_engine import OpenAILlmEngine
from .llm_engine.qwq_llm_engine import QwQLlmEngine
from .llm_engine import LlmEngine
from .response_parser import ResponseParser
from .system_metrics import SystemMetrics
from .web.api import Api
from .web.static import Static
from .web.websockets import Websockets
from .web_agent import WebAgent
from .web_io_buffer import WebIOBuffer
from .working_memory import WorkingMemory
class Main:
@@ -30,49 +26,17 @@ class Main:
# Initialize LLM engines based on config
self._llms = {}
if config.local_enabled:
self._llms['local'] = LocalLlmEngine(
config.local_model,
config.local_temperature,
config.local_token_limit,
self._action_schema,
config.local_api_key,
)
if config.openai_enabled:
self._llms['openai'] = OpenAILlmEngine(
config.openai_model,
config.openai_temperature,
config.openai_token_limit,
config.openai_api_key,
)
if config.hf_enabled:
self._llms['hf'] = HfLlmEngine(
config.hf_model,
config.hf_temperature,
config.hf_api_key,
)
if config.mistral_enabled:
self._llms['mistral'] = MistralLlmEngine(
config.mistral_model,
config.mistral_temperature,
config.mistral_token_limit,
config.mistral_api_key,
)
if config.qwq_enabled:
self._llms['qwq'] = QwQLlmEngine(
config.qwq_model,
config.qwq_temperature,
self._action_schema,
# Use the config.llms property which returns only enabled LLMs
for llm_name, executable_path in config.llms.items():
self._llms[llm_name] = LlmEngine(
executable_path=executable_path,
action_schema_path=str(config.action_schema)
)
if not self._llms:
raise ValueError("No LLM engines enabled in configuration")
self._io_buffer = WebIOBuffer()
self._io_buffer = ChatIOBuffer()
self._working_memory = WorkingMemory()
self._agent = WebAgent(
system_prompt=self._system_prompt,

172
sia/auto_approver.py
View File

@@ -1,13 +1,14 @@
from dataclasses import dataclass
from threading import Thread, Event
from typing import Callable, TypeAlias, TypedDict
from .web_agent import WebAgent, LlmState
from typing import Callable, TypeAlias
from .web_agent import WebAgent, AgentState
class AutoApproverConfig(TypedDict):
@dataclass
class AutoApproverConfig:
context_enabled: bool
response_enabled: bool
context_timeout: float
response_timeout: float
llm_name: str
ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]
@@ -21,48 +22,47 @@ class AutoApprover:
Initialize auto approver with a WebAgent instance.
"""
self.agent = agent
self._llm_name = next(iter(agent.llms.keys()))
self._context_timeout = 5.0
self._response_timeout = 10.0
self._context_enabled = False
self._response_enabled = False
self._config = AutoApproverConfig(
context_enabled=False,
response_enabled=False,
context_timeout=1.0,
response_timeout=1.0,
)
self._stop_event = Event()
self._context_thread: Thread | None = None
self._response_thread: Thread | None = None
self._config_change_handlers: list[ConfigChangeHandler] = []
self.agent.add_llm_change_handler(self._handle_llm_state_change)
self.agent.add_context_change_handler(self._handle_context_change)
self._previous_state = agent.state
self.agent.add_state_change_handler(self._handle_state_change)
@property
def config(self) -> AutoApproverConfig:
return AutoApproverConfig(
context_enabled=self._context_enabled,
response_enabled=self._response_enabled,
context_timeout=self._context_timeout,
response_timeout=self._response_timeout,
llm_name=self._llm_name
)
return self._config
def set_config(self, config: AutoApproverConfig) -> None:
if config['llm_name'] not in self.agent.llms:
raise ValueError(f"Unknown LLM: {config['llm_name']}")
notify_config_change = self._notify_config_change
self._notify_config_change = lambda: None
try:
self.context_enabled = False
self.response_enabled = False
self.context_timeout = config['context_timeout']
self.response_timeout = config['response_timeout']
self.llm_name = config['llm_name']
self.context_enabled = config['context_enabled']
self.response_enabled = config['response_enabled']
finally:
self._notify_config_change = notify_config_change
@config.setter
def config(self, config: AutoApproverConfig) -> None:
if config == self._config:
return
if config.context_timeout != self._config.context_timeout:
if config.context_enabled and self._config.context_enabled:
raise ValueError("Cannot change context timeout while enabled")
if config.response_timeout != self._config.response_timeout:
if config.response_enabled and self._config.response_enabled:
raise ValueError("Cannot change response timeout while enabled")
self._stop_context_thread()
self._stop_response_thread()
# If both context and response are enabled, start response approval only.
# Better to run inference with a full prefix than approve an emmpty response.
if config.response_enabled and not self._config.response_enabled:
# response approval was just enabled
self._config = config
self._start_response_thread()
elif config.context_enabled and not self._config.context_enabled:
# context approval was just enabled
self._config = config
self._start_context_thread()
self._config = config
self._notify_config_change()
def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
@@ -73,81 +73,17 @@ class AutoApprover:
for handler in self._config_change_handlers:
handler(current_config)
@property
def context_timeout(self) -> float:
return self._context_timeout
@context_timeout.setter
def context_timeout(self, timeout: float) -> None:
if self._context_enabled:
raise ValueError("Cannot change timeout while auto-approval is enabled")
self._context_timeout = timeout
self._notify_config_change()
@property
def response_timeout(self) -> float:
return self._response_timeout
@response_timeout.setter
def response_timeout(self, timeout: float) -> None:
if self._response_enabled:
raise ValueError("Cannot change timeout while auto-approval is enabled")
self._response_timeout = timeout
self._notify_config_change()
@property
def context_enabled(self) -> bool:
return self._context_enabled
@context_enabled.setter
def context_enabled(self, enabled: bool) -> None:
if enabled == self._context_enabled:
return
self._context_enabled = enabled
self._stop_context_thread()
if enabled and self.agent.llms[self._llm_name] == LlmState.IDLE:
self._start_context_thread()
self._notify_config_change()
@property
def response_enabled(self) -> bool:
return self._response_enabled
@response_enabled.setter
def response_enabled(self, enabled: bool) -> None:
if enabled == self._response_enabled:
return
self._response_enabled = enabled
self._stop_response_thread()
if enabled and self.agent.llms[self._llm_name] == LlmState.IDLE:
self._start_response_thread()
self._notify_config_change()
@property
def llm_name(self) -> str:
return self._llm_name
@llm_name.setter
def llm_name(self, name: str) -> None:
if name not in self.agent.llms:
raise ValueError(f"Unknown LLM: {name}")
self._llm_name = name
self._notify_config_change()
def _handle_llm_state_change(self, llm_name: str, state: LlmState) -> None:
if llm_name != self._llm_name:
return
if state == LlmState.IDLE and self._response_enabled:
self._start_response_thread()
else:
self._stop_response_thread()
def _handle_context_change(self, context: str, generated: bool) -> None:
if generated and self._context_enabled:
self._start_context_thread()
else:
self._stop_context_thread()
def _handle_state_change(self, state: AgentState) -> None:
if state == AgentState.IDLE:
if self._previous_state == AgentState.INFERENCE:
# finished inference
if self._config.response_enabled:
self._start_response_thread()
elif self._previous_state == AgentState.PROCESSING_RESPONSE:
# finished processing response
if self._config.context_enabled:
self._start_context_thread()
self._previous_state = state
def _stop_context_thread(self) -> None:
if self._context_thread:
@@ -170,13 +106,13 @@ class AutoApprover:
self._response_thread.start()
def _context_approval_thread(self) -> None:
if self._stop_event.wait(self._context_timeout):
if self._stop_event.wait(self._config.context_timeout):
return
if self._context_enabled:
self.agent.run_inference(self._llm_name)
if self._config.context_enabled:
self.agent.run_inference()
def _response_approval_thread(self) -> None:
if self._stop_event.wait(self._response_timeout):
if self._stop_event.wait(self._config.response_timeout):
return
if self._response_enabled:
self.agent.approve_response()
if self._config.response_enabled:
self.agent.approve_response()

14
sia/base_agent.py
View File

@@ -4,7 +4,6 @@ import xml.etree.ElementTree as ET
from .llm_engine import LlmEngine
from .response_parser import ResponseParser
from .system_metrics import SystemMetrics
from .util import pretty_print_element
from .working_memory import WorkingMemory
class BaseAgent(ABC):
@@ -36,8 +35,13 @@ class BaseAgent(ABC):
def system_prompt(self) -> str:
"""Get the system prompt."""
return f"{self._system_prompt}\n{self._action_schema}"
@property
def working_memory(self) -> WorkingMemory:
"""Get the working memory."""
return self._working_memory
def _compile_context(self, llmEngine: LlmEngine) -> str:
def _compile_context(self, llmEngine: LlmEngine) -> ET:
"""
Compile the current context for LLM inference.
Includes system metrics and working memory entries.
@@ -60,15 +64,13 @@ class BaseAgent(ABC):
for entry in memory_context:
context.append(entry)
context_str = pretty_print_element(context)
# Calculate token usage percentage
token_count = llmEngine.token_count(self.system_prompt, context_str)
token_count = llmEngine.token_count(self.system_prompt, context)
token_limit = llmEngine.token_limit()
context_usage = (float(token_count) / float(token_limit)) * 100.0
# Update context usage metric
context.set("context", f"{str(round(context_usage, 2))}%")
return pretty_print_element(context)
return context

102
sia/chat_io_buffer.py Normal file
View File

@@ -0,0 +1,102 @@
from dataclasses import dataclass
from typing import List, Callable
from enum import Enum
import datetime
from .io_buffer import IOBuffer
from .util import format_timestamp
class MessageType(Enum):
USER = "user"
ASSISTANT = "assistant"
@dataclass
class ChatMessage:
id: str # Formatted timestamp
content: str
message_type: MessageType
class ChatIOBuffer(IOBuffer):
def __init__(self):
self._messages: List[ChatMessage] = []
self._last_read_timestamp: str = format_timestamp(datetime.datetime.now())
self._change_handlers: List[Callable] = []
@property
def messages(self) -> List[ChatMessage]:
return self._messages.copy()
@property
def last_read_timestamp(self) -> str:
return self._last_read_timestamp
def read(self) -> str:
unread_messages = self._unread_messages()
self._last_read_timestamp = self._format_now()
self._notify_handlers([], [])
return unread_messages
def write(self, content: str) -> None:
message = ChatMessage(
id=self._format_now(),
content=content,
message_type=MessageType.ASSISTANT
)
self._messages.append(message)
self._notify_handlers(
new_messages=[message],
deleted_message_ids=[]
)
def buffer_length(self) -> int:
return len(self._unread_messages())
def add_user_message(self, content: str):
message = ChatMessage(
id=self._format_now(),
content=content,
message_type=MessageType.USER
)
self._messages.append(message)
self._notify_handlers(
new_messages=[message],
deleted_message_ids=[]
)
def delete_message(self, message_id: str):
for i, msg in enumerate(self._messages):
if msg.id == message_id:
self._messages.pop(i)
self._notify_handlers(
new_messages=[],
deleted_message_ids=[message_id]
)
break
def clear(self) -> None:
deleted_ids = [msg.id for msg in self._messages]
self._messages.clear()
self._last_read_timestamp = self._format_now()
self._notify_handlers(
new_messages=[],
deleted_message_ids=deleted_ids
)
def add_change_handler(self, handler: Callable[[
List[ChatMessage], # new messages
List[str], # deleted message ids
str # last read timestamp
], None]) -> None:
self._change_handlers.append(handler)
def _notify_handlers(self, new_messages: List[ChatMessage],
deleted_message_ids: List[str]) -> None:
for handler in self._change_handlers:
handler(new_messages, deleted_message_ids, self._last_read_timestamp)
def _unread_messages(self) -> List[ChatMessage]:
return "\n".join(msg.content for msg in self._messages
if msg.message_type == MessageType.USER and msg.id > self._last_read_timestamp)
def _format_now(self) -> str:
return format_timestamp(datetime.datetime.now())

262
sia/config.py
View File

@@ -1,9 +1,10 @@
from dataclasses import dataclass
from dotenv import load_dotenv
from pathlib import Path
from typing import Optional
from typing import Dict
import argparse
import os
import tomli as tomllib
@dataclass
class Config:
@@ -59,170 +60,27 @@ class Config:
parser.add_argument(
'--static-files',
type=Path,
default=self._parse_optional_path('SIA_STATIC_FILES', '/root/static/'),
default=os.getenv('SIA_STATIC_FILES', '/root/static/'),
help='Path to static web files (default: /root/static/, env: SIA_STATIC_FILES)'
)
# Local LLM configuration
# LLM configuration
parser.add_argument(
'--local-enable',
action='store_true',
default=self._parse_bool_env('SIA_LOCAL_ENABLED', False),
help='Enable local LLM engine (env: SIA_LOCAL_ENABLED)'
)
parser.add_argument(
'--local-model',
type=str,
default=os.getenv('SIA_LOCAL_MODEL', '/root/models/current'),
help='Path to local model directory (default: /root/models/current, env: SIA_LOCAL_MODEL)'
)
parser.add_argument(
'--local-temperature',
type=float,
default=float(os.getenv('SIA_LOCAL_TEMPERATURE', '0.1')),
help='Local LLM temperature (default: 0.1, env: SIA_LOCAL_TEMPERATURE)'
)
parser.add_argument(
'--local-token-limit',
type=int,
default=int(os.getenv('SIA_LOCAL_TOKEN_LIMIT', '2048')),
help='Local LLM token limit (env: SIA_LOCAL_TOKEN_LIMIT)'
)
parser.add_argument(
'--local-api-key',
type=str,
default=os.getenv('SIA_LOCAL_API_KEY'),
help='API key for local models (env: SIA_LOCAL_API_KEY)'
)
# OpenAI configuration
parser.add_argument(
'--openai-enable',
action='store_true',
default=self._parse_bool_env('SIA_OPENAI_ENABLED', False),
help='Enable OpenAI LLM engine (env: SIA_OPENAI_ENABLED)'
)
parser.add_argument(
'--openai-model',
type=str,
default=os.getenv('SIA_OPENAI_MODEL', 'gpt-4o'),
help='OpenAI model name (default: gpt-4o, env: SIA_OPENAI_MODEL)'
)
parser.add_argument(
'--openai-temperature',
type=float,
default=float(os.getenv('SIA_OPENAI_TEMPERATURE', '0.1')),
help='OpenAI temperature (default: 0.1, env: SIA_OPENAI_TEMPERATURE)'
)
parser.add_argument(
'--openai-token-limit',
type=int,
default=int(os.getenv('SIA_OPENAI_TOKEN_LIMIT', '4096')),
help='OpenAI token limit (env: SIA_OPENAI_TOKEN_LIMIT)'
)
parser.add_argument(
'--openai-api-key',
type=str,
default=os.getenv('SIA_OPENAI_API_KEY'),
help='OpenAI API key (env: SIA_OPENAI_API_KEY)'
)
# Hugging Face configuration
parser.add_argument(
'--hf-enable',
action='store_true',
default=self._parse_bool_env('SIA_HF_ENABLED', False),
help='Enable Hugging Face LLM engine (env: SIA_HF_ENABLED)'
)
parser.add_argument(
'--hf-model',
type=str,
default=os.getenv('SIA_HF_MODEL'),
help='Hugging Face model name (env: SIA_HF_MODEL)'
)
parser.add_argument(
'--hf-temperature',
type=float,
default=float(os.getenv('SIA_HF_TEMPERATURE', '0.1')),
help='Hugging Face temperature (default: 0.1, env: SIA_HF_TEMPERATURE)'
)
parser.add_argument(
'--hf-api-key',
type=str,
default=os.getenv('SIA_HF_API_KEY'),
help='Hugging Face API key (env: SIA_HF_API_KEY)'
)
# Mistral configuration
parser.add_argument(
'--mistral-enable',
action='store_true',
default=self._parse_bool_env('SIA_MISTRAL_ENABLED', False),
help='Enable Mistral LLM engine (env: SIA_MISTRAL_ENABLED)'
)
parser.add_argument(
'--mistral-model',
type=str,
default=os.getenv('SIA_MISTRAL_MODEL'),
help='Mistral model name (env: SIA_MISTRAL_MODEL)'
)
parser.add_argument(
'--mistral-temperature',
type=float,
default=float(os.getenv('SIA_MISTRAL_TEMPERATURE', '0.1')),
help='Mistral temperature (default: 0.1, env: SIA_MISTRAL_TEMPERATURE)'
)
parser.add_argument(
'--mistral-token-limit',
type=int,
default=int(os.getenv('SIA_MISTRAL_TOKEN_LIMIT', '4096')),
help='Mistral token limit (env: SIA_MISTRAL_TOKEN_LIMIT)'
)
parser.add_argument(
'--mistral-api-key',
type=str,
default=os.getenv('SIA_MISTRAL_API_KEY'),
help='Mistral API key (env: SIA_MISTRAL_API_KEY)'
)
# QwQ configuration
parser.add_argument(
'--qwq-enable',
action='store_true',
default=self._parse_bool_env('SIA_QWQ_ENABLED', False),
help='Enable QwQ LLM engine (env: SIA_QWQ_ENABLED)'
)
parser.add_argument(
'--qwq-model',
type=str,
default=os.getenv('SIA_QWQ_MODEL', '/root/models/current'),
help='Path to QwQ model or HF model ID (default: /root/models/current, env: SIA_QWQ_MODEL)'
)
parser.add_argument(
'--qwq-temperature',
type=float,
default=float(os.getenv('SIA_QWQ_TEMPERATURE', '0.6')),
help='QwQ temperature (default: 0.1, env: SIA_QWQ_TEMPERATURE)'
)
parser.add_argument(
'--qwq-token-limit',
type=int,
default=int(os.getenv('SIA_QWQ_TOKEN_LIMIT', '4096')),
help='QwQ token limit (0 for model default, env: SIA_QWQ_TOKEN_LIMIT)'
'--llm-config',
type=Path,
default=os.getenv('SIA_LLM_CONFIG', '/root/sia/llm_config.toml'),
help='Path to TOML configuration file for LLMs (default: /root/sia/llm_config.toml, env: SIA_LLM_CONFIG)'
)
self.args = parser.parse_args()
with open(self.llm_config, "rb") as f:
self._llm_configs = tomllib.load(f)
def _parse_bool_env(self, env_var: str, default: bool) -> bool:
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]:
val = os.getenv(env_var)
if val is None:
return default
return Path(val)
# Core properties
@property
@@ -258,100 +116,12 @@ class Config:
def static_files(self) -> Path:
return self.args.static_files
# Local LLM properties
# LLM properties
@property
def local_enabled(self) -> bool:
return self.args.local_enable
@property
def local_model(self) -> str:
return self.args.local_model
@property
def local_temperature(self) -> float:
return self.args.local_temperature
@property
def local_token_limit(self) -> int:
return self.args.local_token_limit
def llm_config(self) -> Path:
return self.args.llm_config
@property
def local_api_key(self) -> Optional[str]:
return self.args.local_api_key
# OpenAI properties
@property
def openai_enabled(self) -> bool:
return self.args.openai_enable
@property
def openai_model(self) -> str:
return self.args.openai_model
@property
def openai_temperature(self) -> float:
return self.args.openai_temperature
@property
def openai_token_limit(self) -> int:
return self.args.openai_token_limit
@property
def openai_api_key(self) -> Optional[str]:
return self.args.openai_api_key
# Hugging Face properties
@property
def hf_enabled(self) -> bool:
return self.args.hf_enable
@property
def hf_model(self) -> str:
return self.args.hf_model
@property
def hf_temperature(self) -> float:
return self.args.hf_temperature
@property
def hf_api_key(self) -> Optional[str]:
return self.args.hf_api_key
# Mistral properties
@property
def mistral_enabled(self) -> bool:
return self.args.mistral_enable
@property
def mistral_model(self) -> str:
return self.args.mistral_model
@property
def mistral_temperature(self) -> float:
return self.args.mistral_temperature
@property
def mistral_token_limit(self) -> int:
return self.args.mistral_token_limit
@property
def mistral_api_key(self) -> Optional[str]:
return self.args.mistral_api_key
# QwQ properties
@property
def qwq_enabled(self) -> bool:
return self.args.qwq_enable
@property
def qwq_model(self) -> str:
return self.args.qwq_model
@property
def qwq_temperature(self) -> float:
return self.args.qwq_temperature
@property
def qwq_token_limit(self) -> Optional[int]:
# Return None if 0 to use model default
return self.args.qwq_token_limit if self.args.qwq_token_limit > 0 else None
def llms(self) -> Dict[str, str]:
"""Get only the enabled LLM configurations."""
return self._llm_configs

5
sia/iteration_logger.py
View File

@@ -23,7 +23,7 @@ class IterationLogger:
def log_iteration(
self,
timestamp: datetime,
context: str,
context: ET,
response: str,
):
"""
@@ -41,8 +41,7 @@ class IterationLogger:
root.set("system_prompt_hash", self._system_prompt_hash)
root.set("action_schema_hash", self._action_schema_hash)
context_elem = ET.SubElement(root, "context")
context_elem.text = context
root.append(context)
response_elem = ET.SubElement(root, "response")
response_elem.text = response

212
sia/llm_engine.py Normal file
View File

@@ -0,0 +1,212 @@
from typing import Iterator
from .util import pretty_print_element
import subprocess
import sys
import xml.etree.ElementTree as ET
class LlmEngine:
"""
LlmEngine manages communication with LLM engine subprocesses.
Each LLM type runs in its own subprocess with a tailored environment.
"""
EOT = '\x04' # EOT character (ASCII 4) as bytes
def __init__(self, executable_path: str, action_schema_path: str):
"""
Initialize the LLM engine subprocess.
Args:
executable_path (str): Path to the LLM engine executable
action_schema_path (str): Path to the XML action schema
"""
self.executable_path = executable_path
self.action_schema_path = action_schema_path
self.action_schema = open(action_schema_path, 'r').read()
self.process = None
self.restart()
def _read_until_eot(self) -> str:
"""
Read from subprocess stdout until EOT character.
Returns:
str: Complete response without the EOT character
"""
response = []
while True:
# Read available data
data = self.process.stdout.read(1024) # Read up to 1024 bytes at a time
if not data: # process died
self.restart()
raise RuntimeError(f"LLM subprocess terminated unexpectedly")
data = data.decode('utf-8')
# Check if EOT is in the data
if self.EOT in data:
eot_index = data.index(self.EOT)
response.append(data[:eot_index]) # Add data before EOT
break
else:
response.append(data)
return "".join(response)
def token_limit(self) -> int:
"""
Get the maximum token limit of the LLM.
Returns:
int: Maximum token limit
"""
self.process.stdin.write(b"<token_limit/>\n")
self.process.stdin.flush()
response = self._read_until_eot()
return int(response.strip())
def token_count(self, system_prompt: str, main_context: ET.Element, prefix: str = "") -> int:
"""
Count the number of tokens in the prompt.
Args:
system_prompt (str): System prompt text
main_context (ET.Element): Main context as ElementTree
prefix (str): Optional prefix for continuing generation
Returns:
int: Token count
"""
# Create the XML document
root = ET.Element("token_count")
# Add system prompt
system_prompt_elem = ET.SubElement(root, "system")
system_prompt_elem.text = self._append_action_schema(system_prompt)
# Add context element
context_elem = ET.SubElement(root, "context")
context_elem.text = pretty_print_element(main_context)
# Add prefix if provided
if prefix:
prefix_elem = ET.SubElement(root, "prefix")
prefix_elem.text = prefix
# Send to subprocess - convert to bytes
xml_str = ET.tostring(root, encoding='utf-8')
self.process.stdin.write(xml_str + b"\n")
self.process.stdin.flush()
# Read response
response = self._read_until_eot()
return int(response.strip())
def infer(self, system_prompt: str, main_context: ET.Element, prefix: str = "") -> Iterator[str]:
"""
Generate text from the LLM.
Args:
system_prompt (str): System prompt text
main_context (ET.Element): Main context as ElementTree
prefix (str): Optional prefix for continuing generation
Returns:
Iterator[str]: Generated text, yielded as it's produced
"""
# Create the XML document
root = ET.Element("infer_xml")
# Add action schema path
schema_path_elem = ET.SubElement(root, "schema")
schema_path_elem.text = self.action_schema_path
# Add system prompt in CDATA
system_prompt_elem = ET.SubElement(root, "system")
system_prompt_elem.text = self._append_action_schema(system_prompt)
# Add context element
context_elem = ET.SubElement(root, "context")
context_elem.text = pretty_print_element(main_context)
# Add prefix if provided
if prefix:
prefix_elem = ET.SubElement(root, "prefix")
prefix_elem.text = prefix
# Send to subprocess - convert to bytes
xml_str = ET.tostring(root, encoding='utf-8')
self.process.stdin.write(xml_str + b"\n")
self.process.stdin.flush()
while True:
# Read available data
data = self.process.stdout.read(1024) # Read up to 1024 bytes at a time
if not data: # Process died
self.restart()
raise RuntimeError("LLM subprocess terminated unexpectedly")
data = data.decode('utf-8')
if self.EOT in data:
eot_index = data.index(self.EOT)
if eot_index > 0:
yield data[:eot_index]
break
else:
yield data
def restart(self):
"""Start the LLM engine subprocess."""
# Ensure any existing process is terminated
if self.process:
self._terminate_process()
# Start the subprocess with pipes for stdin/stdout and direct stderr
self.process = subprocess.Popen(
["/bin/bash", "-c", self.executable_path],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=sys.stderr,
text=False, # Use binary mode to avoid buffering and encoding issues
bufsize=0,
)
# Check if the process started successfully
if self.process.poll() is not None:
raise RuntimeError(f"Failed to start LLM engine at {self.executable_path}")
def _terminate_process(self):
"""Terminate the LLM engine subprocess safely."""
if self.process:
try:
self.process.stdin.close()
self.process.stdout.close()
self.process.terminate()
try:
self.process.wait(timeout=5) # Wait for process to terminate
except subprocess.TimeoutExpired:
# Force kill if termination takes too long
self.process.kill()
self.process.wait(timeout=2)
finally:
self.process = None
def _append_action_schema(self, system_prompt: str) -> str:
"""
Append the action schema to the system prompt.
Args:
system_prompt (str): Original system prompt
Returns:
str: Updated system prompt with action schema
"""
return f"{system_prompt}\n\n--- Action Schema ---\n{self.action_schema}"
def __del__(self):
"""Cleanup when the object is destroyed."""
self._terminate_process()

15
sia/llm_engine/__init__.py
View File

@@ -1,15 +0,0 @@
from typing import Callable, Iterator
from abc import ABC, abstractmethod
class LlmEngine(ABC):
@abstractmethod
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
pass
@abstractmethod
def token_count(self, system_prompt: str, main_context: str) -> int:
pass
@abstractmethod
def token_limit(self) -> int:
pass

87
sia/llm_engine/hf_llm_engine.py
View File

@@ -1,87 +0,0 @@
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoConfig
from typing import Iterator, Optional, Callable
from . import LlmEngine
class HfLlmEngine(LlmEngine):
"""
LLM Engine implementation using HuggingFace's InferenceClient.
"""
def __init__(
self,
model: str,
temperature: float,
api_token: Optional[str],
):
"""
Initialize the HuggingFace Inference API LLM Engine.
Args:
model: HuggingFace model ID to use
temperature: Sampling temperature
api_token: HuggingFace API token
"""
self._model = model
self._temperature = temperature
self._tokenizer = AutoTokenizer.from_pretrained(model, token=api_token)
self._config = AutoConfig.from_pretrained(model, token=api_token)
self._client = InferenceClient(token=api_token)
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> 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
continuation_text: Part of the response that is already generated
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context},
{"role": "assistant", "content": continuation_text},
]
stream = self._client.chat_completion(
model=self._model,
messages=messages,
temperature=self._temperature,
add_generation_prompt=False,
stream=True
)
try:
for response in stream:
if should_stop():
stream.close()
break
if content := response.choices[0].delta.content:
yield content
finally:
stream.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return len(self._tokenizer.encode(prompt))
def token_limit(self) -> int:
"""
Get the model's context window size.
Returns:
int: Maximum number of tokens the model can process
"""
return self._config.max_position_embeddings

137
sia/llm_engine/local_llm_engine.py
View File

@@ -1,137 +0,0 @@
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
from typing import Iterator, Optional, Callable
from xml_schema_validator import XmlLogitsProcessor
import sys
import torch
from . import LlmEngine
from .. import util
class LocalLlmEngine(LlmEngine):
def __init__(
self,
model_path: str,
temperature: float,
token_limit: int,
xml_schema_text: Optional[str] = None,
api_token: Optional[str] = None,
):
"""
Initialize the LLM Engine with a model path.
Args:
model_path: Path to the model weights to be used.
temperature: Temperature for sampling
token_limit: Maximum number of tokens to generate
xml_schema_text: Optional XML schema to validate against
api_token: Huggingface API key
"""
self._temperature = temperature
self._token_limit = token_limit
self._tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
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,
token=api_token,
)
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,
)
if xml_schema_text:
self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
else:
self._logits_processor = None
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> 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
continuation_text: Part of the response that is already generated
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context},
{"role": "assistant", "content": continuation_text},
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False,
add_generation_prompt=False,
)
streamer = TextIteratorStreamer(
self._tokenizer,
skip_prompt=True
)
generation_kwargs = {
"text_inputs": prompt,
"do_sample": True,
"temperature": self._temperature,
"max_new_tokens": self.token_limit(),
"streamer": streamer,
}
if self._logits_processor:
generation_kwargs["logits_processor"] = [self._logits_processor.copy()]
generation_thread = Thread(
target=self._pipeline,
kwargs=generation_kwargs
)
generation_thread.start()
for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
yield text
if should_stop():
break
generation_thread.join()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""
Count tokens for the given system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string
Returns:
int: Total number of tokens
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return len(self._tokenizer.encode(prompt))
def token_limit(self) -> int:
"""
Get the model's context window size.
Returns:
int: Maximum number of tokens the model can process
"""
if self._token_limit is not None:
return self._token_limit
else:
return self._pipeline.model.config.max_position_embeddings

83
sia/llm_engine/mistral_llm_engine.py
View File

@@ -1,83 +0,0 @@
from typing import Iterator, Optional, Callable
from mistralai import Mistral
from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
from . import LlmEngine
from ..util import skip_prefix
class MistralLlmEngine(LlmEngine):
def __init__(
self,
model: str,
temperature: float,
token_limit: int,
api_key: str,
):
self._model = model
self._temperature = temperature
self._token_limit = token_limit
self._api_key = api_key
self._client = Mistral(api_key=api_key)
self._tokenizer = MistralTokenizer.v3()
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
messages = [
{
"role": "system",
"content": system_prompt,
},
{
"role": "user",
"content": main_context,
},
{
"role": "assistant",
"content": continuation_text,
"prefix": True,
},
] if continuation_text else [
{
"role": "system",
"content": system_prompt,
},
{
"role": "user",
"content": main_context,
},
]
stream_response = self._client.chat.stream(
model=self._model,
messages=messages,
temperature=self._temperature,
)
try:
def content_generator():
for chunk in stream_response:
if should_stop():
stream_response.response.close()
break
if content := chunk.data.choices[0].delta.content:
yield content
yield from skip_prefix(content_generator(), continuation_text)
finally:
stream_response.response.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
messages = [
SystemMessage(content=system_prompt),
UserMessage(content=main_context),
]
tokenized = self._tokenizer.encode_chat_completion(
ChatCompletionRequest(
messages=messages,
model=self._model
)
)
return len(tokenized.tokens)
def token_limit(self) -> int:
return self._token_limit

77
sia/llm_engine/openai_llm_engine.py
View File

@@ -1,77 +0,0 @@
from typing import Callable, Iterator
import openai
import tiktoken
from . import LlmEngine
class OpenAILlmEngine(LlmEngine):
"""
LLM Engine implementation using OpenAI's API.
Supports streaming responses from chat completion models.
"""
def __init__(
self,
model: str,
temperature: float,
token_limit: int,
api_key: str,
):
"""
Initialize the OpenAI LLM Engine.
Args:
model: OpenAI model to use
temperature: Temperature for sampling
api_key: OpenAI API key
token_limit: Maximum number of tokens to generate
"""
self._model = model
self._temperature = temperature
self._token_limit = token_limit
self._client = openai.Client(
api_key=api_key,
)
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
if continuation_text:
print("OpenAI LLM Engine: continuation_text is not supported")
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
stream = self._client.chat.completions.create(
model=self._model,
messages=messages,
temperature=self._temperature,
stream=True,
)
try:
for chunk in stream:
if should_stop():
break
if content := chunk.choices[0].delta.content:
yield content
finally:
stream.close()
#stream.response.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""
Calculate the total token count for the system prompt and context.
Args:
system_prompt: The system prompt string
main_context: The main context string
Returns:
int: Total number of tokens
"""
encoding = tiktoken.encoding_for_model(self._model)
return len(encoding.encode(system_prompt)) + len(encoding.encode(main_context))
def token_limit(self) -> int:
return self._token_limit

227
sia/llm_engine/qwq_llm_engine.ipynb
View File

@@ -1,227 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"source /root/venvs/sia/bin/activate\n",
"apt-get update && apt-get install -y cuda-toolkit\n",
"pip install flash-attn --no-build-isolation"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Unsloth should be imported before transformers to ensure all optimizations are applied.\n",
"from unsloth import FastLanguageModel"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"from threading import Thread\n",
"from transformers import AutoTokenizer, TextIteratorStreamer, pipeline\n",
"from xml_schema_validator import XmlLogitsProcessor"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"temperature = 0.6\n",
"model_path = \"/root/models/current\"\n",
"xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Load tokenizer\n",
"tokenizer = AutoTokenizer.from_pretrained(\n",
" model_path,\n",
" legacy=False,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Load model\n",
"model, _returned_tokenizer = FastLanguageModel.from_pretrained(\n",
" model_path,\n",
" gpu_memory_utilization = 0.5, # Reduce if out of memory\n",
" load_in_4bit=True,\n",
" attn_implementation=\"flash_attention_2\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# enable unsloth optimizations\n",
"FastLanguageModel.for_inference(model)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"# Create inference pipeline with memory-efficient settings\n",
"pipeline = pipeline(\n",
" \"text-generation\",\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
" return_full_text=False,\n",
" torch_dtype=torch.float16,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_text)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"messages = [\n",
" {\"role\": \"system\", \"content\": \"Always respond in a <write_stdout></write_stdout> xml block.\"},\n",
" {\"role\": \"user\", \"content\": \"Hi, how are you?\"},\n",
" {\"role\": \"assistant\", \"content\": \"\"},\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"text = tokenizer.apply_chat_template(\n",
" messages,\n",
" tokenize=False,\n",
" add_generation_prompt=False,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"streamer = TextIteratorStreamer(\n",
" tokenizer,\n",
" skip_prompt=True,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"generation_kwargs = {\n",
" \"text_inputs\": text,\n",
" \"do_sample\": True,\n",
" \"temperature\": temperature,\n",
" \"streamer\": streamer,\n",
" \"use_cache\": True,\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"generation_kwargs[\"logits_processor\"] = [logits_processor.copy()]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"generation_thread = Thread(\n",
" target=pipeline,\n",
" kwargs=generation_kwargs\n",
")\n",
"\n",
"generation_thread.start()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for text in streamer:\n",
" print(text, end=\"\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"generation_thread.join()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "sia",
"language": "python",
"name": "sia"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

135
sia/llm_engine/qwq_llm_engine.py
View File

@@ -1,135 +0,0 @@
# Unsloth should be imported before transformers to ensure all optimizations are applied.
from unsloth import FastLanguageModel
from pathlib import Path
from threading import Thread
from transformers import AutoTokenizer, TextIteratorStreamer, pipeline
from typing import Callable, Iterator, Optional
from xml_schema_validator import XmlLogitsProcessor
from . import LlmEngine
from .. import util
class QwQLlmEngine(LlmEngine):
def __init__(
self,
model_path: Path,
temperature: float,
xml_schema_text: Optional[str] = None,
):
"""
Initialize the QwQ LLM Engine.
Args:
model_path: Local path to the model
temperature: Sampling temperature
xml_schema_text: Optional XML schema to validate against
"""
self._temperature = temperature
# Load tokenizer
self._tokenizer = AutoTokenizer.from_pretrained(
model_path,
)
# Load model
self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
model_path,
gpu_memory_utilization = 0.5, # Reduce if out of memory
)
# enable unsloth optimizations
FastLanguageModel.for_inference(self._model)
# Create inference pipeline
self._pipeline = pipeline(
"text-generation",
model=self._model,
tokenizer=self._tokenizer,
return_full_text=False,
)
if xml_schema_text:
self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
else:
self._logits_processor = None
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> 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
continuation_text: Part of the response that is already generated
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context},
{"role": "assistant", "content": continuation_text},
]
text = self._tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=False,
)
streamer = TextIteratorStreamer(
self._tokenizer,
skip_prompt=True,
)
generation_kwargs = {
"text_inputs": text,
"do_sample": True,
"temperature": self._temperature,
"streamer": streamer,
"use_cache": True,
}
if self._logits_processor:
generation_kwargs["logits_processor"] = [self._logits_processor.copy()]
generation_thread = Thread(
target=self._pipeline,
kwargs=generation_kwargs
)
generation_thread.start()
for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
yield text
if should_stop():
break
generation_thread.join()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""
Count tokens for the given system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string
Returns:
int: Total number of tokens
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
return len(self._tokenizer.encode(prompt))
def token_limit(self) -> int:
return self._pipeline.model.config.max_position_embeddings

255
sia/llm_engine/qwq_vllm.ipynb
View File

@@ -1,255 +0,0 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# vLLM Streaming Implementation\n",
"\n",
"This notebook demonstrates how to implement streaming capability with vLLM, comparable to the unsloth implementation.\n",
"\n",
"First, let's make sure we have vLLM installed:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO 04-25 19:36:31 [__init__.py:239] Automatically detected platform cuda.\n"
]
}
],
"source": [
"from pathlib import Path\n",
"from vllm import SamplingParams\n",
"from transformers import AutoTokenizer\n",
"import sys"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"temperature = 0.6\n",
"model_path = \"/root/models/current\"\n",
"xml_schema_text = Path(\"/root/sia/action_schema.xsd\").read_text()"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"# Load tokenizer\n",
"tokenizer = AutoTokenizer.from_pretrained(\n",
" model_path,\n",
" legacy=False,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"messages = [\n",
" {\"role\": \"system\", \"content\": \"Always respond in a <write_stdout></write_stdout> xml block.\"},\n",
" {\"role\": \"user\", \"content\": \"Hi, how are you?\"},\n",
" {\"role\": \"assistant\", \"content\": \"\"},\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"prompt = tokenizer.apply_chat_template(\n",
" messages,\n",
" tokenize=False,\n",
" add_generation_prompt=False,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"# Define sampling parameters\n",
"sampling_params = SamplingParams(\n",
" temperature=temperature,\n",
" top_p=0.95,\n",
" max_tokens=512,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO 04-25 19:36:40 [config.py:585] This model supports multiple tasks: {'generate', 'score', 'reward', 'embed', 'classify'}. Defaulting to 'generate'.\n",
"WARNING 04-25 19:36:42 [config.py:664] bitsandbytes quantization is not fully optimized yet. The speed can be slower than non-quantized models.\n",
"WARNING 04-25 19:36:42 [arg_utils.py:1854] --quantization bitsandbytes is not supported by the V1 Engine. Falling back to V0. \n",
"INFO 04-25 19:36:42 [llm_engine.py:241] Initializing a V0 LLM engine (v0.8.2) with config: model='/root/models/current', speculative_config=None, tokenizer='/root/models/current', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=4096, download_dir=None, load_format=bitsandbytes, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=bitsandbytes, enforce_eager=False, kv_cache_dtype=auto, device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='xgrammar', reasoning_backend=None), observability_config=ObservabilityConfig(show_hidden_metrics=False, otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=None, served_model_name=/root/models/current, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=None, chunked_prefill_enabled=False, use_async_output_proc=True, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={\"splitting_ops\":[],\"compile_sizes\":[],\"cudagraph_capture_sizes\":[256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],\"max_capture_size\":256}, use_cached_outputs=False, \n",
"INFO 04-25 19:36:42 [cuda.py:291] Using Flash Attention backend.\n",
"INFO 04-25 19:36:43 [parallel_state.py:954] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0\n",
"INFO 04-25 19:36:43 [model_runner.py:1110] Starting to load model /root/models/current...\n",
"INFO 04-25 19:36:43 [loader.py:1155] Loading weights with BitsAndBytes quantization. May take a while ...\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "8b9f3cb293484cac932e6cedd841c813",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading safetensors checkpoint shards: 0% Completed | 0/4 [00:00<?, ?it/s]\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "54f8aa5eefdb43d8bc07274044a8bc1c",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Loading safetensors checkpoint shards: 0% Completed | 0/4 [00:00<?, ?it/s]\n"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO 04-25 19:36:51 [model_runner.py:1146] Model loading took 18.0523 GB and 8.113452 seconds\n",
"INFO 04-25 19:36:55 [worker.py:267] Memory profiling takes 3.23 seconds\n",
"INFO 04-25 19:36:55 [worker.py:267] the current vLLM instance can use total_gpu_memory (47.53GiB) x gpu_memory_utilization (0.90) = 42.78GiB\n",
"INFO 04-25 19:36:55 [worker.py:267] model weights take 18.05GiB; non_torch_memory takes 0.06GiB; PyTorch activation peak memory takes 1.59GiB; the rest of the memory reserved for KV Cache is 23.08GiB.\n",
"INFO 04-25 19:36:55 [executor_base.py:111] # cuda blocks: 5907, # CPU blocks: 1024\n",
"INFO 04-25 19:36:55 [executor_base.py:116] Maximum concurrency for 4096 tokens per request: 23.07x\n",
"INFO 04-25 19:36:58 [model_runner.py:1442] Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Capturing CUDA graph shapes: 100%|██████████| 35/35 [01:01<00:00, 1.75s/it]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO 04-25 19:38:00 [model_runner.py:1570] Graph capturing finished in 61 secs, took 1.98 GiB\n",
"INFO 04-25 19:38:00 [llm_engine.py:447] init engine (profile, create kv cache, warmup model) took 68.57 seconds\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\n"
]
}
],
"source": [
"from vllm import LLM, SamplingParams\n",
"import time\n",
"\n",
"# Initialize LLM\n",
"llm = LLM(\n",
" model=model_path,\n",
" tensor_parallel_size=1,\n",
" max_model_len=4096,\n",
" quantization=\"bitsandbytes\",\n",
" load_format=\"bitsandbytes\",\n",
" trust_remote_code=True,\n",
" # Enable streaming\n",
" enable_lora=False,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Starting generation with token-by-token output:\n",
"<think>\n",
"Okay, the user greeted me with \"Hi, how are you?\" I need to respond appropriately. Let me see... The instructions say to always use the <write_stdout> XML tag. So first, I should acknowledge their greeting and state that I'm an AI, then ask how I can assist them. Keep it friendly and helpful. Let me make sure I don't add any extra information beyond that. Just a simple response. Alright, that should work.\n",
"</think>\n",
"\n",
"<write_stdout>\n",
"Hello! I'm just a computer program, but I'm here to help you. How can I assist you today?\n",
"</write_stdout>"
]
}
],
"source": [
"previous_text = \"\"\n",
"print(\"Starting generation with token-by-token output:\")\n",
"\n",
"# Try with direct iteration over the generator\n",
"for output in llm.generate(prompt, sampling_params, use_tqdm=False):\n",
" if hasattr(output, 'outputs') and output.outputs and len(output.outputs) > 0:\n",
" generated_text = output.outputs[0].text\n",
" if len(generated_text) > len(previous_text):\n",
" new_text = generated_text[len(previous_text):]\n",
" sys.stdout.write(new_text)\n",
" sys.stdout.flush()\n",
" previous_text = generated_text"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "sia",
"language": "python",
"name": "sia"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

62
sia/util.py
View File

@@ -1,67 +1,5 @@
import datetime
import xml.dom.minidom
import xml.etree.ElementTree as ET
from typing import Iterator, Optional
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
def pretty_print_element(elem: ET.Element, level: int = 0, max_line: int = 80) -> str:
"""Convert ElementTree element to pretty-printed string with custom formatting."""

182
sia/web/api.py
View File

@@ -1,14 +1,14 @@
from pathlib import Path
from aiohttp import web
import json
from pathlib import Path
import asyncio
import json
from ..auto_approver import AutoApprover
from ..auto_approver import AutoApprover, AutoApproverConfig
from ..chat_io_buffer import ChatIOBuffer
from ..entry.entry_factory import EntryFactory
from ..iteration_parser import IterationParser
from ..web_agent import WebAgent
from ..web_io_buffer import WebIOBuffer
from ..working_memory import WorkingMemory
class Api:
@@ -17,7 +17,7 @@ class Api:
work_dir: Path,
app: web.Application,
agent: WebAgent,
io_buffer: WebIOBuffer,
io_buffer: ChatIOBuffer,
working_memory: WorkingMemory,
auto_approver: AutoApprover
):
@@ -32,23 +32,24 @@ class Api:
def _init_routes(self):
"""Initialize REST API and WebSocket routes."""
self._app.router.add_post("/api/inference/{llm}", self._run_inference)
self._app.router.add_post("/api/inference/{llm}/stop", self._stop_inference)
self._app.router.add_get("/api/llms", self._get_llms)
self._app.router.add_get("/api/llms/active", self._get_active_llm)
self._app.router.add_post("/api/llms/active", self._set_active_llm)
self._app.router.add_post("/api/inference", self._run_inference)
self._app.router.add_post("/api/inference/stop", self._stop_inference)
self._app.router.add_get("/api/response", self._get_response)
self._app.router.add_post("/api/response", self._set_response)
self._app.router.add_post("/api/response/approve", self._approve_response)
self._app.router.add_get("/api/response", self._get_response)
self._app.router.add_post("/api/context", self._modify_context)
self._app.router.add_post("/api/input", self._send_input)
self._app.router.add_post("/api/clear", self._clear_output)
self._app.router.add_get("/api/output/{llm}", self._get_output)
self._app.router.add_get("/api/llms", self._get_llms)
self._app.router.add_get("/api/auto_approver/config", self._get_auto_approver_config)
self._app.router.add_post("/api/auto_approver/config", self._set_auto_approver_config)
self._app.router.add_post("/api/auto_approver/context_enabled", self._set_context_enabled)
self._app.router.add_post("/api/auto_approver/response_enabled", self._set_response_enabled)
self._app.router.add_post("/api/auto_approver/context_timeout", self._set_context_timeout)
self._app.router.add_post("/api/auto_approver/response_timeout", self._set_response_timeout)
self._app.router.add_post("/api/auto_approver/llm", self._set_llm_name)
self._app.router.add_post("/api/chat/message", self._add_chat_message)
self._app.router.add_delete("/api/chat/message/{id}", self._delete_chat_message)
self._app.router.add_delete("/api/chat", self._clear_chat)
self._app.router.add_get("/api/auto_approver", self._get_auto_approver_config)
self._app.router.add_post("/api/auto_approver", self._set_auto_approver_config)
self._app.router.add_get("/api/memory", self._get_memory)
self._app.router.add_post("/api/memory/entry", self._create_entry)
self._app.router.add_put("/api/memory/entry/{id}", self._save_entry)
@@ -57,16 +58,33 @@ class Api:
self._app.router.add_post("/api/memory/entry/{id}/update", self._update_entry)
self._app.router.add_post("/api/memory/load_iteration", self._load_iteration)
async def _run_inference(self, request: web.Request) -> web.Response:
"""Start inference on specified LLM."""
async def _get_llms(self, request: web.Request) -> web.Response:
return web.Response(
text=json.dumps(
[{"name": name} for name in self._agent.llms]
),
content_type="application/json"
)
async def _get_active_llm(self, request: web.Request) -> web.Response:
return web.Response(
text=json.dumps(
{"active_llm": self._agent.active_llm}
),
content_type="application/json"
)
async def _set_active_llm(self, request: web.Request) -> web.Response:
try:
llm_name = request.match_info["llm"]
data = await request.json()
text = data.get("response")
context = data.get("context")
self._agent._response_buffer.set_text(text)
self._agent.modify_context(context)
await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference, llm_name)
self._agent.active_llm = data["active_llm"]
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
async def _run_inference(self, request: web.Request) -> web.Response:
try:
await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference)
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
@@ -74,29 +92,7 @@ class Api:
async def _stop_inference(self, request: web.Request) -> web.Response:
"""Stop inference on specified LLM."""
try:
llm_name = request.match_info["llm"]
self._agent.stop_inference(llm_name)
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
async def _set_response(self, request: web.Request) -> web.Response:
"""Edit the shared response buffer"""
try:
data = await request.json()
text = data.get("response")
self._agent._response_buffer.set_text(text)
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
async def _approve_response(self, request: web.Request) -> web.Response:
"""Approve current buffer content"""
data = await request.json()
try:
response = data.get("response")
self._agent.response_buffer.set_text(response)
self._agent.approve_response()
self._agent.stop_inference()
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
@@ -110,53 +106,50 @@ class Api:
content_type="application/json"
)
async def _modify_context(self, request: web.Request) -> web.Response:
"""Modify the current context."""
async def _set_response(self, request: web.Request) -> web.Response:
"""Edit the shared response buffer"""
try:
data = await request.json()
context = data.get("context")
self._agent.modify_context(context)
text = data.get("response")
self._agent._response_buffer.set_text(text)
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
async def _send_input(self, request: web.Request) -> web.Response:
"""Send input to the IO buffer."""
async def _approve_response(self, request: web.Request) -> web.Response:
"""Approve current buffer content"""
try:
data = await request.json()
input_text = data.get("input")
self._io_buffer.append_stdin(input_text)
self._agent.approve_response()
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
async def _clear_output(self, request: web.Request) -> web.Response:
"""Clear the stdout buffer."""
self._io_buffer.clear_stdout()
return web.Response(status=200)
async def _get_output(self, request: web.Request) -> web.Response:
"""Get complete output for specified LLM."""
async def _add_chat_message(self, request: web.Request) -> web.Response:
"""Add a chat message"""
try:
llm_name = request.match_info["llm"]
output = self._agent.get_output(llm_name)
return web.Response(
text=json.dumps({"output": output}),
content_type="application/json"
)
data = await request.json()
message = data.get("message")
self._io_buffer.add_user_message(message)
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
async def _delete_chat_message(self, request: web.Request) -> web.Response:
"""Delete a chat message"""
try:
message_id = request.match_info["id"]
self._io_buffer.delete_message(message_id)
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
async def _clear_chat(self, request: web.Request) -> web.Response:
"""Clear chat messages"""
try:
self._io_buffer.clear()
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
async def _get_llms(self, request: web.Request) -> web.Response:
"""Get all LLMs and their current states."""
states = self._agent.llms
return web.Response(
text=json.dumps(
[{"name": name, "state": state.name}
for name, state in states.items()]
),
content_type="application/json"
)
async def _get_auto_approver_config(self, request: web.Request) -> web.Response:
"""Get current auto approver configuration."""
@@ -169,7 +162,7 @@ class Api:
"""Update auto approver configuration."""
try:
data = await request.json()
self._auto_approver.set_config(data)
self._auto_approver.config = AutoApproverConfig(**data)
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
@@ -219,18 +212,6 @@ class Api:
except Exception as e:
return web.Response(status=400, text=str(e))
async def _set_llm_name(self, request: web.Request) -> web.Response:
"""Set LLM name for auto-approval."""
try:
data = await request.json()
name = data.get("name")
if name is None:
return web.Response(status=400, text="Missing name parameter")
self._auto_approver.llm_name = name
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
async def _get_memory(self, request: web.Request) -> web.Response:
"""Get complete working memory state."""
entries = self._working_memory.get_entries()
@@ -304,13 +285,12 @@ class Api:
if not content:
return web.Response(status=400, text="Missing content in request body")
(context, response, entries) = IterationParser.parse_iteration(content, self._work_dir, self._io_buffer)
(_context, response, entries) = IterationParser.parse_iteration(content, self._work_dir, self._io_buffer)
for entry in entries:
self._working_memory.add_entry(entry)
self._agent.modify_context(context, True)
self._agent.response_buffer.set_text(response)
return web.Response(status=200)
except Exception as e:
return web.Response(status=400, text=str(e))
return web.Response(status=400, text=str(e))

5
sia/web/auto_approver_websocket.py
View File

@@ -1,4 +1,5 @@
from aiohttp import web, WSMsgType
from dataclasses import asdict
from typing import Dict, Set
from ..auto_approver import AutoApprover, AutoApproverConfig
@@ -28,7 +29,7 @@ class AutoApproverWebSocket:
async def _handle_config_change(self, config: AutoApproverConfig):
"""Handle config changes from the AutoApprover."""
await self._broadcast_message({
"config": config
"config": asdict(config)
})
async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
@@ -41,7 +42,7 @@ class AutoApproverWebSocket:
try:
# Send initial config
await ws.send_json({
"config": self._auto_approver.config
"config": asdict(self._auto_approver.config)
})
async for msg in ws:

64
sia/web/chat_websocket.py Normal file
View File

@@ -0,0 +1,64 @@
from aiohttp import web, WSMsgType
from typing import Dict, Set
from .util import wrap_async
from ..chat_io_buffer import ChatIOBuffer
class ChatWebSocket:
def __init__(self, io_buffer: ChatIOBuffer):
self._io_buffer = io_buffer
self._clients: Set[web.WebSocketResponse] = set()
self._io_buffer.add_change_handler(wrap_async(self._handle_io_buffer_change))
async def _broadcast_message(self, message: Dict):
"""Broadcast message to all connected clients."""
disconnected = set()
for ws in self._clients:
try:
await ws.send_json(message)
except ConnectionResetError:
disconnected.add(ws)
self._clients -= disconnected
async def _handle_io_buffer_change(self, new_messages, deleted_message_ids, last_read_timestamp):
"""Handle changes to the ChatIOBuffer."""
await self._broadcast_message({
"messages": [
{
"id": message.id,
"content": message.content,
"message_type": message.message_type.value
}
for message in new_messages
],
"deleted_message_ids": deleted_message_ids,
"last_read_timestamp": last_read_timestamp
})
async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
"""Handle new WebSocket connections."""
ws = web.WebSocketResponse(heartbeat=30)
await ws.prepare(request)
try:
# Send initial state
await ws.send_json({
"messages": [
{
"id": message.id,
"content": message.content,
"message_type": message.message_type.value
}
for message in self._io_buffer.messages
],
"last_read_timestamp": self._io_buffer.last_read_timestamp
})
self._clients.add(ws)
async for msg in ws:
if msg.type == WSMsgType.ERROR:
print(f"WebSocket connection closed with error: {ws.exception()}")
finally:
self._clients.remove(ws)
return ws

55
sia/web/context_websocket.py
View File

@@ -1,55 +0,0 @@
from aiohttp import web, WSMsgType
from typing import Dict, Set
from .util import wrap_async
from ..web_agent import WebAgent
class ContextWebSocket:
"""
WebSocket handler for context changes.
Broadcasts context updates to all connected clients.
"""
def __init__(self, agent: WebAgent):
self._agent = agent
self._clients: Set[web.WebSocketResponse] = set()
self._agent.add_context_change_handler(wrap_async(self._handle_context_change))
async def _broadcast_message(self, message: Dict):
"""Broadcast message to all connected clients."""
disconnected = set()
for ws in self._clients:
try:
await ws.send_json(message)
except ConnectionResetError:
disconnected.add(ws)
self._clients -= disconnected
async def _handle_context_change(self, context: str, generated: bool):
"""Handle context changes from the WebAgent."""
await self._broadcast_message({
"context": context,
"generated": generated
})
async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
"""Handle new WebSocket connections."""
ws = web.WebSocketResponse(heartbeat=30)
await ws.prepare(request)
self._clients.add(ws)
try:
# Send initial context
await ws.send_json({
"context": self._agent.context,
"generated": True
})
async for msg in ws:
if msg.type == WSMsgType.ERROR:
print(f"WebSocket connection closed with error: {ws.exception()}")
finally:
self._clients.remove(ws)
return ws

2
sia/web/memory_websocket.py
View File

@@ -1,8 +1,6 @@
from aiohttp import web, WSMsgType
from typing import Dict, Set
from ..entry import Entry
from ..web_agent import WebAgent
from ..working_memory import WorkingMemory
from .util import wrap_async

33
sia/web/llm_websocket.py → sia/web/state_websocket.py
View File

@@ -2,18 +2,19 @@ from aiohttp import web, WSMsgType
from typing import Dict, Set
from .util import wrap_async
from ..web_agent import WebAgent, LlmState
from ..web_agent import WebAgent, AgentState
class LlmWebSocket:
class StateWebSocket:
"""
WebSocket handler for LLM state changes.
WebSocket handler for agent state changes.
Broadcasts state updates to all connected clients.
"""
def __init__(self, agent: WebAgent):
self._agent = agent
self._clients: Set[web.WebSocketResponse] = set()
self._agent.add_llm_change_handler(wrap_async(self._handle_state_change))
self._agent.add_state_change_handler(wrap_async(self._handle_change))
self._agent.add_selected_llm_change_handler(wrap_async(self._handle_change))
async def _broadcast_message(self, message: Dict):
"""Broadcast message to all connected clients."""
@@ -24,12 +25,12 @@ class LlmWebSocket:
except ConnectionResetError:
disconnected.add(ws)
self._clients -= disconnected
async def _handle_state_change(self, llm_name: str, new_state: LlmState):
"""Handle state changes from the WebAgent."""
async def _handle_change(self, arg):
"""Handle changes to the active LLM."""
await self._broadcast_message({
"llm": llm_name,
"state": new_state.name
"state": self._agent.state.name,
"active_llm": self._agent.active_llm
})
async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
@@ -38,16 +39,12 @@ class LlmWebSocket:
await ws.prepare(request)
try:
# Send initial states for all LLMs
states = self._agent.llms
for llm_name, state in states.items():
await ws.send_json({
"llm": llm_name,
"state": state.name
})
# Send initial state
await ws.send_json({
"state": self._agent.state.name,
"active_llm": self._agent.active_llm
})
self._clients.add(ws)
async for msg in ws:
if msg.type == WSMsgType.ERROR:
print(f"WebSocket connection closed with error: {ws.exception()}")

19
sia/web/websockets.py
View File

@@ -1,28 +1,25 @@
from aiohttp import web
from ..auto_approver import AutoApprover
from ..chat_io_buffer import ChatIOBuffer
from ..web_agent import WebAgent
from ..web_io_buffer import WebIOBuffer
from ..working_memory import WorkingMemory
from .auto_approver_websocket import AutoApproverWebSocket
from .context_websocket import ContextWebSocket
from .llm_websocket import LlmWebSocket
from .chat_websocket import ChatWebSocket
from .memory_websocket import MemoryWebSocket
from .response_websocket import ResponseWebSocket
from .stdout_websocket import StdoutWebSocket
from .state_websocket import StateWebSocket
class Websockets:
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory):
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: ChatIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory):
self._auto_approver_ws = AutoApproverWebSocket(auto_approver)
self._context_ws = ContextWebSocket(agent)
self._llm_ws = LlmWebSocket(agent)
self._chat_ws = ChatWebSocket(io_buffer)
self._memory_ws = MemoryWebSocket(working_memory)
self._response_ws = ResponseWebSocket(agent)
self._stdout_ws = StdoutWebSocket(io_buffer)
self._state_ws = StateWebSocket(agent)
app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection)
app.router.add_get("/ws/context", self._context_ws.handle_connection)
app.router.add_get("/ws/llms", self._llm_ws.handle_connection)
app.router.add_get("/ws/chat", self._chat_ws.handle_connection)
app.router.add_get("/ws/memory", self._memory_ws.handle_connection)
app.router.add_get("/ws/response", self._response_ws.handle_connection)
app.router.add_get("/ws/stdout", self._stdout_ws.handle_connection)
app.router.add_get("/ws/state", self._state_ws.handle_connection)

145
sia/web_agent.py
View File

@@ -1,13 +1,10 @@
from collections import defaultdict
from datetime import datetime, timezone
from enum import Enum, auto
from sys import exit
from threading import Lock
from typing import Callable, Dict, List, Optional
from typing import Callable, Dict, List
from .base_agent import BaseAgent
from .command import Command
from .command_result import CommandResult
from .iteration_logger import IterationLogger
from .llm_engine import LlmEngine
from .response_buffer import ResponseBuffer
@@ -15,9 +12,10 @@ from .response_parser import ResponseParser
from .system_metrics import SystemMetrics
from .working_memory import WorkingMemory
class LlmState(Enum):
class AgentState(Enum):
IDLE = auto()
INFERENCE = auto()
PROCESSING_RESPONSE = auto()
class WebAgent(BaseAgent):
def __init__(
@@ -37,97 +35,81 @@ class WebAgent(BaseAgent):
metrics,
parser
)
self._llms = llms
self._selected_llm = list(llms.keys())[0]
self._state: AgentState = AgentState.IDLE
self._iteration_logger = iteration_logger
self._llm_states: Dict[str, LlmState] = {name: LlmState.IDLE for name in llms}
self._response_buffer: ResponseBuffer = ResponseBuffer()
self._validation_error: Optional[str] = None
self._command_result: Optional[CommandResult] = None
self._context = self._compile_context(next(iter(self._llms.values())))
self._stop_flags: Dict[str, bool] = {name: False for name in llms}
# Locks
self._llm_lock = Lock()
self._output_lock = Lock()
# Event handlers
self._llm_change_handlers: List[Callable[[str, LlmState], None]] = []
self._token_handlers: List[Callable[[str, str], None]] = []
self._context_change_handlers: List[Callable[[str, bool], None]] = []
self._state_change_handlers: List[Callable[[AgentState], None]] = []
self._selected_llm_change_handlers: List[Callable[[str], None]] = []
# Change handlers
self._working_memory.add_change_handler(self._handle_memory_update)
@property
def llms(self) -> Dict[str, LlmState]:
"""Get current state of all LLMs"""
with self._llm_lock:
return self._llm_states.copy()
self._update_compiled_context()
self._working_memory.add_change_handler(self._update_compiled_context)
@property
def response_buffer(self) -> ResponseBuffer:
return self._response_buffer
@property
def context(self) -> str:
return self._context
def llms(self) -> List[str]:
return self._llms.keys()
@property
def command_result(self) -> Optional[CommandResult]:
return self._command_result
def active_llm(self) -> str:
return self._selected_llm
@active_llm.setter
def active_llm(self, llm_name: str) -> None:
if llm_name not in self._llms:
raise ValueError(f"Invalid LLM name: {llm_name}")
if self._selected_llm == llm_name:
return
self._selected_llm = llm_name
self._update_compiled_context()
self._notify_selected_llm_change()
def add_selected_llm_change_handler(self, handler: Callable[[str], None]) -> None:
self._selected_llm_change_handlers.append(handler)
@property
def validation_error(self) -> Optional[str]:
return self._validation_error
def state(self) -> AgentState:
return self._state
def add_llm_change_handler(self, handler: Callable[[str, LlmState], None]) -> None:
"""Add handler for LLM state changes"""
if handler not in self._llm_change_handlers:
self._llm_change_handlers.append(handler)
def add_state_change_handler(self, handler: Callable[[AgentState], None]) -> None:
self._state_change_handlers.append(handler)
def add_context_change_handler(self, handler: Callable[[str, bool], None]) -> None:
"""Add handler for context changes"""
if handler not in self._context_change_handlers:
self._context_change_handlers.append(handler)
def modify_context(self, context: str, generated: bool = False) -> None:
"""Update context and reset all LLM states"""
with self._llm_lock:
self._context = context
for handler in self._context_change_handlers:
handler(context, generated)
def run_inference(self, llm_name: str) -> None:
def run_inference(self) -> None:
"""Start inference on specified LLM"""
if llm_name not in self._llms:
raise ValueError(f"Unknown LLM: {llm_name}")
with self._llm_lock:
if self._llm_states[llm_name] != LlmState.IDLE:
raise RuntimeError(f"LLM {llm_name} is not ready for inference")
self._set_llm_state(llm_name, LlmState.INFERENCE)
self._stop_flags[llm_name] = False
llm = self._llms[llm_name]
def should_stop() -> bool:
return self._stop_flags[llm_name]
response_token_iter = llm.infer(self.system_prompt, self.context, self._response_buffer.get_text(), should_stop)
if self._state != AgentState.IDLE:
raise RuntimeError(f"Not ready for inference")
self._update_state(AgentState.INFERENCE)
llm = self._llms[self.active_llm]
response_token_iter = llm.infer(
self.system_prompt,
self._compiled_context,
self._response_buffer.get_text()
)
for token in response_token_iter:
with self._output_lock:
self._response_buffer.append_text(token)
with self._llm_lock:
self._set_llm_state(llm_name, LlmState.IDLE)
self._response_buffer.append_text(token)
self._update_state(AgentState.IDLE)
def stop_inference(self, llm_name: str) -> None:
def stop_inference(self) -> None:
"""Stop ongoing inference for specified LLM"""
if llm_name not in self._llms:
raise ValueError(f"Unknown LLM: {llm_name}")
self._stop_flags[llm_name] = True
with self._llm_lock:
self._set_llm_state(llm_name, LlmState.IDLE)
self._llms[self.active_llm].restart()
self._update_state(AgentState.IDLE)
def approve_response(self) -> None:
"""Process approved response from specified LLM"""
self._update_state(AgentState.PROCESSING_RESPONSE)
timestamp = datetime.now(timezone.utc)
self._iteration_logger.log_iteration(timestamp, self._context, self._response_buffer.get_text())
self._iteration_logger.log_iteration(
timestamp,
self._compiled_context,
self._response_buffer.get_text()
)
parse_result = self._parser.parse(timestamp, self._response_buffer.get_text())
self._response_buffer.clear()
if isinstance(parse_result, Command):
@@ -141,14 +123,13 @@ class WebAgent(BaseAgent):
parse_result.update()
self._working_memory.update()
self._working_memory.add_entry(parse_result)
self._update_state(AgentState.IDLE)
def _set_llm_state(self, llm_name: str, state: LlmState) -> None:
def _update_state(self, state: AgentState) -> None:
"""Update LLM state and notify handlers"""
self._llm_states[llm_name] = state
for handler in self._llm_change_handlers:
handler(llm_name, state)
def _handle_memory_update(self) -> None:
"""Handle memory updates and update context"""
context = self._compile_context(next(iter(self._llms.values())))
self.modify_context(context, True)
self._state = state
for handler in self._state_change_handlers:
handler(state)
def _update_compiled_context(self):
self._compiled_context = self._compile_context(self._llms[self.active_llm])

4
sia/working_memory.py
View File

@@ -85,8 +85,8 @@ class WorkingMemory:
def __del__(self):
"""Clean up all entries when memory is deleted."""
if hasattr(self, '_entries'):
self.clear()
self._change_handlers.clear()
self.clear()
def get_entry(self, id: str) -> Optional[Entry]:
"""