Fixed context usage calculation

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