fix(langfuse.py): support time to first token logging on langfuse

This commit is contained in:
Krrish Dholakia 2024-02-23 10:04:12 -08:00
parent 957614f0df
commit 70a523eee7
2 changed files with 17 additions and 1 deletions

View file

@ -820,6 +820,8 @@ class Logging:
## DYNAMIC LANGFUSE KEYS ##
self.langfuse_public_key = langfuse_public_key
self.langfuse_secret = langfuse_secret
## TIME TO FIRST TOKEN LOGGING ##
self.completion_start_time: Optional[datetime.datetime] = None
def update_environment_variables(
self, model, user, optional_params, litellm_params, **additional_params
@ -840,6 +842,7 @@ class Logging:
"user": user,
"call_type": str(self.call_type),
"litellm_call_id": self.litellm_call_id,
"completion_start_time": self.completion_start_time,
**self.optional_params,
**additional_params,
}
@ -1069,6 +1072,11 @@ class Logging:
start_time = self.start_time
if end_time is None:
end_time = datetime.datetime.now()
if self.completion_start_time is None:
self.completion_start_time = end_time
self.model_call_details["completion_start_time"] = (
self.completion_start_time
)
self.model_call_details["log_event_type"] = "successful_api_call"
self.model_call_details["end_time"] = end_time
self.model_call_details["cache_hit"] = cache_hit
@ -1358,7 +1366,7 @@ class Logging:
f"is complete_streaming_response in kwargs: {kwargs.get('complete_streaming_response', None)}"
)
if complete_streaming_response is None:
break
continue
else:
print_verbose("reaches langfuse for streaming logging!")
result = kwargs["complete_streaming_response"]