diff --git a/requirements.txt b/requirements.txt
index da16518..c39ba5f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,5 +3,6 @@ aiohttp
bs4
openai
python-dotenv
+tiktoken
torch
transformers
\ No newline at end of file
diff --git a/sia/__main__.py b/sia/__main__.py
index 23d8146..5541aa7 100644
--- a/sia/__main__.py
+++ b/sia/__main__.py
@@ -25,14 +25,6 @@ mimetypes.add_type("application/javascript", ".jsx")
mimetypes.add_type("text/javascript", ".js")
mimetypes.add_type("text/javascript", ".jsx")
-class TestLLM:
- def infer(self, prompt: str, context: str):
- yield ""
- time.sleep(2)
- yield "test reasoning"
- time.sleep(2)
- yield ""
-
class Main:
@classmethod
async def create(cls, config: Config):
@@ -44,20 +36,24 @@ class Main:
match self._config.llm_engine:
case "local":
- self._llm = LocalLlmEngine(self._config.model)
+ self._llm = LocalLlmEngine(
+ self._config.model,
+ self._config.temperature,
+ self._config.token_limit
+ )
case "hf":
self._llm = HfLlmEngine(
- model_id=self._config.model,
- api_token=self._config.api_token,
- temperature=self._config.temperature
+ self._config.model,
+ self._config.temperature,
+ self._config.api_token,
)
case "openai":
self._llm = OpenAILlmEngine(
- model=self._config.model,
- api_key=self._config.api_token
+ self._config.model,
+ self._config.temperature,
+ self._config.api_token,
+ self._config.token_limit
)
- case "test":
- self._llm = TestLLM()
case _:
raise ValueError(f"Invalid LLM engine: {self._config.llm_engine}")
self._io_buffer = WebIOBuffer()
@@ -65,7 +61,7 @@ class Main:
system_prompt=self._system_prompt,
action_schema=self._action_schema,
working_memory=WorkingMemory(),
- system_metrics=SystemMetrics(),
+ metrics=SystemMetrics(),
llm=self._llm,
validator=XMLValidator(self._action_schema),
parser=ResponseParser(self._io_buffer)
diff --git a/sia/base_agent.py b/sia/base_agent.py
index b10f86f..c490ccf 100644
--- a/sia/base_agent.py
+++ b/sia/base_agent.py
@@ -18,27 +18,36 @@ class BaseAgent(ABC):
and coordinating components for LLM inference.
"""
- def __init__(self,
- action_schema: str,
- working_memory: WorkingMemory,
- system_metrics: SystemMetrics,
- llm: LlmEngine,
- validator: XMLValidator,
- parser: ResponseParser):
+ def __init__(
+ self,
+ system_prompt: str,
+ action_schema: str,
+ working_memory: WorkingMemory,
+ metrics: SystemMetrics,
+ llm: LlmEngine,
+ validator: XMLValidator,
+ parser: ResponseParser
+ ):
"""
Initialize agent with required components.
"""
+ self._system_prompt = system_prompt
+ self._action_schema = action_schema
self._working_memory = working_memory
- self._metrics = system_metrics
+ self._metrics = metrics
self._llm = llm
self._validator = validator
self._parser = parser
- self._action_schema = action_schema
def __del__(self):
"""Clean up resources on deletion."""
if hasattr(self, '_metrics'):
self._metrics.stop()
+
+ @property
+ def system_prompt(self) -> str:
+ """Get the system prompt."""
+ return f"{self._system_prompt}\n{self._action_schema}"
def _compile_context(self) -> str:
"""
@@ -48,14 +57,10 @@ class BaseAgent(ABC):
Returns:
str: Complete context as XML string
"""
- # Get memory context and calculate size
memory_context = self._working_memory.generate_context()
- context_size = len(memory_context) / 100
-
- # Get system metrics
metrics_data = self._metrics.get_metrics()
- # Create context element with metrics
+ # Create context element
context = ET.Element("context")
context.set("time", metrics_data["timestamp"])
context.set("cpu", str(metrics_data["cpu"]))
@@ -64,11 +69,20 @@ class BaseAgent(ABC):
context.set("memory_total", str(metrics_data["memory_total"]))
context.set("disk_used", str(metrics_data["disk_used"]))
context.set("disk_total", str(metrics_data["disk_total"]))
- context.set("context", str(round(context_size * 100)))
context.set("stdin", str(self._parser.io_buffer.buffer_length()))
-
- # Add memory entries
+ context.set("context", "100")
+
for entry in memory_context:
context.append(entry)
-
- return pretty_print_element(context)
+
+ context_str = pretty_print_element(context)
+
+ # Calculate token usage percentage
+ token_count = self._llm.token_count(self.system_prompt, context_str)
+ token_limit = self._llm.token_limit()
+ context_usage = (float(token_count) / float(token_limit)) * 100.0
+
+ # Update context usage metric
+ context.set("context", str(round(context_usage, 2)))
+
+ return pretty_print_element(context)
\ No newline at end of file
diff --git a/sia/config.py b/sia/config.py
index 96f33a2..ee53758 100644
--- a/sia/config.py
+++ b/sia/config.py
@@ -81,6 +81,12 @@ class Config:
default=float(os.getenv('SIA_TEMPERATURE', '0.7')),
help='LLM temperature parameter (default: 0.7, env: SIA_TEMPERATURE)'
)
+ parser.add_argument(
+ '--token-limit',
+ type=int,
+ default=os.getenv('SIA_TOKEN_LIMIT'),
+ help='Token limit for the LLM (env: SIA_TOKEN_LIMIT)'
+ )
self.args = parser.parse_args()
def _parse_bool_env(self, env_var: str, default: bool) -> bool:
@@ -145,4 +151,9 @@ class Config:
@property
def temperature(self) -> float:
"""LLM temperature parameter."""
- return self.args.temperature
\ No newline at end of file
+ return self.args.temperature
+
+ @property
+ def token_limit(self) -> int:
+ """Token limit for the LLM."""
+ return self.args.token_limit
\ No newline at end of file
diff --git a/sia/hf_llm_engine.py b/sia/hf_llm_engine.py
index 3182030..98a1207 100644
--- a/sia/hf_llm_engine.py
+++ b/sia/hf_llm_engine.py
@@ -1,5 +1,6 @@
-from typing import Iterator, Optional
from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoConfig
+from typing import Iterator, Optional
from .llm_engine import LlmEngine
@@ -10,35 +11,24 @@ class HfLlmEngine(LlmEngine):
def __init__(
self,
- model_id: str = "mistralai/Mistral-7B-Instruct-v0.2",
- api_token: Optional[str] = None,
- temperature: float = 0.7,
- max_new_tokens: int = 1024,
+ model: str,
+ temperature: float,
+ api_token: Optional[str],
):
"""
Initialize the HuggingFace Inference API LLM Engine.
Args:
- model_id: HuggingFace model ID to use (default: Mistral-7B-Instruct)
- api_token: HuggingFace API token. If None, will try to read from HF_TOKEN env var
- temperature: Sampling temperature (default: 0.7)
- max_new_tokens: Maximum number of tokens to generate (default: 1024)
+ model: HuggingFace model ID to use
+ temperature: Sampling temperature
+ api_token: HuggingFace API token
"""
- self.model_id = model_id
- self.client = InferenceClient(token=api_token)
-
- # Generation parameters
- self.temperature = temperature
- self.max_new_tokens = max_new_tokens
-
- def set_model_path(self, model_id: str):
- """
- Update the model being used.
+ self._model = model
+ self._temperature = temperature
- Args:
- model_id: New HuggingFace model ID to use
- """
- self.model_id = model_id
+ 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) -> Iterator[str]:
"""
@@ -51,21 +41,41 @@ class HfLlmEngine(LlmEngine):
Returns:
Iterator[str]: An iterator that yields the generated text.
"""
+ token_count=self.token_count(system_prompt, main_context)
+
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
def stream_wrapper():
- stream = self.client.chat_completion(
- model=self.model_id,
+ stream = self._client.chat_completion(
+ model=self._model,
messages=messages,
- temperature=self.temperature,
- max_tokens=self.max_new_tokens,
+ temperature=self._temperature,
stream=True
)
for response in stream:
if content := response.choices[0].delta.content:
yield content
- return stream_wrapper()
\ No newline at end of file
+ return stream_wrapper()
+
+ 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
diff --git a/sia/llm_engine.py b/sia/llm_engine.py
index 522c790..ffee779 100644
--- a/sia/llm_engine.py
+++ b/sia/llm_engine.py
@@ -2,7 +2,14 @@ from typing import Iterator
from abc import ABC, abstractmethod
class LlmEngine(ABC):
-
@abstractmethod
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
+ pass
+
+ @abstractmethod
+ def token_count(self, system_prompt: str, main_context: str) -> int:
+ pass
+
+ @abstractmethod
+ def token_limit(self) -> int:
pass
\ No newline at end of file
diff --git a/sia/local_llm_engine.py b/sia/local_llm_engine.py
index f3997d2..c5683ec 100644
--- a/sia/local_llm_engine.py
+++ b/sia/local_llm_engine.py
@@ -1,5 +1,5 @@
from threading import Thread
-from typing import Iterator
+from typing import Iterator, Optional
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
import torch
@@ -8,23 +8,23 @@ from . import util
from .llm_engine import LlmEngine
class LocalLlmEngine(LlmEngine):
- def __init__(self, model_path: str):
+ def __init__(
+ self,
+ model_path: str,
+ temperature: float,
+ token_limit: int,
+ ):
"""
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
"""
- self.set_model_path(model_path)
-
- def set_model_path(self, model_path: str):
- """
- Load the model from the specified path.
-
- Args:
- model_path: Path to the model weights to load.
- """
- self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+ self._temperature = temperature
+ self._token_limit = token_limit
+ self._tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
model_path,
return_dict=True,
@@ -33,14 +33,14 @@ class LocalLlmEngine(LlmEngine):
device_map="auto",
trust_remote_code=True,
)
- if self.tokenizer.pad_token_id is None:
- self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+ if 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(
+ self._pipeline = pipeline(
"text-generation",
model=model,
- tokenizer=self.tokenizer,
+ tokenizer=self._tokenizer,
torch_dtype=torch.bfloat16,
device_map="auto",
return_full_text=False,
@@ -61,19 +61,49 @@ class LocalLlmEngine(LlmEngine):
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
- prompt = self.tokenizer.apply_chat_template(
+ prompt = self._tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
streamer = TextIteratorStreamer(
- self.tokenizer,
+ self._tokenizer,
skip_prompt=True
)
pipeline_kwargs = dict(
text_inputs=prompt,
do_sample=True,
- max_new_tokens=1024,
+ temperature=self._temperature,
+ max_new_tokens=self._token_limit,
streamer=streamer
)
- thread = Thread(target=self.pipeline, kwargs=pipeline_kwargs)
+ thread = Thread(target=self._pipeline, kwargs=pipeline_kwargs)
thread.start()
return util.stop_before_value(streamer, '<|eot_id|>')
+
+ 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
+ """
+ return self._pipeline.model.config.max_position_embeddings
\ No newline at end of file
diff --git a/sia/openai_llm_engine.py b/sia/openai_llm_engine.py
index 88607a8..43464a7 100644
--- a/sia/openai_llm_engine.py
+++ b/sia/openai_llm_engine.py
@@ -1,6 +1,6 @@
-from typing import Iterator, Optional
+from typing import Iterator
import openai
-import json
+import tiktoken
from .llm_engine import LlmEngine
@@ -13,19 +13,24 @@ class OpenAILlmEngine(LlmEngine):
def __init__(
self,
model: str,
+ temperature: float,
api_key: str,
+ token_limit: int = 0,
):
"""
Initialize the OpenAI LLM Engine.
Args:
- model: OpenAI model to use (default: gpt-4)
- api_key: OpenAI API key. If None, will try to read from OPENAI_API_KEY env var
+ 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
-
- # Initialize OpenAI client
- self.client = openai.Client(
+ self._temperature = temperature
+ self._token_limit = token_limit
+
+ self._client = openai.Client(
api_key=api_key,
)
@@ -45,13 +50,30 @@ class OpenAILlmEngine(LlmEngine):
{"role": "user", "content": main_context}
]
- stream = self.client.chat.completions.create(
+ stream = self._client.chat.completions.create(
model=self._model,
messages=messages,
+ temperature=self._temperature,
stream=True,
- temperature=0.3,
)
for chunk in stream:
if content := chunk.choices[0].delta.content:
- yield content
\ No newline at end of file
+ yield content
+
+ 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
\ No newline at end of file
diff --git a/sia/web_agent.py b/sia/web_agent.py
index 898bcae..deb10a5 100644
--- a/sia/web_agent.py
+++ b/sia/web_agent.py
@@ -33,20 +33,28 @@ class WebAgent(BaseAgent):
Broadcasts state changes to registered handlers.
"""
- def __init__(self,
- system_prompt: str,
- action_schema: str,
- working_memory: WorkingMemory,
- system_metrics: SystemMetrics,
- llm: LlmEngine,
- validator: XMLValidator,
- parser: ResponseParser):
+ def __init__(
+ self,
+ system_prompt: str,
+ action_schema: str,
+ working_memory: WorkingMemory,
+ metrics: SystemMetrics,
+ llm: LlmEngine,
+ validator: XMLValidator,
+ parser: ResponseParser
+ ):
"""
Initialize web agent with required components.
"""
- super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser)
+ super().__init__(
+ system_prompt,
+ action_schema,
+ working_memory,
+ metrics, llm,
+ validator,
+ parser
+ )
self._response = ""
- self._system_prompt = system_prompt
self._state = WebAgentState.CONTEXT_APPROVAL
self._validation_error: Optional[str] = None
self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
@@ -148,7 +156,7 @@ class WebAgent(BaseAgent):
def _approve_context_thread(self) -> None:
self._set_response("")
- response_token_iter = self._llm.infer(f"{self._system_prompt}\n{self._action_schema}", self.context)
+ response_token_iter = self._llm.infer(self.system_prompt, self.context)
response = ""
for token in response_token_iter:
response += token
diff --git a/test/base_agent_test.py b/test/base_agent_test.py
index 69532a8..80c01c4 100644
--- a/test/base_agent_test.py
+++ b/test/base_agent_test.py
@@ -49,13 +49,14 @@ class MockEntry(Entry):
class TestBaseAgent(BaseAgent):
"""Concrete implementation of BaseAgent for testing."""
- def run(self) -> None:
- pass
+ pass
class BaseAgentTest(unittest.TestCase):
def setUp(self):
"""Set up test cases with mocked components."""
self.mock_llm = Mock(spec=LlmEngine)
+ self.mock_llm.token_count.return_value = 100
+ self.mock_llm.token_limit.return_value = 1000
self.mock_validator = Mock(spec=XMLValidator)
self.io_buffer = WebIOBuffer()
self.working_memory = WorkingMemory()
@@ -75,9 +76,10 @@ class BaseAgentTest(unittest.TestCase):
# Create test agent
self.agent = TestBaseAgent(
+ system_prompt="test prompt",
action_schema="test schema",
working_memory=self.working_memory,
- system_metrics=self.mock_metrics,
+ metrics=self.mock_metrics,
llm=self.mock_llm,
validator=self.mock_validator,
parser=self.parser
@@ -130,7 +132,7 @@ class BaseAgentTest(unittest.TestCase):
self.assertEqual(root.get("disk_total"), "10000")
# Check context size is 0 (empty memory)
- self.assertEqual(root.get("context"), "0")
+ self.assertEqual(root.get("context"), "10.0")
# Check no memory entries
self.assertEqual(len(list(root)), 0)
@@ -144,7 +146,7 @@ class BaseAgentTest(unittest.TestCase):
root = ET.fromstring(context)
# Check context size reflects one entry
- self.assertEqual(root.get("context"), "1")
+ self.assertEqual(root.get("context"), "10.0")
# Verify entry content
reasoning_elem = root.find("reasoning")
@@ -166,7 +168,7 @@ class BaseAgentTest(unittest.TestCase):
root = ET.fromstring(context)
# Check context size reflects three entries
- self.assertEqual(root.get("context"), "3")
+ self.assertEqual(root.get("context"), "10.0")
# Verify entry order maintained
reasoning_elems = root.findall("reasoning")
diff --git a/test/local_llm_engine_test.py b/test/local_llm_engine_test.py
index f78fce8..3887fee 100644
--- a/test/local_llm_engine_test.py
+++ b/test/local_llm_engine_test.py
@@ -7,20 +7,43 @@ from . import test_data
from sia.llm_engine import LlmEngine
from sia.local_llm_engine import LocalLlmEngine
-class LlmEngineTest(unittest.TestCase):
+class LocalLlmEngineTest(unittest.TestCase):
def setUp(self):
self.model_path = "/root/model"
def test_initialization(self):
- llm_engine = LocalLlmEngine(self.model_path)
+ llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
self.assertIsInstance(llm_engine, LlmEngine)
def test_infer(self):
main_context = "This is a test"
- llm_engine = LocalLlmEngine(self.model_path)
+ llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
tokens = llm_engine.infer(test_data.echo_system_prompt, main_context)
print_tokens, result_tokens = tee(tokens)
for token in print_tokens:
print(token, end="", flush=True)
result = ''.join(result_tokens)
self.assertEqual(result, f"{main_context}{main_context}")
+
+ def test_token_count(self):
+ """Test token counting returns reasonable numbers"""
+ llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
+
+ # Test with short inputs
+ count = llm_engine.token_count("Short prompt", "Brief context")
+ self.assertGreater(count, 5)
+ self.assertLess(count, 50)
+
+ # Test with longer inputs
+ long_prompt = "A detailed system prompt with multiple sentences. " * 5
+ long_context = "An extensive context containing various details. " * 10
+ count = llm_engine.token_count(long_prompt, long_context)
+ self.assertGreater(count, 100)
+ self.assertLess(count, 1000)
+
+ def test_token_limit(self):
+ """Test token limit is within expected range for modern LLMs"""
+ llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
+ limit = llm_engine.token_limit()
+ self.assertGreaterEqual(limit, 2048)
+ self.assertLessEqual(limit, 1048576)
\ No newline at end of file
diff --git a/test/web_agent_test.py b/test/web_agent_test.py
index 426c813..8f4df57 100644
--- a/test/web_agent_test.py
+++ b/test/web_agent_test.py
@@ -60,13 +60,15 @@ class WebAgentTest(unittest.TestCase):
yield "test reasoning"
self.mock_llm.infer.side_effect = mock_infer
+ self.mock_llm.token_count.return_value = 100
+ self.mock_llm.token_limit.return_value = 1000
# Create agent with all components
self.agent = WebAgent(
system_prompt="test prompt",
action_schema="test schema",
working_memory=self.working_memory,
- system_metrics=self.mock_metrics,
+ metrics=self.mock_metrics,
llm=self.mock_llm,
validator=self.mock_validator,
parser=self.parser
diff --git a/test/web_socket_manager_test.py b/test/web_socket_manager_test.py
index e92680f..9b3cc96 100644
--- a/test/web_socket_manager_test.py
+++ b/test/web_socket_manager_test.py
@@ -31,7 +31,9 @@ class WebSocketManagerTest(AioHTTPTestCase):
"memory_used": 1000,
"memory_total": 2000,
"disk_used": 5000,
- "disk_total": 10000
+ "disk_total": 10000,
+ "token_count": 100,
+ "token_limit": 1000
}
# Create minimal mocks for other components
@@ -41,6 +43,8 @@ class WebSocketManagerTest(AioHTTPTestCase):
self.mock_llm.infer.side_effect = mock_infer
self.mock_validator = Mock(spec=XMLValidator)
self.mock_validator.validate.return_value = None
+ self.mock_llm.token_count.return_value = 100
+ self.mock_llm.token_limit.return_value = 1000
# Create parser with real IO buffer
self.parser = ResponseParser(self.io_buffer)
@@ -50,7 +54,7 @@ class WebSocketManagerTest(AioHTTPTestCase):
system_prompt="test prompt",
action_schema="test schema",
working_memory=self.working_memory,
- system_metrics=self.mock_metrics,
+ metrics=self.mock_metrics,
llm=self.mock_llm,
validator=self.mock_validator,
parser=self.parser