dify/api/core/ops/entities/trace_entity.py

123 lines
3.4 KiB
Python

from datetime import datetime
from enum import Enum
from typing import Any, Optional, Union
from pydantic import BaseModel, ConfigDict, field_validator
class BaseTraceInfo(BaseModel):
message_id: Optional[str] = None
message_data: Optional[Any] = None
inputs: Optional[Union[str, dict[str, Any], list]] = None
outputs: Optional[Union[str, dict[str, Any], list]] = None
start_time: Optional[datetime] = None
end_time: Optional[datetime] = None
metadata: dict[str, Any]
@field_validator("inputs", "outputs")
@classmethod
def ensure_type(cls, v):
if v is None:
return None
if isinstance(v, str | dict | list):
return v
return ""
class WorkflowTraceInfo(BaseTraceInfo):
workflow_data: Any
conversation_id: Optional[str] = None
workflow_app_log_id: Optional[str] = None
workflow_id: str
tenant_id: str
workflow_run_id: str
workflow_run_elapsed_time: Union[int, float]
workflow_run_status: str
workflow_run_inputs: dict[str, Any]
workflow_run_outputs: dict[str, Any]
workflow_run_version: str
error: Optional[str] = None
total_tokens: int
file_list: list[str]
query: str
metadata: dict[str, Any]
class MessageTraceInfo(BaseTraceInfo):
conversation_model: str
message_tokens: int
answer_tokens: int
total_tokens: int
error: Optional[str] = None
file_list: Optional[Union[str, dict[str, Any], list]] = None
message_file_data: Optional[Any] = None
conversation_mode: str
class ModerationTraceInfo(BaseTraceInfo):
flagged: bool
action: str
preset_response: str
query: str
class SuggestedQuestionTraceInfo(BaseTraceInfo):
total_tokens: int
status: Optional[str] = None
error: Optional[str] = None
from_account_id: Optional[str] = None
agent_based: Optional[bool] = None
from_source: Optional[str] = None
model_provider: Optional[str] = None
model_id: Optional[str] = None
suggested_question: list[str]
level: str
status_message: Optional[str] = None
workflow_run_id: Optional[str] = None
model_config = ConfigDict(protected_namespaces=())
class DatasetRetrievalTraceInfo(BaseTraceInfo):
documents: Any
class ToolTraceInfo(BaseTraceInfo):
tool_name: str
tool_inputs: dict[str, Any]
tool_outputs: str
metadata: dict[str, Any]
message_file_data: Any
error: Optional[str] = None
tool_config: dict[str, Any]
time_cost: Union[int, float]
tool_parameters: dict[str, Any]
file_url: Union[str, None, list]
class GenerateNameTraceInfo(BaseTraceInfo):
conversation_id: Optional[str] = None
tenant_id: str
trace_info_info_map = {
"WorkflowTraceInfo": WorkflowTraceInfo,
"MessageTraceInfo": MessageTraceInfo,
"ModerationTraceInfo": ModerationTraceInfo,
"SuggestedQuestionTraceInfo": SuggestedQuestionTraceInfo,
"DatasetRetrievalTraceInfo": DatasetRetrievalTraceInfo,
"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"