fix streaming

This commit is contained in:
Vince Loewe 2024-02-29 13:23:51 -08:00
parent d4d160353e
commit fa6211616e
2 changed files with 21 additions and 5 deletions

View file

@ -7,6 +7,7 @@ sys.path.insert(0, os.path.abspath("../.."))
from litellm import completion from litellm import completion
import litellm import litellm
litellm.failure_callback = ["lunary"]
litellm.success_callback = ["lunary"] litellm.success_callback = ["lunary"]
litellm.set_verbose = True litellm.set_verbose = True

View file

@ -1336,9 +1336,10 @@ class Logging:
if callback == "lunary": if callback == "lunary":
print_verbose("reaches lunary for logging!") print_verbose("reaches lunary for logging!")
model = self.model model = self.model
kwargs = self.model_call_details
input = self.model_call_details.get( input = kwargs.get(
"messages", self.model_call_details.get("input", None) "messages", kwargs.get("input", None)
) )
type = ( type = (
@ -1347,15 +1348,27 @@ class Logging:
else "llm" else "llm"
) )
# this only logs streaming once, complete_streaming_response exists i.e when stream ends
if self.stream:
print_verbose("reaches lunary for streaming logging!")
print(kwargs)
if "complete_streaming_response" not in kwargs:
break
else:
print_verbose(
"reaches lunary for streaming logging!"
)
result = kwargs["complete_streaming_response"]
lunaryLogger.log_event( lunaryLogger.log_event(
type=type, type=type,
kwargs=self.model_call_details, kwargs=kwargs,
event="end", event="end",
model=model, model=model,
input=input, input=input,
user_id=self.model_call_details.get("user_id", self.model_call_details.get("user", None)), user_id=kwargs.get("user", None),
#user_props=self.model_call_details.get("user_props", None), #user_props=self.model_call_details.get("user_props", None),
extra=self.model_call_details.get("optional_params", {}), extra=kwargs.get("optional_params", {}),
response_obj=result, response_obj=result,
start_time=start_time, start_time=start_time,
end_time=end_time, end_time=end_time,
@ -1873,6 +1886,8 @@ class Logging:
else "llm" else "llm"
) )
lunaryLogger.log_event( lunaryLogger.log_event(
type=_type, type=_type,
event="error", event="error",