MCPcopy Create free account
hub / github.com/Open-Bee/DataStudio / _generate_async

Method _generate_async

datastudio/models/openai_api.py:379–488  ·  view source on GitHub ↗

Run async generation over all messages. Creates a shared :class:`aiohttp.ClientSession` with connection pooling and fires all requests via ``asyncio.gather``. Args: processed_messages: List of pre-processed messages (OpenAI format). all_kwargs: Addit

(
        self,
        processed_messages: List[Any],
        all_kwargs: dict,
    )

Source from the content-addressed store, hash-verified

377 return self.fail_msg
378
379 async def _generate_async(
380 self,
381 processed_messages: List[Any],
382 all_kwargs: dict,
383 ) -> List[str]:
384 """Run async generation over all messages.
385
386 Creates a shared :class:`aiohttp.ClientSession` with connection
387 pooling and fires all requests via ``asyncio.gather``.
388
389 Args:
390 processed_messages: List of pre-processed messages (OpenAI format).
391 all_kwargs: Additional generation arguments.
392
393 Returns:
394 List of response strings.
395 """
396 # Configure timeout
397 connect_timeout, read_timeout = self.timeout
398 timeout = aiohttp.ClientTimeout(
399 total=None,
400 sock_connect=connect_timeout,
401 sock_read=read_timeout,
402 )
403
404 # Configure connector with connection pooling
405 pool_size = min(self.max_connections, self.concurrency_limit)
406 connector = aiohttp.TCPConnector(
407 limit=pool_size,
408 ttl_dns_cache=300,
409 force_close=False,
410 keepalive_timeout=300,
411 )
412
413 total = len(processed_messages)
414 if self.logger:
415 self.logger.info(
416 f"[OpenAI API] Processing {total} requests with concurrency={self.concurrency_limit}, "
417 f"connections={pool_size}"
418 )
419
420 error_counter = {"ok": 0, "active": 0, "timeout": 0, "error": 0, "fail": 0}
421 pbar = async_tqdm(total=total, desc="Processing API Requests", leave=True)
422 pbar.set_postfix(ok=0, active=0, timeout=0, error=0, fail=0)
423
424 # Warm-up: gradually increase concurrency to avoid TCP connection storm
425 initial_concurrency = min(64, self.concurrency_limit)
426 need_warmup = (
427 self.concurrency_limit > initial_concurrency and total > initial_concurrency
428 )
429
430 if need_warmup:
431 semaphore = asyncio.Semaphore(initial_concurrency)
432 else:
433 semaphore = asyncio.Semaphore(self.concurrency_limit)
434
435 async with aiohttp.ClientSession(
436 timeout=timeout, connector=connector

Callers 2

_runMethod · 0.95
generateMethod · 0.95

Calls 4

infoMethod · 0.80
closeMethod · 0.80
warningMethod · 0.80

Tested by

no test coverage detected