(feat) langfuse - log embeddings

This commit is contained in:
ishaan-jaff 2024-01-31 10:22:29 -08:00
parent a0022410fa
commit f4def61a90

View file

@ -9,6 +9,7 @@ dotenv.load_dotenv() # Loading env variables using dotenv
import traceback import traceback
from packaging.version import Version from packaging.version import Version
from litellm._logging import verbose_logger from litellm._logging import verbose_logger
import litellm
class LangFuseLogger: class LangFuseLogger:
@ -63,11 +64,15 @@ class LangFuseLogger:
pass pass
# end of processing langfuse ######################## # end of processing langfuse ########################
if kwargs.get("call_type") == "embedding" or isinstance(
response_obj, litellm.EmbeddingResponse
):
input = prompt
output = response_obj["data"]
else:
input = prompt input = prompt
output = response_obj["choices"][0]["message"].json() output = response_obj["choices"][0]["message"].json()
print_verbose( print_verbose(f"OUTPUT IN LANGFUSE: {output}; original: {response_obj}")
f"OUTPUT IN LANGFUSE: {output}; original: {response_obj['choices'][0]['message']}"
)
self._log_langfuse_v2( self._log_langfuse_v2(
user_id, user_id,
metadata, metadata,
@ -173,8 +178,12 @@ class LangFuseLogger:
tags = [] tags = []
supports_tags = Version(langfuse.version.__version__) >= Version("2.6.3") supports_tags = Version(langfuse.version.__version__) >= Version("2.6.3")
generation_name = metadata.get("generation_name", None)
if generation_name is None:
# just log `litellm-{call_type}` as the generation name
generation_name = f"litellm-{kwargs.get('call_type', 'completion')}"
trace_params = { trace_params = {
"name": metadata.get("generation_name", "litellm-completion"), "name": generation_name,
"input": input, "input": input,
"output": output, "output": output,
"user_id": metadata.get("trace_user_id", user_id), "user_id": metadata.get("trace_user_id", user_id),
@ -190,7 +199,7 @@ class LangFuseLogger:
trace = self.Langfuse.trace(**trace_params) trace = self.Langfuse.trace(**trace_params)
trace.generation( trace.generation(
name=metadata.get("generation_name", "litellm-completion"), name=generation_name,
id=metadata.get("generation_id", None), id=metadata.get("generation_id", None),
startTime=start_time, startTime=start_time,
endTime=end_time, endTime=end_time,