fix(utils.py): don't run post-call rules on a coroutine function

This commit is contained in:
Krrish Dholakia 2024-03-27 13:16:27 -07:00
parent d08da5b05a
commit 9b7383ac67
2 changed files with 28 additions and 5 deletions

View file

@ -127,3 +127,14 @@ def test_post_call_rule_streaming():
print(type(e))
print(vars(e))
assert e.message == "This violates LiteLLM Proxy Rules. Response too short"
def test_post_call_processing_error_async_response():
response = asyncio.run(
acompletion(
model="command-nightly", # Just used as an example
messages=[{"content": "Hello, how are you?", "role": "user"}],
api_base="https://openai-proxy.berriai.repl.co", # Just used as an example
custom_llm_provider="openai",
)
)

View file

@ -2465,6 +2465,14 @@ def client(original_function):
)
raise e
def check_coroutine(value) -> bool:
if inspect.iscoroutine(value):
return True
elif inspect.iscoroutinefunction(value):
return True
else:
return False
def post_call_processing(original_response, model):
try:
if original_response is None:
@ -2475,11 +2483,15 @@ def client(original_function):
call_type == CallTypes.completion.value
or call_type == CallTypes.acompletion.value
):
model_response = original_response["choices"][0]["message"][
"content"
]
### POST-CALL RULES ###
rules_obj.post_call_rules(input=model_response, model=model)
is_coroutine = check_coroutine(original_function)
if is_coroutine == True:
pass
else:
model_response = original_response["choices"][0]["message"][
"content"
]
### POST-CALL RULES ###
rules_obj.post_call_rules(input=model_response, model=model)
except Exception as e:
raise e