Feat/fix ops trace (#5672)

Co-authored-by: takatost <takatost@gmail.com>
This commit is contained in:
Joe
2024-06-28 00:24:37 +08:00
committed by GitHub
parent f0ea540b34
commit e8b8f6c6dd
17 changed files with 372 additions and 64 deletions

View File

@@ -147,6 +147,7 @@ class LangFuseDataTrace(BaseTraceInstance):
# add span
if trace_info.message_id:
span_data = LangfuseSpan(
id=node_execution_id,
name=f"{node_name}_{node_execution_id}",
input=inputs,
output=outputs,
@@ -160,6 +161,7 @@ class LangFuseDataTrace(BaseTraceInstance):
)
else:
span_data = LangfuseSpan(
id=node_execution_id,
name=f"{node_name}_{node_execution_id}",
input=inputs,
output=outputs,
@@ -173,6 +175,30 @@ class LangFuseDataTrace(BaseTraceInstance):
self.add_span(langfuse_span_data=span_data)
process_data = json.loads(node_execution.process_data) if node_execution.process_data else {}
if process_data and process_data.get("model_mode") == "chat":
total_token = metadata.get("total_tokens", 0)
# add generation
generation_usage = GenerationUsage(
totalTokens=total_token,
)
node_generation_data = LangfuseGeneration(
name=f"generation_{node_execution_id}",
trace_id=trace_id,
parent_observation_id=node_execution_id,
start_time=created_at,
end_time=finished_at,
input=inputs,
output=outputs,
metadata=metadata,
level=LevelEnum.DEFAULT if status == 'succeeded' else LevelEnum.ERROR,
status_message=trace_info.error if trace_info.error else "",
usage=generation_usage,
)
self.add_generation(langfuse_generation_data=node_generation_data)
def message_trace(
self, trace_info: MessageTraceInfo, **kwargs
):
@@ -186,7 +212,7 @@ class LangFuseDataTrace(BaseTraceInstance):
if message_data.from_end_user_id:
end_user_data: EndUser = db.session.query(EndUser).filter(
EndUser.id == message_data.from_end_user_id
).first().session_id
).first()
user_id = end_user_data.session_id
trace_data = LangfuseTrace(
@@ -220,6 +246,7 @@ class LangFuseDataTrace(BaseTraceInstance):
output=trace_info.answer_tokens,
total=trace_info.total_tokens,
unit=UnitEnum.TOKENS,
totalCost=message_data.total_price,
)
langfuse_generation_data = LangfuseGeneration(
@@ -303,7 +330,7 @@ class LangFuseDataTrace(BaseTraceInstance):
start_time=trace_info.start_time,
end_time=trace_info.end_time,
metadata=trace_info.metadata,
level=LevelEnum.DEFAULT if trace_info.error == "" else LevelEnum.ERROR,
level=LevelEnum.DEFAULT if trace_info.error == "" or trace_info.error is None else LevelEnum.ERROR,
status_message=trace_info.error,
)