@samitouri / QOSami_Agents_Zero-2 / commits / 76547079

rate limiter polishing, file download polishing

frdel committed Jul 31, 2025 at 14:46 UTC 765470796304a044ff17fd240b5a848d44e4e081

7 files changed +19 -90

agent.py

+8 -51

@@ -613,15 +613,6 @@ class Agent:

613 ):

614 model = self.get_utility_model()

615

616 - # rate limiter

617 - limiter = await self.rate_limiter(

618 - self.config.utility_model, f"SYSTEM: {system}\nUSER: {message}", background

619 - )

620 -

621 - # add output tokens to rate limiter in tokens callback

622 - async def tokens_callback(delta: str, tokens: int):

623 - await self.handle_intervention()

624 - limiter.add(output=tokens)

616

617 # propagate stream to callback if set

618 async def stream_callback(chunk: str, total: str):

@@ -632,7 +623,7 @@ class Agent:

623 system_message=system,

624 user_message=message,

625 response_callback=stream_callback,

635 - tokens_callback=tokens_callback,

626 + rate_limiter_callback=self.rate_limiter_callback if not background else None,

627 )

628

629 return response

@@ -642,63 +633,29 @@ class Agent:

633 messages: list[BaseMessage],

634 response_callback: Callable[[str, str], Awaitable[None]] | None = None,

635 reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,

636 + background: bool = False,

637 ):

638 response = ""

639

640 # model class

641 model = self.get_chat_model()

642

651 - # rate limiter

652 - limiter = await self.rate_limiter(

653 - self.config.chat_model, ChatPromptTemplate.from_messages(messages).format()

654 - )

655 -

656 - # add output tokens to rate limiter in tokens callback

657 - async def tokens_callback(delta: str, tokens: int):

658 - await self.handle_intervention()

659 - limiter.add(output=tokens)

660 -

643 # call model

644 response, reasoning = await model.unified_call(

645 messages=messages,

646 reasoning_callback=reasoning_callback,

647 response_callback=response_callback,

666 - tokens_callback=tokens_callback,

648 + rate_limiter_callback=self.rate_limiter_callback if not background else None,

649 )

650

651 return response, reasoning

652

671 - async def rate_limiter(

672 - self, model_config: models.ModelConfig, input: str, background: bool = False

653 + async def rate_limiter_callback(

654 + self, message:str, key:str, total:int, limit:int

655 ):

674 - # rate limiter log

675 - wait_log = None

676 -

677 - async def wait_callback(msg: str, key: str, total: int, limit: int):

678 - nonlocal wait_log

679 - if not wait_log:

680 - wait_log = self.context.log.log(

681 - type="util",

682 - update_progress="none",

683 - heading=msg,

684 - model=f"{model_config.provider}\\{model_config.name}",

685 - )

686 - wait_log.update(heading=msg, key=key, value=total, limit=limit)

687 - if not background:

688 - self.context.log.set_progress(msg, -1)

689 -

690 - # rate limiter

691 - limiter = models.get_rate_limiter(

692 - model_config.provider,

693 - model_config.name,

694 - model_config.limit_requests,

695 - model_config.limit_input,

696 - model_config.limit_output,

697 - )

698 - limiter.add(input=tokens.approximate_tokens(input))

699 - limiter.add(requests=1)

700 - await limiter.wait(callback=wait_callback)

701 - return limiter

656 + # show the rate limit waiting in a progress bar, no need to spam the chat history

657 + self.context.log.set_progress(message, True)

658 + return False

659

660 async def handle_intervention(self, progress: str = ""):

661 while self.context.paused:

models.py

+6 -5

@@ -108,7 +108,7 @@ 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):

111 +async def apply_rate_limiter(model_config: ModelConfig|None, input_text: str, rate_limiter_callback: Callable[[str, str, int, int], Awaitable[bool]] | None = None):

112 if not model_config:

113 return

114 limiter = get_rate_limiter(

@@ -120,15 +120,15 @@ async def apply_rate_limiter(model_config: ModelConfig|None, input_text: str):

120 )

121 limiter.add(input=approximate_tokens(input_text))

122 limiter.add(requests=1)

123 - await limiter.wait()

123 + await limiter.wait(rate_limiter_callback)

124 return limiter

125

126 -def apply_rate_limiter_sync(model_config: ModelConfig|None, input_text: str):

126 +def apply_rate_limiter_sync(model_config: ModelConfig|None, input_text: str, rate_limiter_callback: Callable[[str, str, int, int], Awaitable[bool]] | None = None):

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

131 + return asyncio.run(apply_rate_limiter(model_config, input_text, rate_limiter_callback))

132

133

134 class LiteLLMChatWrapper(SimpleChatModel):

@@ -286,6 +286,7 @@ class LiteLLMChatWrapper(SimpleChatModel):

