mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-11 12:20:24 +02:00
0c29743538
* use max_tokens to do better rate limit handling * fix unti tests * address greptile comment, thanks greptile
21 lines
711 B
Python
21 lines
711 B
Python
import multiprocessing
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
from typing import TypeVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def run_with_timeout_multiproc(
|
|
task: Callable[..., T], timeout: int, kwargs: dict[str, Any]
|
|
) -> T:
|
|
# Use multiprocessing to prevent a thread from blocking the main thread
|
|
with multiprocessing.Pool(processes=1) as pool:
|
|
async_result = pool.apply_async(task, kwds=kwargs)
|
|
try:
|
|
# Wait at most timeout seconds for the function to complete
|
|
result = async_result.get(timeout=timeout)
|
|
return result
|
|
except multiprocessing.TimeoutError:
|
|
raise TimeoutError(f"Function timed out after {timeout} seconds")
|