From 4ce421bbce2a4ed4c9ede97f0b5dc11756da33c1 Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Thu, 14 Nov 2024 18:23:33 +0100 Subject: [PATCH] Fixed context usage calculation --- requirements.txt | 1 + sia/__main__.py | 30 ++++++-------- sia/base_agent.py | 52 +++++++++++++++--------- sia/config.py | 13 +++++- sia/hf_llm_engine.py | 66 ++++++++++++++++++------------- sia/llm_engine.py | 9 ++++- sia/local_llm_engine.py | 70 +++++++++++++++++++++++---------- sia/openai_llm_engine.py | 42 +++++++++++++++----- sia/web_agent.py | 30 ++++++++------ test/base_agent_test.py | 14 ++++--- test/local_llm_engine_test.py | 29 ++++++++++++-- test/web_agent_test.py | 4 +- test/web_socket_manager_test.py | 8 +++- 13 files changed, 249 insertions(+), 119 deletions(-) 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