286 response_callback: Callable[[str, str], Awaitable[None]] | None = None,

287 reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,

288 tokens_callback: Callable[[str, int], Awaitable[None]] | None = None,

289 + rate_limiter_callback: Callable[[str, str, int, int], Awaitable[bool]] | None = None,

290 **kwargs: Any,

291 ) -> Tuple[str, str]:

292

@@ -303,7 +304,7 @@ class LiteLLMChatWrapper(SimpleChatModel):

304 msgs_conv = self._convert_messages(messages)

305

306 # Apply rate limiting if configured

306 - limiter = await apply_rate_limiter(self.a0_model_conf, str(msgs_conv))

307 + limiter = await apply_rate_limiter(self.a0_model_conf, str(msgs_conv), rate_limiter_callback)

308

309 # call model

310 _completion = await acompletion(

python/api/get_work_dir_files.py

+1 -3

@@ -1,8 +1,6 @@

1 from python.helpers.api import ApiHandler, Request, Response

2 from python.helpers.file_browser import FileBrowser

3 from python.helpers import runtime

4 -import python.api.get_work_dir_files as get_work_dir_files_module

5 -

5 class GetWorkDirFiles(ApiHandler):

@@ -21,7 +19,7 @@ class GetWorkDirFiles(ApiHandler):

20 # browser = FileBrowser()

21 # result = browser.get_files(current_path)

24 - result = await runtime.call_development_function(get_work_dir_files_module.get_files, current_path)

22 + result = await runtime.call_development_function(get_files, current_path)

24 return {"data": result}

python/helpers/document_query.py

-6

@@ -140,13 +140,7 @@ class DocumentQueryStore:

140 PrintStyle.error(f"No chunks created for document: {document_uri}")

141 return False, []

142

143 - # Apply rate limiter

143 try:

145 - docs_text = "".join(chunk.page_content for chunk in docs)

146 - await self.agent.rate_limiter(

147 - model_config=self.agent.config.embeddings_model, input=docs_text

148 - )

149 -

144 # Initialize vector db if not already initialized

145 if not self.vector_db:

146 self.vector_db = self.init_vector_db()

python/helpers/memory.py

-11

@@ -300,11 +300,6 @@ class Memory:

300 ):

301 comparator = Memory._get_comparator(filter) if filter else None

302

303 - # rate limiter

304 - await self.agent.rate_limiter(

305 - model_config=self.agent.config.embeddings_model, input=query

306 - )

307 -

303 return await self.db.asearch(

304 query,

305 search_type="similarity_score_threshold",

@@ -376,12 +371,6 @@ class Memory:

371 if not doc.metadata.get("area", ""):

372 doc.metadata["area"] = Memory.Area.MAIN.value

373

379 - # rate limiter

380 - docs_txt = "".join(self.format_docs_plain(docs))

381 - await self.agent.rate_limiter(

382 - model_config=self.agent.config.embeddings_model, input=docs_txt

383 - )

384 -

374 await self.db.aadd_documents(documents=docs, ids=ids)

375 self._save_db() # persist

376 return ids

python/helpers/rate_limiter.py

+4 -3

@@ -32,7 +32,7 @@ class RateLimiter:

33 async def wait(

34 self,

35 - callback: Callable[[str, str, int, int], Awaitable[None]] | None = None,

35 + callback: Callable[[str, str, int, int], Awaitable[bool]] | None = None,

36 ):

37 while True:

38 await self.cleanup()

@@ -46,8 +46,9 @@ class RateLimiter:

46 if total > limit:

47 if callback:

48 msg = f"Rate limit exceeded for {key} ({total}/{limit}), waiting..."

49 - await callback(msg, key, total, limit)

50 - should_wait = True

49 + should_wait = not await callback(msg, key, total, limit)

50 + else:

51 + should_wait = True

52 break

54 if not should_wait:

python/helpers/vector_db.py

-11

@@ -77,11 +77,6 @@ class VectorDB:

77 ):

78 comparator = get_comparator(filter) if filter else None

80 - # rate limiter

81 - await self.agent.rate_limiter(

82 - model_config=self.agent.config.embeddings_model, input=query

83 - )

84 -

80 return await self.db.asearch(

81 query,

82 search_type="similarity_score_threshold",

@@ -109,12 +104,6 @@ class VectorDB:

104 for doc, id in zip(docs, ids):

105 doc.metadata["id"] = id # add ids to documents metadata

106

112 - # rate limiter

113 - docs_txt = "".join(format_docs_plain(docs))

114 - await self.agent.rate_limiter(

115 - model_config=self.agent.config.embeddings_model, input=docs_txt

116 - )

117 -

107 self.db.add_documents(documents=docs, ids=ids)

108 return ids

109