AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-12 10:48:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: update ops trace (#7102)

Browse Source
This commit is contained in:
Joe 2024-08-09 15:22:16 +08:00 committed by GitHub
parent 7201b56a6d
commit 425174e82f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 118 additions and 129 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
api/core/app/apps/advanced_chat/app_generator.py
View File

@ -89,7 +89,8 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
)
# get tracing instance
trace_manager = TraceQueueManager(app_id=app_model.id)
user_id = user.id if isinstance(user, Account) else user.session_id
trace_manager = TraceQueueManager(app_model.id, user_id)
if invoke_from == InvokeFrom.DEBUGGER:
# always enable retriever resource in debugger mode

3
api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py
View File

@ -48,7 +48,8 @@ from core.model_runtime.entities.message_entities import (
)
from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
from core.model_runtime.utils.encoders import jsonable_encoder
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask, TraceTaskName
from core.ops.entities.trace_entity import TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
from core.prompt.utils.prompt_message_util import PromptMessageUtil
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from events.message_event import message_was_created

3
api/core/app/task_pipeline/workflow_cycle_manage.py
View File

@ -22,7 +22,8 @@ from core.app.entities.task_entities import (
from core.app.task_pipeline.workflow_iteration_cycle_manage import WorkflowIterationCycleManage
from core.file.file_obj import FileVar
from core.model_runtime.utils.encoders import jsonable_encoder
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask, TraceTaskName
from core.ops.entities.trace_entity import TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
from core.tools.tool_manager import ToolManager
from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeType
from core.workflow.nodes.tool.entities import ToolNodeData

3
api/core/callback_handler/agent_tool_callback_handler.py
View File

@ -4,7 +4,8 @@ from typing import Any, Optional, TextIO, Union
from pydantic import BaseModel
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask, TraceTaskName
from core.ops.entities.trace_entity import TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
from core.tools.entities.tool_entities import ToolInvokeMessage
_TEXT_COLOR_MAPPING = {

3
api/core/llm_generator/llm_generator.py
View File

@ -14,7 +14,8 @@ from core.model_manager import ModelManager
from core.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage
from core.model_runtime.entities.model_entities import ModelType
from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask, TraceTaskName
from core.ops.entities.trace_entity import TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
from core.ops.utils import measure_time
from core.prompt.utils.prompt_template_parser import PromptTemplateParser

3
api/core/moderation/input_moderation.py
View File

@ -4,7 +4,8 @@ from typing import Optional
from core.app.app_config.entities import AppConfig
from core.moderation.base import ModerationAction, ModerationException
from core.moderation.factory import ModerationFactory
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask, TraceTaskName
from core.ops.entities.trace_entity import TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
from core.ops.utils import measure_time
logger = logging.getLogger(__name__)

12
api/core/ops/entities/trace_entity.py
View File

@ -1,4 +1,5 @@
from datetime import datetime
from enum import Enum
from typing import Any, Optional, Union
from pydantic import BaseModel, ConfigDict, field_validator
@ -106,3 +107,14 @@ trace_info_info_map = {
'ToolTraceInfo': ToolTraceInfo,
'GenerateNameTraceInfo': GenerateNameTraceInfo,
}
class TraceTaskName(str, Enum):
CONVERSATION_TRACE = 'conversation'
WORKFLOW_TRACE = 'workflow'
MESSAGE_TRACE = 'message'
MODERATION_TRACE = 'moderation'
SUGGESTED_QUESTION_TRACE = 'suggested_question'
DATASET_RETRIEVAL_TRACE = 'dataset_retrieval'
TOOL_TRACE = 'tool'
GENERATE_NAME_TRACE = 'generate_conversation_name'

10
api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py
View File

@ -50,6 +50,7 @@ class LangfuseTrace(BaseModel):
"""
Langfuse trace model
"""
id: Optional[str] = Field(
default=None,
description="The id of the trace can be set, defaults to a random id. Used to link traces to external systems "
@ -109,6 +110,7 @@ class LangfuseSpan(BaseModel):
"""
Langfuse span model
"""
id: Optional[str] = Field(
default=None,
description="The id of the span can be set, otherwise a random id is generated. Spans are upserted on id.",
@ -185,10 +187,9 @@ class UnitEnum(str, Enum):
class GenerationUsage(BaseModel):
promptTokens: Optional[int] = None
completionTokens: Optional[int] = None
totalTokens: Optional[int] = None
total: Optional[int] = None
input: Optional[int] = None
output: Optional[int] = None
total: Optional[int] = None
unit: Optional[UnitEnum] = None
inputCost: Optional[float] = None
outputCost: Optional[float] = None
@ -230,9 +231,7 @@ class LangfuseGeneration(BaseModel):
default=None,
description="The time at which the generation ended. Automatically set by generation.end().",
)
model: Optional[str] = Field(
default=None, description="The name of the model used for the generation."
)
model: Optional[str] = Field(default=None, description="The name of the model used for the generation.")
model_parameters: Optional[dict[str, Any]] = Field(
default=None,
description="The parameters of the model used for the generation; can be any key-value pairs.",
@ -277,4 +276,3 @@ class LangfuseGeneration(BaseModel):
def ensure_dict(cls, v, info: ValidationInfo):
field_name = info.field_name
return validate_input_output(v, field_name)

103
api/core/ops/langfuse_trace/langfuse_trace.py
View File

@ -16,6 +16,7 @@ from core.ops.entities.trace_entity import (
ModerationTraceInfo,
SuggestedQuestionTraceInfo,
ToolTraceInfo,
TraceTaskName,
WorkflowTraceInfo,
)
from core.ops.langfuse_trace.entities.langfuse_trace_entity import (
@ -68,9 +69,9 @@ class LangFuseDataTrace(BaseTraceInstance):
user_id = trace_info.metadata.get("user_id")
if trace_info.message_id:
trace_id = trace_info.message_id
name = f"message_{trace_info.message_id}"
name = TraceTaskName.MESSAGE_TRACE.value
trace_data = LangfuseTrace(
id=trace_info.message_id,
id=trace_id,
user_id=user_id,
name=name,
input=trace_info.workflow_run_inputs,
@ -78,11 +79,13 @@ class LangFuseDataTrace(BaseTraceInstance):
metadata=trace_info.metadata,
session_id=trace_info.conversation_id,
tags=["message", "workflow"],
created_at=trace_info.start_time,
updated_at=trace_info.end_time,
)
self.add_trace(langfuse_trace_data=trace_data)
workflow_span_data = LangfuseSpan(
id=trace_info.workflow_app_log_id if trace_info.workflow_app_log_id else trace_info.workflow_run_id,
name=f"workflow_{trace_info.workflow_app_log_id}" if trace_info.workflow_app_log_id else f"workflow_{trace_info.workflow_run_id}",
id=(trace_info.workflow_app_log_id if trace_info.workflow_app_log_id else trace_info.workflow_run_id),
name=TraceTaskName.WORKFLOW_TRACE.value,
input=trace_info.workflow_run_inputs,
output=trace_info.workflow_run_outputs,
trace_id=trace_id,
@ -97,7 +100,7 @@ class LangFuseDataTrace(BaseTraceInstance):
trace_data = LangfuseTrace(
id=trace_id,
user_id=user_id,
name=f"workflow_{trace_info.workflow_app_log_id}" if trace_info.workflow_app_log_id else f"workflow_{trace_info.workflow_run_id}",
name=TraceTaskName.WORKFLOW_TRACE.value,
input=trace_info.workflow_run_inputs,
output=trace_info.workflow_run_outputs,
metadata=trace_info.metadata,
@ -134,14 +137,12 @@ class LangFuseDataTrace(BaseTraceInstance):
node_type = node_execution.node_type
status = node_execution.status
if node_type == "llm":
inputs = json.loads(node_execution.process_data).get(
"prompts", {}
) if node_execution.process_data else {}
inputs = (
json.loads(node_execution.process_data).get("prompts", {}) if node_execution.process_data else {}
)
else:
inputs = json.loads(node_execution.inputs) if node_execution.inputs else {}
outputs = (
json.loads(node_execution.outputs) if node_execution.outputs else {}
)
outputs = json.loads(node_execution.outputs) if node_execution.outputs else {}
created_at = node_execution.created_at if node_execution.created_at else datetime.now()
elapsed_time = node_execution.elapsed_time
finished_at = created_at + timedelta(seconds=elapsed_time)
@ -163,28 +164,30 @@ class LangFuseDataTrace(BaseTraceInstance):
if trace_info.message_id:
span_data = LangfuseSpan(
id=node_execution_id,
name=f"{node_name}_{node_execution_id}",
name=node_type,
input=inputs,
output=outputs,
trace_id=trace_id,
start_time=created_at,
end_time=finished_at,
metadata=metadata,
level=LevelEnum.DEFAULT if status == 'succeeded' else LevelEnum.ERROR,
level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR),
status_message=trace_info.error if trace_info.error else "",
parent_observation_id=trace_info.workflow_app_log_id if trace_info.workflow_app_log_id else trace_info.workflow_run_id,
parent_observation_id=(
trace_info.workflow_app_log_id if trace_info.workflow_app_log_id else trace_info.workflow_run_id
),
)
else:
span_data = LangfuseSpan(
id=node_execution_id,
name=f"{node_name}_{node_execution_id}",
name=node_type,
input=inputs,
output=outputs,
trace_id=trace_id,
start_time=created_at,
end_time=finished_at,
metadata=metadata,
level=LevelEnum.DEFAULT if status == 'succeeded' else LevelEnum.ERROR,
level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR),
status_message=trace_info.error if trace_info.error else "",
)
@ -195,11 +198,11 @@ class LangFuseDataTrace(BaseTraceInstance):
total_token = metadata.get("total_tokens", 0)
# add generation
generation_usage = GenerationUsage(
totalTokens=total_token,
total=total_token,
)
node_generation_data = LangfuseGeneration(
name=f"generation_{node_execution_id}",
name="llm",
trace_id=trace_id,
parent_observation_id=node_execution_id,
start_time=created_at,
@ -207,16 +210,14 @@ class LangFuseDataTrace(BaseTraceInstance):
input=inputs,
output=outputs,
metadata=metadata,
level=LevelEnum.DEFAULT if status == 'succeeded' else LevelEnum.ERROR,
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
):
def message_trace(self, trace_info: MessageTraceInfo, **kwargs):
# get message file data
file_list = trace_info.file_list
metadata = trace_info.metadata
@ -225,9 +226,9 @@ class LangFuseDataTrace(BaseTraceInstance):
user_id = message_data.from_account_id
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()
end_user_data: EndUser = (
db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first()
)
if end_user_data is not None:
user_id = end_user_data.session_id
metadata["user_id"] = user_id
@ -235,7 +236,7 @@ class LangFuseDataTrace(BaseTraceInstance):
trace_data = LangfuseTrace(
id=message_id,
user_id=user_id,
name=f"message_{message_id}",
name=TraceTaskName.MESSAGE_TRACE.value,
input={
"message": trace_info.inputs,
"files": file_list,
@ -258,7 +259,6 @@ class LangFuseDataTrace(BaseTraceInstance):
# start add span
generation_usage = GenerationUsage(
totalTokens=trace_info.total_tokens,
input=trace_info.message_tokens,
output=trace_info.answer_tokens,
total=trace_info.total_tokens,
@ -267,7 +267,7 @@ class LangFuseDataTrace(BaseTraceInstance):
)
langfuse_generation_data = LangfuseGeneration(
name=f"generation_{message_id}",
name="llm",
trace_id=message_id,
start_time=trace_info.start_time,
end_time=trace_info.end_time,
@ -275,7 +275,7 @@ class LangFuseDataTrace(BaseTraceInstance):
input=trace_info.inputs,
output=message_data.answer,
metadata=metadata,
level=LevelEnum.DEFAULT if message_data.status != 'error' else LevelEnum.ERROR,
level=(LevelEnum.DEFAULT if message_data.status != "error" else LevelEnum.ERROR),
status_message=message_data.error if message_data.error else "",
usage=generation_usage,
)
@ -284,7 +284,7 @@ class LangFuseDataTrace(BaseTraceInstance):
def moderation_trace(self, trace_info: ModerationTraceInfo):
span_data = LangfuseSpan(
name="moderation",
name=TraceTaskName.MODERATION_TRACE.value,
input=trace_info.inputs,
output={
"action": trace_info.action,
@ -303,22 +303,21 @@ class LangFuseDataTrace(BaseTraceInstance):
def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo):
message_data = trace_info.message_data
generation_usage = GenerationUsage(
totalTokens=len(str(trace_info.suggested_question)),
total=len(str(trace_info.suggested_question)),
input=len(trace_info.inputs),
output=len(trace_info.suggested_question),
total=len(trace_info.suggested_question),
unit=UnitEnum.CHARACTERS,
)
generation_data = LangfuseGeneration(
name="suggested_question",
name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value,
input=trace_info.inputs,
output=str(trace_info.suggested_question),
trace_id=trace_info.message_id,
start_time=trace_info.start_time,
end_time=trace_info.end_time,
metadata=trace_info.metadata,
level=LevelEnum.DEFAULT if message_data.status != 'error' else LevelEnum.ERROR,
level=(LevelEnum.DEFAULT if message_data.status != "error" else LevelEnum.ERROR),
status_message=message_data.error if message_data.error else "",
usage=generation_usage,
)
@ -327,7 +326,7 @@ class LangFuseDataTrace(BaseTraceInstance):
def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo):
dataset_retrieval_span_data = LangfuseSpan(
name="dataset_retrieval",
name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value,
input=trace_info.inputs,
output={"documents": trace_info.documents},
trace_id=trace_info.message_id,
@ -347,7 +346,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 == "" or trace_info.error is None else LevelEnum.ERROR,
level=(LevelEnum.DEFAULT if trace_info.error == "" or trace_info.error is None else LevelEnum.ERROR),
status_message=trace_info.error,
)
@ -355,7 +354,7 @@ class LangFuseDataTrace(BaseTraceInstance):
def generate_name_trace(self, trace_info: GenerateNameTraceInfo):
name_generation_trace_data = LangfuseTrace(
name="generate_name",
name=TraceTaskName.GENERATE_NAME_TRACE.value,
input=trace_info.inputs,
output=trace_info.outputs,
user_id=trace_info.tenant_id,
@ -366,7 +365,7 @@ class LangFuseDataTrace(BaseTraceInstance):
self.add_trace(langfuse_trace_data=name_generation_trace_data)
name_generation_span_data = LangfuseSpan(
name="generate_name",
name=TraceTaskName.GENERATE_NAME_TRACE.value,
input=trace_info.inputs,
output=trace_info.outputs,
trace_id=trace_info.conversation_id,
@ -377,9 +376,7 @@ class LangFuseDataTrace(BaseTraceInstance):
self.add_span(langfuse_span_data=name_generation_span_data)
def add_trace(self, langfuse_trace_data: Optional[LangfuseTrace] = None):
format_trace_data = (
filter_none_values(langfuse_trace_data.model_dump()) if langfuse_trace_data else {}
)
format_trace_data = filter_none_values(langfuse_trace_data.model_dump()) if langfuse_trace_data else {}
try:
self.langfuse_client.trace(**format_trace_data)
logger.debug("LangFuse Trace created successfully")
@ -387,9 +384,7 @@ class LangFuseDataTrace(BaseTraceInstance):
raise ValueError(f"LangFuse Failed to create trace: {str(e)}")
def add_span(self, langfuse_span_data: Optional[LangfuseSpan] = None):
format_span_data = (
filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {}
)
format_span_data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {}
try:
self.langfuse_client.span(**format_span_data)
logger.debug("LangFuse Span created successfully")
@ -397,19 +392,13 @@ class LangFuseDataTrace(BaseTraceInstance):
raise ValueError(f"LangFuse Failed to create span: {str(e)}")
def update_span(self, span, langfuse_span_data: Optional[LangfuseSpan] = None):
format_span_data = (
filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {}
)
format_span_data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {}
span.end(**format_span_data)
def add_generation(
self, langfuse_generation_data: Optional[LangfuseGeneration] = None
):
def add_generation(self, langfuse_generation_data: Optional[LangfuseGeneration] = None):
format_generation_data = (
filter_none_values(langfuse_generation_data.model_dump())
if langfuse_generation_data
else {}
filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {}
)
try:
self.langfuse_client.generation(**format_generation_data)
@ -417,13 +406,9 @@ class LangFuseDataTrace(BaseTraceInstance):
except Exception as e:
raise ValueError(f"LangFuse Failed to create generation: {str(e)}")
def update_generation(
self, generation, langfuse_generation_data: Optional[LangfuseGeneration] = None
):
def update_generation(self, generation, langfuse_generation_data: Optional[LangfuseGeneration] = None):
format_generation_data = (
filter_none_values(langfuse_generation_data.model_dump())
if langfuse_generation_data
else {}
filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {}
)
generation.end(**format_generation_data)

51
api/core/ops/langsmith_trace/langsmith_trace.py
View File

@ -15,6 +15,7 @@ from core.ops.entities.trace_entity import (
ModerationTraceInfo,
SuggestedQuestionTraceInfo,
ToolTraceInfo,
TraceTaskName,
WorkflowTraceInfo,
)
from core.ops.langsmith_trace.entities.langsmith_trace_entity import (
@ -39,9 +40,7 @@ class LangSmithDataTrace(BaseTraceInstance):
self.langsmith_key = langsmith_config.api_key
self.project_name = langsmith_config.project
self.project_id = None
self.langsmith_client = Client(
api_key=langsmith_config.api_key, api_url=langsmith_config.endpoint
)
self.langsmith_client = Client(api_key=langsmith_config.api_key, api_url=langsmith_config.endpoint)
self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001")
def trace(self, trace_info: BaseTraceInfo):
@ -64,7 +63,7 @@ class LangSmithDataTrace(BaseTraceInstance):
if trace_info.message_id:
message_run = LangSmithRunModel(
id=trace_info.message_id,
name=f"message_{trace_info.message_id}",
name=TraceTaskName.MESSAGE_TRACE.value,
inputs=trace_info.workflow_run_inputs,
outputs=trace_info.workflow_run_outputs,
run_type=LangSmithRunType.chain,
@ -73,8 +72,8 @@ class LangSmithDataTrace(BaseTraceInstance):
extra={
"metadata": trace_info.metadata,
},
tags=["message"],
error=trace_info.error
tags=["message", "workflow"],
error=trace_info.error,
)
self.add_run(message_run)
@ -82,7 +81,7 @@ class LangSmithDataTrace(BaseTraceInstance):
file_list=trace_info.file_list,
total_tokens=trace_info.total_tokens,
id=trace_info.workflow_app_log_id if trace_info.workflow_app_log_id else trace_info.workflow_run_id,
name=f"workflow_{trace_info.workflow_app_log_id}" if trace_info.workflow_app_log_id else f"workflow_{trace_info.workflow_run_id}",
name=TraceTaskName.WORKFLOW_TRACE.value,
inputs=trace_info.workflow_run_inputs,
run_type=LangSmithRunType.tool,
start_time=trace_info.workflow_data.created_at,
@ -126,22 +125,18 @@ class LangSmithDataTrace(BaseTraceInstance):
node_type = node_execution.node_type
status = node_execution.status
if node_type == "llm":
inputs = json.loads(node_execution.process_data).get(
"prompts", {}
) if node_execution.process_data else {}
inputs = (
json.loads(node_execution.process_data).get("prompts", {}) if node_execution.process_data else {}
)
else:
inputs = json.loads(node_execution.inputs) if node_execution.inputs else {}
outputs = (
json.loads(node_execution.outputs) if node_execution.outputs else {}
)
outputs = json.loads(node_execution.outputs) if node_execution.outputs else {}
created_at = node_execution.created_at if node_execution.created_at else datetime.now()
elapsed_time = node_execution.elapsed_time
finished_at = created_at + timedelta(seconds=elapsed_time)
execution_metadata = (
json.loads(node_execution.execution_metadata)
if node_execution.execution_metadata
else {}
json.loads(node_execution.execution_metadata) if node_execution.execution_metadata else {}
)
node_total_tokens = execution_metadata.get("total_tokens", 0)
@ -168,7 +163,7 @@ class LangSmithDataTrace(BaseTraceInstance):
langsmith_run = LangSmithRunModel(
total_tokens=node_total_tokens,
name=f"{node_name}_{node_execution_id}",
name=node_type,
inputs=inputs,
run_type=run_type,
start_time=created_at,
@ -178,7 +173,9 @@ class LangSmithDataTrace(BaseTraceInstance):
extra={
"metadata": metadata,
},
parent_run_id=trace_info.workflow_app_log_id if trace_info.workflow_app_log_id else trace_info.workflow_run_id,
parent_run_id=trace_info.workflow_app_log_id
if trace_info.workflow_app_log_id
else trace_info.workflow_run_id,
tags=["node_execution"],
)
@ -198,9 +195,9 @@ class LangSmithDataTrace(BaseTraceInstance):
metadata["user_id"] = user_id
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()
end_user_data: EndUser = (
db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first()
)
if end_user_data is not None:
end_user_id = end_user_data.session_id
metadata["end_user_id"] = end_user_id
@ -210,7 +207,7 @@ class LangSmithDataTrace(BaseTraceInstance):
output_tokens=trace_info.answer_tokens,
total_tokens=trace_info.total_tokens,
id=message_id,
name=f"message_{message_id}",
name=TraceTaskName.MESSAGE_TRACE.value,
inputs=trace_info.inputs,
run_type=LangSmithRunType.chain,
start_time=trace_info.start_time,
@ -230,7 +227,7 @@ class LangSmithDataTrace(BaseTraceInstance):
input_tokens=trace_info.message_tokens,
output_tokens=trace_info.answer_tokens,
total_tokens=trace_info.total_tokens,
name=f"llm_{message_id}",
name="llm",
inputs=trace_info.inputs,
run_type=LangSmithRunType.llm,
start_time=trace_info.start_time,
@ -248,7 +245,7 @@ class LangSmithDataTrace(BaseTraceInstance):
def moderation_trace(self, trace_info: ModerationTraceInfo):
langsmith_run = LangSmithRunModel(
name="moderation",
name=TraceTaskName.MODERATION_TRACE.value,
inputs=trace_info.inputs,
outputs={
"action": trace_info.action,
@ -271,7 +268,7 @@ class LangSmithDataTrace(BaseTraceInstance):
def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo):
message_data = trace_info.message_data
suggested_question_run = LangSmithRunModel(
name="suggested_question",
name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value,
inputs=trace_info.inputs,
outputs=trace_info.suggested_question,
run_type=LangSmithRunType.tool,
@ -288,7 +285,7 @@ class LangSmithDataTrace(BaseTraceInstance):
def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo):
dataset_retrieval_run = LangSmithRunModel(
name="dataset_retrieval",
name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value,
inputs=trace_info.inputs,
outputs={"documents": trace_info.documents},
run_type=LangSmithRunType.retriever,
@ -323,7 +320,7 @@ class LangSmithDataTrace(BaseTraceInstance):
def generate_name_trace(self, trace_info: GenerateNameTraceInfo):
name_run = LangSmithRunModel(
name="generate_name",
name=TraceTaskName.GENERATE_NAME_TRACE.value,
inputs=trace_info.inputs,
outputs=trace_info.outputs,
run_type=LangSmithRunType.tool,

13
api/core/ops/ops_trace_manager.py
View File

@ -5,7 +5,6 @@ import queue
import threading
import time
from datetime import timedelta
from enum import Enum
from typing import Any, Optional, Union
from uuid import UUID
@ -24,6 +23,7 @@ from core.ops.entities.trace_entity import (
ModerationTraceInfo,
SuggestedQuestionTraceInfo,
ToolTraceInfo,
TraceTaskName,
WorkflowTraceInfo,
)
from core.ops.langfuse_trace.langfuse_trace import LangFuseDataTrace
@ -253,17 +253,6 @@ class OpsTraceManager:
return trace_instance(tracing_config).api_check()
class TraceTaskName(str, Enum):
CONVERSATION_TRACE = 'conversation_trace'
WORKFLOW_TRACE = 'workflow_trace'
MESSAGE_TRACE = 'message_trace'
MODERATION_TRACE = 'moderation_trace'
SUGGESTED_QUESTION_TRACE = 'suggested_question_trace'
DATASET_RETRIEVAL_TRACE = 'dataset_retrieval_trace'
TOOL_TRACE = 'tool_trace'
GENERATE_NAME_TRACE = 'generate_name_trace'
class TraceTask:
def __init__(
self,

3
api/core/rag/retrieval/dataset_retrieval.py
View File

@ -14,7 +14,8 @@ from core.model_manager import ModelInstance, ModelManager
from core.model_runtime.entities.message_entities import PromptMessageTool
from core.model_runtime.entities.model_entities import ModelFeature, ModelType
from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask, TraceTaskName
from core.ops.entities.trace_entity import TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
from core.ops.utils import measure_time
from core.rag.data_post_processor.data_post_processor import DataPostProcessor
from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler

3
api/services/message_service.py
View File

@ -7,7 +7,8 @@ from core.llm_generator.llm_generator import LLMGenerator
from core.memory.token_buffer_memory import TokenBufferMemory
from core.model_manager import ModelManager
from core.model_runtime.entities.model_entities import ModelType
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask, TraceTaskName
from core.ops.entities.trace_entity import TraceTaskName
from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
from core.ops.utils import measure_time
from extensions.ext_database import db
from libs.infinite_scroll_pagination import InfiniteScrollPagination