rate limiter polishing
frdel committed
Jul 31, 2025 at 13:31 UTC
27b2a691fdc0212ba2620520744a8e7eaee15ab1
4 files changed
+108
-133
.gitignore
+1
-3
@@ -39,6 +39,4 @@ instruments/**
39
40
# Global rule to include .gitkeep files anywhere
41
!**/.gitkeep
42
-agent_history.gif
43
-
44
-tests/**
42
+agent_history.gif
\ No newline at end of file
models.py
+39
-130
@@ -22,7 +22,6 @@ from python.helpers.dotenv import load_dotenv
22
from python.helpers.providers import get_provider_config
23
from python.helpers.rate_limiter import RateLimiter
24
from python.helpers.tokens import approximate_tokens
25
-import asyncio
25
26
from langchain_core.language_models.chat_models import SimpleChatModel
27
from langchain_core.outputs.chat_generation import ChatGenerationChunk
@@ -109,6 +108,28 @@ def get_rate_limiter(
108
limiter.limits["output"] = output or 0
109
return limiter
110
111
+async def apply_rate_limiter(model_config: ModelConfig|None, input_text: str):
112
+ if not model_config:
113
+ return
114
+ limiter = get_rate_limiter(
115
+ model_config.provider,
116
+ model_config.name,
117
+ model_config.limit_requests,
118
+ model_config.limit_input,
119
+ model_config.limit_output,
120
+ )
121
+ limiter.add(input=approximate_tokens(input_text))
122
+ limiter.add(requests=1)
123
+ await limiter.wait()
124
+ return limiter
125
+
126
+def apply_rate_limiter_sync(model_config: ModelConfig|None, input_text: str):
127
+ if not model_config:
128
+ return
129
+ import asyncio, nest_asyncio
130
+ nest_asyncio.apply()
131
+ return asyncio.run(apply_rate_limiter(model_config, input_text))
132
+
133
134
class LiteLLMChatWrapper(SimpleChatModel):
135
model_name: str
@@ -123,30 +144,13 @@ class LiteLLMChatWrapper(SimpleChatModel):
144
def __init__(self, model: str, provider: str, model_config: Optional[ModelConfig] = None, **kwargs: Any):
145
model_value = f"{provider}/{model}"
146
super().__init__(model_name=model_value, provider=provider, kwargs=kwargs) # type: ignore
126
- # Set rate limiting config as instance attribute after parent init
127
- self.rate_limit_config = model_config
147
+ # Set A0 model config as instance attribute after parent init
148
+ self.a0_model_conf = model_config
149
150
@property
151
def _llm_type(self) -> str:
152
return "litellm-chat"
153
133
- async def _apply_rate_limiting(self, input_text: str, background: bool = False):
134
- """Apply rate limiting if rate_limit_config is provided"""
135
- if not self.rate_limit_config:
136
- return None
137
-
138
- limiter = get_rate_limiter(
139
- self.rate_limit_config.provider,
140
- self.rate_limit_config.name,
141
- self.rate_limit_config.limit_requests,
142
- self.rate_limit_config.limit_input,
143
- self.rate_limit_config.limit_output,
144
- )
145
- limiter.add(input=approximate_tokens(input_text))
146
- limiter.add(requests=1)
147
- await limiter.wait()
148
- return limiter
149
-
154
def _convert_messages(self, messages: List[BaseMessage]) -> List[dict]:
155
result = []
156
# Map LangChain message types to LiteLLM roles
@@ -207,18 +211,14 @@ class LiteLLMChatWrapper(SimpleChatModel):
211
msgs = self._convert_messages(messages)
212
213
# Apply rate limiting if configured
210
- if self.rate_limit_config:
211
- try:
212
- # Convert messages to text for token counting
213
- messages_text = str(msgs)
214
- asyncio.run(self._apply_rate_limiting(messages_text))
215
- except Exception as e:
216
- # Don't let rate limiting break the model call
217
- print(f"Rate limiting warning: {e}")
214
+ apply_rate_limiter_sync(self.a0_model_conf, str(msgs))
215
216
+ # Call the model
217
resp = completion(
218
model=self.model_name, messages=msgs, stop=stop, **{**self.kwargs, **kwargs}
219
)
220
+
221
+ # Parse output
222
parsed = _parse_chunk(resp)
223
return parsed["response_delta"]
224
@@ -234,14 +234,7 @@ class LiteLLMChatWrapper(SimpleChatModel):
234
msgs = self._convert_messages(messages)
235
236
# Apply rate limiting if configured
237
- if self.rate_limit_config:
238
- try:
239
- # Convert messages to text for token counting
240
- messages_text = str(msgs)
241
- asyncio.run(self._apply_rate_limiting(messages_text))
242
- except Exception as e:
243
- # Don't let rate limiting break the model call
244
- print(f"Rate limiting warning: {e}")
237
+ apply_rate_limiter_sync(self.a0_model_conf, str(msgs))
238
239
for chunk in completion(
240
model=self.model_name,
@@ -267,10 +260,8 @@ class LiteLLMChatWrapper(SimpleChatModel):
260
msgs = self._convert_messages(messages)
261
262
# Apply rate limiting if configured
270
- if self.rate_limit_config:
271
- # Convert messages to text for token counting
272
- messages_text = str(msgs)
273
- await self._apply_rate_limiting(messages_text)
263
+ await apply_rate_limiter(self.a0_model_conf, str(msgs))
264
+
265
266
response = await acompletion(
267
model=self.model_name,
@@ -312,11 +303,7 @@ class LiteLLMChatWrapper(SimpleChatModel):
303
msgs_conv = self._convert_messages(messages)
304
305
# Apply rate limiting if configured
315
- limiter = None
316
- if self.rate_limit_config:
317
- # Convert messages to text for token counting
318
- messages_text = str(msgs_conv)
319
- limiter = await self._apply_rate_limiting(messages_text)
306
+ limiter = await apply_rate_limiter(self.a0_model_conf, str(msgs_conv))
307
308
# call model
309
_completion = await acompletion(
@@ -402,54 +389,16 @@ class BrowserCompatibleChatWrapper(LiteLLMChatWrapper):
389
class LiteLLMEmbeddingWrapper(Embeddings):
390
model_name: str
391
kwargs: dict = {}
405
- rate_limit_config: Optional[ModelConfig] = None
392
+ a0_model_conf: Optional[ModelConfig] = None
393
394
def __init__(self, model: str, provider: str, model_config: Optional[ModelConfig] = None, **kwargs: Any):
395
self.model_name = f"{provider}/{model}" if provider != "openai" else model
396
self.kwargs = kwargs
410
- self.rate_limit_config = model_config
397
+ self.a0_model_conf = model_config
398
412
- def _apply_rate_limiting_sync(self, input_text: str):
413
- """Apply rate limiting synchronously if model_config is provided"""
414
- if not self.rate_limit_config:
415
- return None
416
-
417
- limiter = get_rate_limiter(
418
- self.rate_limit_config.provider,
419
- self.rate_limit_config.name,
420
- self.rate_limit_config.limit_requests,
421
- self.rate_limit_config.limit_input,
422
- self.rate_limit_config.limit_output,
423
- )
424
- limiter.add(input=approximate_tokens(input_text))
425
- limiter.add(requests=1)
426
- # Note: Embeddings typically don't have streaming, so we do synchronous rate limiting
427
- import asyncio
428
- try:
429
- asyncio.run(limiter.wait())
430
- except RuntimeError:
431
- # If we're already in an event loop, create a new thread
432
- import threading
433
- import concurrent.futures
434
-
435
- def wait_for_limiter():
436
- new_loop = asyncio.new_event_loop()
437
- asyncio.set_event_loop(new_loop)
438
- try:
439
- new_loop.run_until_complete(limiter.wait())
440
- finally:
441
- new_loop.close()
442
-
443
- with concurrent.futures.ThreadPoolExecutor() as executor:
444
- executor.submit(wait_for_limiter).result()
445
- return limiter
446
-
399
def embed_documents(self, texts: List[str]) -> List[List[float]]:
400
# Apply rate limiting if configured
449
- if self.rate_limit_config:
450
- # Convert texts to combined string for token counting
451
- texts_combined = " ".join(texts)
452
- self._apply_rate_limiting_sync(texts_combined)
401
+ apply_rate_limiter_sync(self.a0_model_conf, " ".join(texts))
402
403
resp = embedding(model=self.model_name, input=texts, **self.kwargs)
404
return [
@@ -459,8 +408,7 @@ class LiteLLMEmbeddingWrapper(Embeddings):
408
409
def embed_query(self, text: str) -> List[float]:
410
# Apply rate limiting if configured
462
- if self.rate_limit_config:
463
- self._apply_rate_limiting_sync(text)
411
+ apply_rate_limiter_sync(self.a0_model_conf, text)
412
413
resp = embedding(model=self.model_name, input=[text], **self.kwargs)
414
item = resp.data[0] # type: ignore
@@ -480,57 +428,18 @@ class LocalSentenceTransformerWrapper(Embeddings):
428
429
self.model = SentenceTransformer(model, **kwargs)
430
self.model_name = model
483
- self.rate_limit_config = model_config
431
+ self.a0_model_conf = model_config
432
485
- def _apply_rate_limiting_sync(self, input_text: str):
486
- """Apply rate limiting synchronously if model_config is provided"""
487
- if not self.rate_limit_config:
488
- return None
489
-
490
- limiter = get_rate_limiter(
491
- self.rate_limit_config.provider,
492
- self.rate_limit_config.name,
493
- self.rate_limit_config.limit_requests,
494
- self.rate_limit_config.limit_input,
495
- self.rate_limit_config.limit_output,
496
- )
497
- limiter.add(input=approximate_tokens(input_text))
498
- limiter.add(requests=1)
499
- # Note: Local models still respect rate limiting for consistency
500
- import asyncio
501
- try:
502
- asyncio.run(limiter.wait())
503
- except RuntimeError:
504
- # If we're already in an event loop, create a new thread
505
- import threading
506
- import concurrent.futures
507
-
508
- def wait_for_limiter():
509
- new_loop = asyncio.new_event_loop()
510
- asyncio.set_event_loop(new_loop)
511
- try:
512
- new_loop.run_until_complete(limiter.wait())
513
- finally:
514
- new_loop.close()
515
-
516
- with concurrent.futures.ThreadPoolExecutor() as executor:
517
- executor.submit(wait_for_limiter).result()
518
- return limiter
519
-
433
def embed_documents(self, texts: List[str]) -> List[List[float]]:
434
# Apply rate limiting if configured
522
- if self.rate_limit_config:
523
- # Convert texts to combined string for token counting
524
- texts_combined = " ".join(texts)
525
- self._apply_rate_limiting_sync(texts_combined)
435
+ apply_rate_limiter_sync(self.a0_model_conf, " ".join(texts))
436
437
embeddings = self.model.encode(texts, convert_to_tensor=False) # type: ignore
438
return embeddings.tolist() if hasattr(embeddings, "tolist") else embeddings # type: ignore
439
440
def embed_query(self, text: str) -> List[float]:
441
# Apply rate limiting if configured
532
- if self.rate_limit_config:
533
- self._apply_rate_limiting_sync(text)
442
+ apply_rate_limiter_sync(self.a0_model_conf, text)
443
444
embedding = self.model.encode([text], convert_to_tensor=False) # type: ignore
445
result = (
python/helpers/settings.py
+36
@@ -48,6 +48,9 @@ class Settings(TypedDict):
48
browser_model_name: str
49
browser_model_api_base: str
50
browser_model_vision: bool
51
+ browser_model_rl_requests: int
52
+ browser_model_rl_input: int
53
+ browser_model_rl_output: int
54
browser_model_kwargs: dict[str, str]
55
56
agent_profile: str
@@ -447,6 +450,36 @@ def convert_out(settings: Settings) -> SettingsOutput:
450
}
451
)
452
453
+ browser_model_fields.append(
454
+ {
455
+ "id": "browser_model_rl_requests",
456
+ "title": "Web Browser model rate limit requests",
457
+ "description": "Rate limit requests for web browser model.",
458
+ "type": "number",
459
+ "value": settings["browser_model_rl_requests"],
460
+ }
461
+ )
462
+
463
+ browser_model_fields.append(
464
+ {
465
+ "id": "browser_model_rl_input",
466
+ "title": "Web Browser model rate limit input",
467
+ "description": "Rate limit input for web browser model.",
468
+ "type": "number",
469
+ "value": settings["browser_model_rl_input"],
470
+ }
471
+ )
472
+
473
+ browser_model_fields.append(
474
+ {
475
+ "id": "browser_model_rl_output",
476
+ "title": "Web Browser model rate limit output",
477
+ "description": "Rate limit output for web browser model.",
478
+ "type": "number",
479
+ "value": settings["browser_model_rl_output"],
480
+ }
481
+ )
482
+
483
browser_model_fields.append(
484
{
485
"id": "browser_model_kwargs",
@@ -1199,6 +1232,9 @@ def get_default_settings() -> Settings:
1232
browser_model_name="openai/gpt-4.1",
1233
browser_model_api_base="",
1234
browser_model_vision=True,
1235
+ browser_model_rl_requests=0,
1236
+ browser_model_rl_input=0,
1237
+ browser_model_rl_output=0,
1238
browser_model_kwargs={"temperature": "0"},
1239
memory_recall_enabled=True,
1240
memory_recall_interval=3,
tests/rate_limiter_test.py
new
+32
@@ -0,0 +1,32 @@
1
+
2
+import sys
3
+import os
4
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
5
+import models
6
+
7
+provider = "openai"
8
+name = "gpt-4.1-mini"
9
+
10
+model = models.get_chat_model(
11
+ provider=provider,
12
+ name=name,
13
+ model_config=models.ModelConfig(
14
+ type=models.ModelType.CHAT,
15
+ provider=provider,
16
+ name=name,
17
+ limit_requests = 5,
18
+ limit_input = 15000,
19
+ limit_output = 1000,
20
+ )
21
+ )
22
+
23
+async def run():
24
+ response, reasoning = await model.unified_call(
25
+ user_message="Tell me a joke"
26
+ )
27
+ print("Response: ", response)
28
+ print("Reasoning: ", reasoning)
29
+
30
+
31
+import asyncio
32
+asyncio.run(run())
\ No newline at end of file