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

Fixed context usage calculation

Browse Source
This commit is contained in:
Niels Geens
2024-11-14 18:23:33 +01:00
parent ede0a642d3
commit 4ce421bbce
13 changed files with 249 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
requirements.txt
View File

@@ -3,5 +3,6 @@ aiohttp
bs4
openai
python-dotenv
tiktoken
torch
transformers

30
sia/__main__.py
View File

@@ -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 "<reasoning>"
time.sleep(2)
yield "test reasoning"
time.sleep(2)
yield "</reasoning>"
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)

46
sia/base_agent.py
View File

@@ -18,28 +18,37 @@ 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:
"""
Compile the current context for LLM inference.
@@ -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()))
context.set("context", "100")
# Add memory entries
for entry in memory_context:
context.append(entry)
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)

11
sia/config.py
View File

@@ -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:
@@ -146,3 +152,8 @@ class Config:
def temperature(self) -> float:
"""LLM temperature parameter."""
return self.args.temperature
@property
def token_limit(self) -> int:
"""Token limit for the LLM."""
return self.args.token_limit

64
sia/hf_llm_engine.py
View File

@@ -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)
self._model = model
self._temperature = temperature
# Generation parameters
self.temperature = temperature
self.max_new_tokens = max_new_tokens
def set_model_path(self, model_id: str):
"""
Update the model being used.
Args:
model_id: New HuggingFace model ID to use
"""
self.model_id = model_id
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,17 +41,18 @@ 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
)
@@ -69,3 +60,22 @@ class HfLlmEngine(LlmEngine):
if content := response.choices[0].delta.content:
yield content
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

9
sia/llm_engine.py
View File

@@ -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

70
sia/local_llm_engine.py
View File

@@ -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

38
sia/openai_llm_engine.py
View File

@@ -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
self._temperature = temperature
self._token_limit = token_limit
# Initialize OpenAI client
self.client = openai.Client(
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
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

30
sia/web_agent.py
View File

@@ -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

14
test/base_agent_test.py
View File

@@ -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")

29
test/local_llm_engine_test.py
View File

@@ -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}<test_tag>{main_context}</test_tag>")
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)

4
test/web_agent_test.py
View File

@@ -60,13 +60,15 @@ class WebAgentTest(unittest.TestCase):
yield "<reasoning>test reasoning</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

8
test/web_socket_manager_test.py
View File

@@ -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