From fe4184962c6eb7fc5d13d131f3b0bda1f90e28c4 Mon Sep 17 00:00:00 2001 From: Gal Kleinman Date: Mon, 28 Aug 2023 17:05:21 +0300 Subject: [PATCH] fix: context propagation on acompletion method --- litellm/main.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 6e26bf3e5..001a6c546 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1,7 +1,7 @@ import os, openai, sys from typing import Any from functools import partial -import dotenv, traceback, random, asyncio, time +import dotenv, traceback, random, asyncio, time, contextvars from copy import deepcopy import litellm from litellm import ( # type: ignore @@ -49,8 +49,12 @@ async def acompletion(*args, **kwargs): # Use a partial function to pass your keyword arguments func = partial(completion, *args, **kwargs) + # Add the context to the function + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + # Call the synchronous function using run_in_executor - return await loop.run_in_executor(None, func) + return await loop.run_in_executor(None, func_with_context) @client