Fixed context usage calculation

requirements.txt

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

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 "<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,
+                    self._config.model,
-                    api_token=self._config.api_token,
+                    self._config.temperature,
-                    temperature=self._config.temperature
+                    self._config.api_token,
                 )
             case "openai":
                 self._llm = OpenAILlmEngine(
-                    model=self._config.model,
+                    self._config.model,
-                    api_key=self._config.api_token
+                    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)

sia/base_agent.py

@@ -18,27 +18,36 @@ class BaseAgent(ABC):
     and coordinating components for LLM inference.
     """
     
-    def __init__(self, 
+    def __init__(
-                 action_schema: str, 
+            self,
-                 working_memory: WorkingMemory,
+            system_prompt: str,
-                 system_metrics: SystemMetrics,
+            action_schema: str, 
-                 llm: LlmEngine,
+            working_memory: WorkingMemory,
-                 validator: XMLValidator,
+            metrics: SystemMetrics,
-                 parser: ResponseParser):
+            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()))
-        
+        context.set("context", "100")
-        # Add memory entries
+
         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)

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
+        return self.args.temperature
+    
+    @property
+    def token_limit(self) -> int:
+        """Token limit for the LLM."""
+        return self.args.token_limit

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",
+        model: str,
-        api_token: Optional[str] = None,
+        temperature: float,
-        temperature: float = 0.7,
+        api_token: Optional[str],
-        max_new_tokens: int = 1024,
     ):
         """
         Initialize the HuggingFace Inference API LLM Engine.
 
         Args:
-            model_id: HuggingFace model ID to use (default: Mistral-7B-Instruct)
+            model: HuggingFace model ID to use
-            api_token: HuggingFace API token. If None, will try to read from HF_TOKEN env var
+            temperature: Sampling temperature
-            temperature: Sampling temperature (default: 0.7)
+            api_token: HuggingFace API token
-            max_new_tokens: Maximum number of tokens to generate (default: 1024)
         """
-        self.model_id = model_id
+        self._model = model
-        self.client = InferenceClient(token=api_token)
+        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:
+        self._tokenizer = AutoTokenizer.from_pretrained(model, token=api_token)
-            model_id: New HuggingFace model ID to use
+        self._config = AutoConfig.from_pretrained(model, token=api_token)
-        """
+        self._client = InferenceClient(token=api_token)
-        self.model_id = model_id
                         
     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(
+            stream = self._client.chat_completion(
-                model=self.model_id,
+                model=self._model,
                 messages=messages,
-                temperature=self.temperature,
+                temperature=self._temperature,
-                max_tokens=self.max_new_tokens,
                 stream=True
             )
             
             for response in stream:
                 if content := response.choices[0].delta.content:
                     yield content
-        return stream_wrapper()
+        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

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

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)
+        self._temperature = temperature
-
+        self._token_limit = token_limit
-    def set_model_path(self, model_path: str):
+        self._tokenizer = AutoTokenizer.from_pretrained(model_path)
-        """
-        Load the model from the specified path.
-
-        Args:
-            model_path: Path to the model weights to load.
-        """
-        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:
+        if self._tokenizer.pad_token_id is None:
-            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+            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

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)
+            model: OpenAI model to use
-            api_key: OpenAI API key. If None, will try to read from OPENAI_API_KEY env var
+            temperature: Temperature for sampling
+            api_key: OpenAI API key
+            token_limit: Maximum number of tokens to generate
         """
         self._model = model
-        
+        self._temperature = temperature
-        # Initialize OpenAI client
+        self._token_limit = token_limit
-        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
+                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

sia/web_agent.py

@@ -33,20 +33,28 @@ class WebAgent(BaseAgent):
     Broadcasts state changes to registered handlers.
     """
     
-    def __init__(self,
+    def __init__(
-                system_prompt: str,
+            self,
-                action_schema: str, 
+            system_prompt: str,
-                working_memory: WorkingMemory,
+            action_schema: str, 
-                system_metrics: SystemMetrics,
+            working_memory: WorkingMemory,
-                llm: LlmEngine,
+            metrics: SystemMetrics,
-                validator: XMLValidator,
+            llm: LlmEngine,
-                parser: ResponseParser):
+            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

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

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}<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)

test/web_agent_test.py

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

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