Merge pull request #1704 from maxdeichmann/add-langfuse-cost-tracking

feat: add langfuse cost tracking
This commit is contained in:
Krish Dholakia 2024-01-31 16:09:50 -08:00 committed by GitHub
commit e6231245c1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -103,6 +103,7 @@ class LangFuseLogger:
optional_params, optional_params,
input, input,
response_obj, response_obj,
print_verbose,
) if self._is_langfuse_v2() else self._log_langfuse_v1( ) if self._is_langfuse_v2() else self._log_langfuse_v1(
user_id, user_id,
metadata, metadata,
@ -192,11 +193,15 @@ class LangFuseLogger:
optional_params, optional_params,
input, input,
response_obj, response_obj,
print_verbose,
): ):
import langfuse import langfuse
tags = [] tags = []
supports_tags = Version(langfuse.version.__version__) >= Version("2.6.3") supports_tags = Version(langfuse.version.__version__) >= Version("2.6.3")
supports_costs = Version(langfuse.version.__version__) >= Version("2.7.3")
print_verbose(f"Langfuse Layer Logging - logging to langfuse v2 ")
generation_name = metadata.get("generation_name", None) generation_name = metadata.get("generation_name", None)
if generation_name is None: if generation_name is None:
@ -209,6 +214,8 @@ class LangFuseLogger:
"user_id": metadata.get("trace_user_id", user_id), "user_id": metadata.get("trace_user_id", user_id),
"id": metadata.get("trace_id", None), "id": metadata.get("trace_id", None),
} }
cost = kwargs["response_cost"]
print_verbose(f"trace: {cost}")
if supports_tags: if supports_tags:
for key, value in metadata.items(): for key, value in metadata.items():
tags.append(f"{key}:{value}") tags.append(f"{key}:{value}")
@ -217,7 +224,6 @@ class LangFuseLogger:
trace_params.update({"tags": tags}) trace_params.update({"tags": tags})
trace = self.Langfuse.trace(**trace_params) trace = self.Langfuse.trace(**trace_params)
trace.generation( trace.generation(
name=generation_name, name=generation_name,
id=metadata.get("generation_id", None), id=metadata.get("generation_id", None),
@ -230,6 +236,7 @@ class LangFuseLogger:
usage={ usage={
"prompt_tokens": response_obj["usage"]["prompt_tokens"], "prompt_tokens": response_obj["usage"]["prompt_tokens"],
"completion_tokens": response_obj["usage"]["completion_tokens"], "completion_tokens": response_obj["usage"]["completion_tokens"],
"total_cost": cost if supports_costs else None,
}, },
metadata=metadata, metadata=metadata,
) )