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

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)

52
sia/base_agent.py
View File

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

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

66
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)
# 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()
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

42
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
# 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

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