mirror of
https://github.com/langgenius/dify.git
synced 2026-10-04 13:47:55 +08:00
feat(trace): add provider-neutral unified tracing (#39451)
Co-authored-by: EvoX <evox@evomap.ai>
This commit is contained in:
1 parent
23150c2319
commit
3a8fbcbd02
37 files changed
+3496
-58
No files matched your search
@@ -3,14 +3,43 @@ from typing import override
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from packaging.version import Version
|
||||
from pydantic import SecretStr
|
||||
from pydantic import SecretStr, ValidationError
|
||||
from pydantic_settings import BaseSettings, PydanticBaseSettingsSource
|
||||
from yarl import URL
|
||||
|
||||
from configs.app_config import DifyConfig
|
||||
from configs.feature import OpsTraceConfig
|
||||
from enums import DeploymentEdition
|
||||
|
||||
|
||||
def test_ops_trace_config_rejects_parent_context_ttl_shorter_than_retry_window() -> None:
|
||||
with pytest.raises(ValidationError, match="must cover the retry window"):
|
||||
OpsTraceConfig(
|
||||
OPS_TRACE_UNIFIED_ENABLED=True,
|
||||
OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=4,
|
||||
OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5,
|
||||
OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=19,
|
||||
)
|
||||
|
||||
|
||||
def test_ops_trace_config_skips_parent_context_validation_when_unified_tracing_is_disabled() -> None:
|
||||
OpsTraceConfig(
|
||||
OPS_TRACE_UNIFIED_ENABLED=False,
|
||||
OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=4,
|
||||
OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5,
|
||||
OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=19,
|
||||
)
|
||||
|
||||
|
||||
def test_ops_trace_config_accepts_parent_context_ttl_covering_retry_window() -> None:
|
||||
OpsTraceConfig(
|
||||
OPS_TRACE_UNIFIED_ENABLED=True,
|
||||
OPS_TRACE_RETRYABLE_DISPATCH_MAX_RETRIES=4,
|
||||
OPS_TRACE_RETRYABLE_DISPATCH_DELAY_SECONDS=5,
|
||||
OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=20,
|
||||
)
|
||||
|
||||
|
||||
class _IsolatedDifyConfig(DifyConfig):
|
||||
"""Load explicit test values and packaging metadata without consulting process state."""
|
||||
|
||||
|
||||
@@ -34,6 +34,23 @@ def _build_pipeline() -> pipeline_module.AdvancedChatAppGenerateTaskPipeline:
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_process_passes_message_id_to_conversation_name_generation() -> None:
|
||||
pipeline = _build_pipeline()
|
||||
pipeline._conversation_id = "conversation-1"
|
||||
pipeline._application_generate_entity = SimpleNamespace(query="hello", trace_manager=None)
|
||||
pipeline._message_cycle_manager = mock.Mock()
|
||||
pipeline._base_task_pipeline = SimpleNamespace(stream=True)
|
||||
pipeline._wrapper_process_stream_response = mock.Mock(return_value=iter(()))
|
||||
pipeline._to_stream_response = mock.Mock(return_value="stream")
|
||||
|
||||
result = pipeline.process()
|
||||
|
||||
assert result == "stream"
|
||||
pipeline._message_cycle_manager.generate_conversation_name.assert_called_once_with(
|
||||
conversation_id="conversation-1", query="hello", message_id="message-1"
|
||||
)
|
||||
|
||||
|
||||
def test_persist_human_input_extra_content_adds_record(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pipeline = _build_pipeline()
|
||||
monkeypatch.setattr(pipeline, "_load_human_input_form_id", lambda **kwargs: "form-1")
|
||||
|
||||
+1
-1
@@ -858,7 +858,7 @@ class TestEasyUiBasedGenerateTaskPipeline:
|
||||
|
||||
assert result == "streamed"
|
||||
pipeline._message_cycle_manager.generate_conversation_name.assert_called_once_with(
|
||||
conversation_id="conv", query="hello"
|
||||
conversation_id="conv", query="hello", message_id="msg"
|
||||
)
|
||||
|
||||
def test_process_routes_to_blocking_for_completion_mode(self):
|
||||
|
||||
+10
-2
@@ -291,7 +291,9 @@ class TestMessageCycleManagerOptimization:
|
||||
),
|
||||
patch("core.app.task_pipeline.message_cycle_manager.Timer", DummyTimer),
|
||||
):
|
||||
thread = message_cycle_manager.generate_conversation_name(conversation_id="conv-1", query="hello")
|
||||
thread = message_cycle_manager.generate_conversation_name(
|
||||
conversation_id="conv-1", query="hello", message_id="message-1"
|
||||
)
|
||||
|
||||
assert isinstance(thread, DummyTimer)
|
||||
assert thread.interval == 1
|
||||
@@ -301,6 +303,7 @@ class TestMessageCycleManagerOptimization:
|
||||
assert thread.kwargs["flask_app"] is flask_app
|
||||
assert thread.kwargs["conversation_id"] == "conv-1"
|
||||
assert thread.kwargs["query"] == "hello"
|
||||
assert thread.kwargs["message_id"] == "message-1"
|
||||
assert message_cycle_manager._application_generate_entity.is_new_conversation is False
|
||||
|
||||
def test_generate_conversation_name_skips_thread_when_auto_generate_disabled(self, message_cycle_manager):
|
||||
@@ -377,13 +380,18 @@ class TestMessageCycleManagerOptimization:
|
||||
mock_redis.get.return_value = None
|
||||
mock_llm_generator.generate_conversation_name.return_value = "generated-title"
|
||||
|
||||
message_cycle_manager._generate_conversation_name_worker(flask_app, "conv-1", "hello")
|
||||
message_cycle_manager._generate_conversation_name_worker(
|
||||
flask_app, "conv-1", "hello", message_id="message-1"
|
||||
)
|
||||
|
||||
assert cycle_db.in_transaction() is False
|
||||
with Session(sqlite_engine) as verification_session:
|
||||
conversation = verification_session.get(Conversation, "conv-1")
|
||||
assert conversation is not None
|
||||
assert conversation.name == "generated-title"
|
||||
mock_llm_generator.generate_conversation_name.assert_called_once_with(
|
||||
"tenant-1", "hello", "conv-1", "app-id", message_id="message-1"
|
||||
)
|
||||
mock_redis.setex.assert_called_once()
|
||||
|
||||
def test_generate_conversation_name_worker_falls_back_when_generation_fails(
|
||||
|
||||
@@ -186,9 +186,13 @@ class TestLLMGenerator:
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
|
||||
with patch("core.llm_generator.llm_generator.TraceQueueManager") as mock_trace:
|
||||
name = LLMGenerator.generate_conversation_name("tenant_id", "test query")
|
||||
name = LLMGenerator.generate_conversation_name(
|
||||
"tenant_id", "test query", "conversation-1", "app-1", message_id="message-1"
|
||||
)
|
||||
assert name == "Test Conversation Name"
|
||||
mock_trace.assert_called_once()
|
||||
mock_trace.assert_called_once_with(app_id="app-1")
|
||||
trace_task = mock_trace.return_value.add_trace_task.call_args.args[0]
|
||||
assert trace_task.message_id == "message-1"
|
||||
|
||||
def test_generate_conversation_name_truncated(self, mock_model_instance):
|
||||
long_query = "a" * 2100
|
||||
|
||||
@@ -17,6 +17,7 @@ from sqlalchemy import Engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
import core.ops.ops_trace_manager as module
|
||||
from configs import dify_config
|
||||
from core.ops.ops_trace_manager import OpsTraceManager, TraceQueueManager, TraceTask, TraceTaskName
|
||||
from core.rag.models.document import Document as RetrievalDocument
|
||||
from graphon.enums import WorkflowExecutionStatus
|
||||
@@ -49,6 +50,15 @@ class DummyTraceInstance:
|
||||
return "https://project.fake"
|
||||
|
||||
|
||||
class DummyUnifiedTraceInstance(DummyTraceInstance):
|
||||
pass
|
||||
|
||||
|
||||
class FailingUnifiedTraceInstance(DummyTraceInstance):
|
||||
def __init__(self, config):
|
||||
raise RuntimeError("unified constructor failed")
|
||||
|
||||
|
||||
class FakeProviderMap:
|
||||
def __init__(self, data):
|
||||
self._data = data
|
||||
@@ -66,6 +76,11 @@ PROVIDER_ENTRY = {
|
||||
"trace_instance": DummyTraceInstance,
|
||||
}
|
||||
|
||||
UNIFIED_PROVIDER_ENTRY = {
|
||||
"config_class": DummyConfig,
|
||||
"trace_instance": DummyUnifiedTraceInstance,
|
||||
}
|
||||
|
||||
|
||||
class DummyTimer:
|
||||
def __init__(self, interval, function):
|
||||
@@ -120,9 +135,11 @@ class RecordingStorage:
|
||||
class RecordingDispatcher:
|
||||
def __init__(self) -> None:
|
||||
self.payloads: list[dict[str, str]] = []
|
||||
self.options: list[dict[str, object]] = []
|
||||
|
||||
def delay(self, payload: dict[str, str]) -> None:
|
||||
self.payloads.append(payload)
|
||||
def apply_async(self, *, args: list[dict[str, str]], **kwargs: object) -> None:
|
||||
self.payloads.append(args[0])
|
||||
self.options.append(kwargs)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -137,6 +154,8 @@ def database(sqlite_engine: Engine, sqlite_session: Session) -> Iterator[Session
|
||||
@pytest.fixture
|
||||
def trace_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
|
||||
monkeypatch.setattr(module, "provider_config_map", FakeProviderMap({"dummy": PROVIDER_ENTRY}))
|
||||
monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap({}))
|
||||
monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", False)
|
||||
OpsTraceManager.ops_trace_instances_cache.clear()
|
||||
OpsTraceManager.decrypted_configs_cache.clear()
|
||||
monkeypatch.setattr(module.threading, "Timer", DummyTimer)
|
||||
@@ -346,6 +365,84 @@ def test_ops_trace_instance_uses_persisted_enabled_state_and_cache(
|
||||
assert OpsTraceManager.get_ops_trace_instance("missing") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("enabled", "registered", "expected_type"),
|
||||
[
|
||||
(False, False, DummyTraceInstance),
|
||||
(False, True, DummyTraceInstance),
|
||||
(True, False, DummyTraceInstance),
|
||||
(True, True, DummyUnifiedTraceInstance),
|
||||
],
|
||||
)
|
||||
def test_ops_trace_instance_routes_by_unified_switch(
|
||||
enabled: bool,
|
||||
registered: bool,
|
||||
expected_type: type[DummyTraceInstance],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
trace_environment: None,
|
||||
encryption_functions,
|
||||
database: Session,
|
||||
) -> None:
|
||||
app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"}))
|
||||
database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={}))
|
||||
database.commit()
|
||||
monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", enabled)
|
||||
entries = {"dummy": UNIFIED_PROVIDER_ENTRY} if registered else {}
|
||||
monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap(entries))
|
||||
|
||||
instance = OpsTraceManager.get_ops_trace_instance(app.id)
|
||||
|
||||
assert type(instance) is expected_type
|
||||
|
||||
|
||||
def test_registered_unified_provider_does_not_fallback_when_construction_fails(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
trace_environment: None,
|
||||
encryption_functions,
|
||||
database: Session,
|
||||
) -> None:
|
||||
app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"}))
|
||||
database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={}))
|
||||
database.commit()
|
||||
monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", True)
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"unified_provider_config_map",
|
||||
FakeProviderMap(
|
||||
{
|
||||
"dummy": {
|
||||
"config_class": DummyConfig,
|
||||
"trace_instance": FailingUnifiedTraceInstance,
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="unified constructor failed"):
|
||||
OpsTraceManager.get_ops_trace_instance(app.id)
|
||||
|
||||
|
||||
def test_unified_and_legacy_instances_have_separate_cache_entries(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
trace_environment: None,
|
||||
encryption_functions,
|
||||
database: Session,
|
||||
) -> None:
|
||||
app = _app(database, tracing=json.dumps({"enabled": True, "tracing_provider": "dummy"}))
|
||||
database.add(TraceAppConfig(app_id=app.id, tracing_provider="dummy", tracing_config={}))
|
||||
database.commit()
|
||||
monkeypatch.setattr(module, "unified_provider_config_map", FakeProviderMap({"dummy": UNIFIED_PROVIDER_ENTRY}))
|
||||
|
||||
monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", False)
|
||||
legacy = OpsTraceManager.get_ops_trace_instance(app.id)
|
||||
monkeypatch.setattr(dify_config, "OPS_TRACE_UNIFIED_ENABLED", True)
|
||||
unified = OpsTraceManager.get_ops_trace_instance(app.id)
|
||||
|
||||
assert type(legacy) is DummyTraceInstance
|
||||
assert type(unified) is DummyUnifiedTraceInstance
|
||||
assert legacy is not unified
|
||||
|
||||
|
||||
def test_message_config_lookup_uses_real_conversation_and_model_config(database: Session) -> None:
|
||||
app = _app(database)
|
||||
config = AppModelConfig(app_id=app.id, model='{"provider":"openai"}')
|
||||
@@ -402,7 +499,10 @@ def test_message_trace_reads_real_conversation_app_and_message_file(
|
||||
database.add(file)
|
||||
database.commit()
|
||||
monkeypatch.setattr(module, "get_message_data", lambda _message_id: _message_data())
|
||||
result = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE, message_id=message.id).message_trace(message.id)
|
||||
result = TraceTask(
|
||||
trace_type=TraceTaskName.MESSAGE_TRACE,
|
||||
message_id=message.id,
|
||||
).preprocess()
|
||||
assert result.message_id == message.id
|
||||
assert result.conversation_mode == AppMode.CHAT
|
||||
assert result.file_list[0].endswith("path/to/file")
|
||||
@@ -565,7 +665,7 @@ def test_trace_helpers_and_streaming_metrics(trace_environment: None) -> None:
|
||||
assert OpsTraceManager.check_trace_config_is_effective({}, "dummy")
|
||||
assert OpsTraceManager.get_trace_config_project_key({}, "dummy") == "fake-key"
|
||||
assert OpsTraceManager.get_trace_config_project_url({}, "dummy") == "https://project.fake"
|
||||
task = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE)
|
||||
task = TraceTask(trace_type=TraceTaskName.MESSAGE_TRACE, message_id="message-1")
|
||||
assert task.conversation_trace(foo="bar") == {"foo": "bar"}
|
||||
assert task._extract_streaming_metrics(_message_data(message_metadata="invalid")) == {}
|
||||
assert task.generate_name_trace("conversation", {"start": 1, "end": 2}, tenant_id=None) == {}
|
||||
@@ -578,6 +678,7 @@ def test_trace_helpers_and_streaming_metrics(trace_environment: None) -> None:
|
||||
)
|
||||
assert generated.outputs == "name"
|
||||
assert generated.tenant_id == "tenant-1"
|
||||
assert generated.message_id == "message-1"
|
||||
|
||||
|
||||
def test_trace_queue_collect_run_and_storage_boundary(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None:
|
||||
@@ -597,7 +698,7 @@ def test_trace_queue_collect_run_and_storage_boundary(monkeypatch: pytest.Monkey
|
||||
recording_storage = RecordingStorage()
|
||||
dispatcher = RecordingDispatcher()
|
||||
monkeypatch.setattr(module.storage, "save", recording_storage.save)
|
||||
monkeypatch.setattr(module.process_trace_tasks, "delay", dispatcher.delay)
|
||||
monkeypatch.setattr(module.process_trace_tasks, "apply_async", dispatcher.apply_async)
|
||||
file_id = UUID("00000000-0000-0000-0000-000000000123")
|
||||
monkeypatch.setattr(module, "uuid4", lambda: file_id)
|
||||
manager.add_trace_task(task)
|
||||
@@ -608,3 +709,55 @@ def test_trace_queue_collect_run_and_storage_boundary(monkeypatch: pytest.Monkey
|
||||
assert path.endswith(f"app-id/{file_id.hex}.json")
|
||||
assert json.loads(data)["app_id"] == "app-id"
|
||||
assert dispatcher.payloads == [{"file_id": file_id.hex, "app_id": "app-id"}]
|
||||
|
||||
|
||||
def test_trace_queue_persists_with_caller_supplied_file_id(
|
||||
monkeypatch: pytest.MonkeyPatch, trace_environment: None
|
||||
) -> None:
|
||||
monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True))
|
||||
manager = TraceQueueManager(app_id="app-id", user_id="user-1")
|
||||
task = TraceTask(
|
||||
trace_type=TraceTaskName.GENERATE_NAME_TRACE,
|
||||
conversation_id="conversation-1",
|
||||
timer={"start": 1, "end": 2},
|
||||
tenant_id="tenant-1",
|
||||
generate_conversation_name="name",
|
||||
inputs="query",
|
||||
)
|
||||
recording_storage = RecordingStorage()
|
||||
monkeypatch.setattr(module.storage, "save", recording_storage.save)
|
||||
|
||||
file_info = manager.persist_trace_task(task, file_id="workflow-final-run-1")
|
||||
|
||||
assert file_info == {"file_id": "workflow-final-run-1", "app_id": "app-id"}
|
||||
path, data = recording_storage.writes[0]
|
||||
payload = json.loads(data)
|
||||
assert path == "ops_trace/app-id/workflow-final-run-1.json"
|
||||
assert UUID(payload["trace_info"]["operation_id"])
|
||||
|
||||
|
||||
def test_trace_queue_persistence_error_propagates(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None:
|
||||
monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True))
|
||||
manager = TraceQueueManager(app_id="app-id", user_id="user-1")
|
||||
task = TraceTask(trace_type=TraceTaskName.GENERATE_NAME_TRACE)
|
||||
|
||||
def fail_save(_path: str, _data: bytes) -> None:
|
||||
raise OSError("storage unavailable")
|
||||
|
||||
monkeypatch.setattr(module.storage, "save", fail_save)
|
||||
|
||||
with pytest.raises(OSError, match="storage unavailable"):
|
||||
manager.persist_trace_task(task, file_id="workflow-final-run-1")
|
||||
|
||||
|
||||
def test_trace_queue_enqueue_error_propagates(monkeypatch: pytest.MonkeyPatch, trace_environment: None) -> None:
|
||||
monkeypatch.setattr(OpsTraceManager, "get_ops_trace_instance", classmethod(lambda cls, _app_id: True))
|
||||
manager = TraceQueueManager(app_id="app-id", user_id="user-1")
|
||||
|
||||
def fail_enqueue(*_args: object, **_kwargs: object) -> None:
|
||||
raise ConnectionError("broker unavailable")
|
||||
|
||||
monkeypatch.setattr(module.process_trace_tasks, "apply_async", fail_enqueue)
|
||||
|
||||
with pytest.raises(ConnectionError, match="broker unavailable"):
|
||||
manager.enqueue_persisted_trace({"file_id": "workflow-final-run-1", "app_id": "app-id"})
|
||||
Whitespace-only changes.
@@ -0,0 +1,108 @@
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from core.helper.trace_id_helper import ParentTraceContext
|
||||
from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace
|
||||
|
||||
|
||||
def span(span_id: str, parent_id: str | None = None) -> CanonicalSpan:
|
||||
return CanonicalSpan(
|
||||
id=span_id,
|
||||
parent_id=parent_id,
|
||||
name=span_id,
|
||||
kind=CanonicalSpanKind.CHAIN,
|
||||
start_time=datetime(2025, 1, 1),
|
||||
end_time=None,
|
||||
status=CanonicalSpanStatus.OK,
|
||||
)
|
||||
|
||||
|
||||
def test_canonical_trace_rejects_unknown_fields() -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
CanonicalTrace(
|
||||
trace_id="trace-1",
|
||||
session_id="session-1",
|
||||
root_span_id="root-1",
|
||||
spans=(),
|
||||
unknown=True, # pyrefly: ignore[unexpected-keyword]
|
||||
)
|
||||
|
||||
|
||||
def test_canonical_span_supports_links() -> None:
|
||||
linked_span = CanonicalSpan(
|
||||
id="linked-1",
|
||||
parent_id="message-2",
|
||||
name="linked",
|
||||
kind=CanonicalSpanKind.CHAIN,
|
||||
start_time=datetime(2025, 1, 1),
|
||||
end_time=None,
|
||||
status=CanonicalSpanStatus.OK,
|
||||
links=("message-1",),
|
||||
)
|
||||
|
||||
assert linked_span.links == ("message-1",)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("root_span_id", "spans"),
|
||||
[
|
||||
("missing", (span("root"),)),
|
||||
("root", (span("root"), span("root", "root"))),
|
||||
("root", (span("root", "outside"),)),
|
||||
("root", (span("root"), span("child"))),
|
||||
("root", (span("root"), span("child", "later"), span("later", "root"))),
|
||||
],
|
||||
)
|
||||
def test_canonical_trace_rejects_invalid_fragment(
|
||||
root_span_id: str,
|
||||
spans: tuple[CanonicalSpan, ...],
|
||||
) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
CanonicalTrace(
|
||||
trace_id="trace-1",
|
||||
session_id="session-1",
|
||||
root_span_id=root_span_id,
|
||||
spans=spans,
|
||||
)
|
||||
|
||||
|
||||
def test_canonical_trace_rejects_conflicting_external_parent_modes() -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
CanonicalTrace(
|
||||
trace_id="trace-1",
|
||||
session_id="session-1",
|
||||
root_span_id="root",
|
||||
spans=(span("root"),),
|
||||
external_parent=ParentTraceContext(
|
||||
parent_workflow_run_id="outer-run",
|
||||
parent_node_execution_id="outer-node",
|
||||
),
|
||||
required_parent_context_id="message-1",
|
||||
)
|
||||
|
||||
|
||||
def test_canonical_span_is_immutable() -> None:
|
||||
span = CanonicalSpan(
|
||||
id="span-1",
|
||||
parent_id=None,
|
||||
name="root",
|
||||
kind=CanonicalSpanKind.CHAIN,
|
||||
start_time=datetime(2025, 1, 1),
|
||||
end_time=None,
|
||||
status=CanonicalSpanStatus.OK,
|
||||
)
|
||||
|
||||
assert span.publishes_parent_context is False
|
||||
|
||||
trace = CanonicalTrace(
|
||||
trace_id="trace-1",
|
||||
session_id="session-1",
|
||||
root_span_id=span.id,
|
||||
spans=(span,),
|
||||
)
|
||||
assert trace.required_parent_context_id is None
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
span.name = "changed" # pyrefly: ignore[read-only]
|
||||
@@ -0,0 +1,107 @@
|
||||
from datetime import datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import cast
|
||||
|
||||
from core.ops.unified_trace.hierarchy import WorkflowExecutionLike, build_workflow_hierarchy
|
||||
|
||||
|
||||
def execution(**overrides: object) -> WorkflowExecutionLike:
|
||||
values: dict[str, object] = {
|
||||
"id": "exec-1",
|
||||
"node_execution_id": None,
|
||||
"node_id": "node-1",
|
||||
"node_type": "tool",
|
||||
"predecessor_node_id": None,
|
||||
"iteration_id": None,
|
||||
"iteration_index": None,
|
||||
"loop_id": None,
|
||||
"loop_index": None,
|
||||
"created_at": datetime(2025, 1, 1),
|
||||
"elapsed_time": 1.0,
|
||||
"status": "succeeded",
|
||||
"metadata": {},
|
||||
}
|
||||
values.update(overrides)
|
||||
return cast(WorkflowExecutionLike, SimpleNamespace(**values))
|
||||
|
||||
|
||||
def test_predecessor_becomes_parent_independent_of_repository_order() -> None:
|
||||
start = execution(id="exec-start", node_id="start")
|
||||
llm = execution(id="exec-llm", node_id="llm", predecessor_node_id="start")
|
||||
|
||||
forward = build_workflow_hierarchy([start, llm])
|
||||
reverse = build_workflow_hierarchy([llm, start])
|
||||
|
||||
assert forward.parent_by_execution_id == {"exec-llm": "exec-start"}
|
||||
assert reverse.parent_by_execution_id == forward.parent_by_execution_id
|
||||
|
||||
|
||||
def test_repeated_graph_node_id_is_not_guessed_as_parent() -> None:
|
||||
first = execution(id="exec-a1", node_id="a")
|
||||
second = execution(id="exec-a2", node_id="a")
|
||||
child = execution(id="exec-b", node_id="b", predecessor_node_id="a")
|
||||
|
||||
result = build_workflow_hierarchy([first, second, child])
|
||||
|
||||
assert "exec-b" not in result.parent_by_execution_id
|
||||
|
||||
|
||||
def test_iteration_child_is_parented_to_stable_wrapper() -> None:
|
||||
container = execution(id="iteration-exec", node_id="iteration", node_type="iteration")
|
||||
child = execution(id="child-exec", node_id="child", iteration_id="iteration", iteration_index=0)
|
||||
|
||||
result = build_workflow_hierarchy([child, container])
|
||||
wrapper = result.wrapper_by_child_execution_id["child-exec"]
|
||||
|
||||
assert wrapper.id == "iteration:iteration-exec:0"
|
||||
assert wrapper.parent_execution_id == "iteration-exec"
|
||||
assert result.parent_by_execution_id["child-exec"] == wrapper.id
|
||||
|
||||
|
||||
def test_loop_wrapper_covers_child_times_and_failure() -> None:
|
||||
container = execution(id="loop-exec", node_id="loop", node_type="loop")
|
||||
first = execution(
|
||||
id="first",
|
||||
node_id="first-node",
|
||||
loop_id="loop",
|
||||
loop_index=2,
|
||||
created_at=datetime(2025, 1, 1, 0, 0, 1),
|
||||
elapsed_time=2,
|
||||
)
|
||||
second = execution(
|
||||
id="second",
|
||||
node_id="second-node",
|
||||
loop_id="loop",
|
||||
loop_index=2,
|
||||
created_at=datetime(2025, 1, 1, 0, 0, 2),
|
||||
elapsed_time=4,
|
||||
status="failed",
|
||||
)
|
||||
|
||||
result = build_workflow_hierarchy([second, container, first])
|
||||
wrapper = result.wrappers[0]
|
||||
|
||||
assert wrapper.id == "loop:loop-exec:2"
|
||||
assert wrapper.start_time == first.created_at
|
||||
assert wrapper.end_time == second.created_at + timedelta(seconds=4)
|
||||
assert wrapper.has_error is True
|
||||
assert wrapper.child_execution_ids == frozenset({"first", "second"})
|
||||
|
||||
|
||||
def test_invalid_wrapper_indexes_do_not_create_wrappers() -> None:
|
||||
container = execution(id="loop-exec", node_id="loop", node_type="loop")
|
||||
negative = execution(id="negative", node_id="negative-node", loop_id="loop", loop_index=-1)
|
||||
boolean = execution(id="boolean", node_id="boolean-node", loop_id="loop", loop_index=True)
|
||||
|
||||
result = build_workflow_hierarchy([container, negative, boolean])
|
||||
|
||||
assert result.wrappers == ()
|
||||
|
||||
|
||||
def test_cycle_edges_are_removed_deterministically() -> None:
|
||||
first = execution(id="a-exec", node_id="a", predecessor_node_id="b")
|
||||
second = execution(id="b-exec", node_id="b", predecessor_node_id="a")
|
||||
|
||||
result = build_workflow_hierarchy([first, second])
|
||||
|
||||
assert result.parent_by_execution_id == {}
|
||||
@@ -0,0 +1,164 @@
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from core.helper.trace_id_helper import ParentTraceContext
|
||||
from core.ops.exceptions import (
|
||||
InvalidTraceParentContextError,
|
||||
PendingTraceParentContextError,
|
||||
TraceParentContextAccessError,
|
||||
)
|
||||
from core.ops.unified_trace.parent_context import (
|
||||
ParentContextCoordinator,
|
||||
ParentDestination,
|
||||
ParentResolutionKind,
|
||||
ProviderParentContext,
|
||||
destination_scope,
|
||||
parent_destination_from_config,
|
||||
)
|
||||
|
||||
|
||||
def parent() -> ParentTraceContext:
|
||||
return ParentTraceContext(parent_workflow_run_id="outer-run", parent_node_execution_id="outer-tool")
|
||||
|
||||
|
||||
def context(**overrides: object) -> ProviderParentContext:
|
||||
values: dict[str, object] = {
|
||||
"provider": "langsmith",
|
||||
"scope": "scope-a",
|
||||
"trace_id": "root-run",
|
||||
"parent_id": "outer-tool",
|
||||
"provider_context": {"dotted_order": "root.tool"},
|
||||
}
|
||||
values.update(overrides)
|
||||
return ProviderParentContext.model_validate(values)
|
||||
|
||||
|
||||
def coordinator(redis: MagicMock, destination: ParentDestination | None) -> ParentContextCoordinator:
|
||||
return ParentContextCoordinator(redis, lambda _workflow_run_id: destination)
|
||||
|
||||
|
||||
def test_parent_destination_uses_non_secret_provider_scope() -> None:
|
||||
destination = parent_destination_from_config(
|
||||
"langsmith",
|
||||
{"api_key": "secret", "endpoint": "https://smith.example", "project": "project-a"},
|
||||
unified=True,
|
||||
)
|
||||
|
||||
assert destination == ParentDestination(
|
||||
provider="langsmith",
|
||||
scope=destination_scope("langsmith", "https://smith.example", "project-a"),
|
||||
unified=True,
|
||||
)
|
||||
assert "secret" not in destination.scope
|
||||
|
||||
|
||||
def test_publish_uses_unified_namespace_and_configured_ttl(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
redis = MagicMock()
|
||||
value = context()
|
||||
monkeypatch.setattr(
|
||||
"core.ops.unified_trace.parent_context.dify_config",
|
||||
SimpleNamespace(OPS_TRACE_PARENT_CONTEXT_TTL_SECONDS=1_800),
|
||||
)
|
||||
|
||||
coordinator(redis, None).publish("outer-tool", value)
|
||||
|
||||
key, ttl, payload = redis.setex.call_args.args
|
||||
assert key == "trace:unified:parent:outer-tool"
|
||||
assert ttl == 1_800
|
||||
assert json.loads(payload)["provider"] == "langsmith"
|
||||
|
||||
|
||||
def test_resolve_returns_compatible_context() -> None:
|
||||
redis = MagicMock()
|
||||
redis.get.return_value = context().model_dump_json().encode()
|
||||
subject = coordinator(redis, ParentDestination(provider="langsmith", scope="scope-a", unified=True))
|
||||
|
||||
result = subject.resolve(parent(), expected_provider="langsmith", expected_scope="scope-a")
|
||||
|
||||
assert result.kind is ParentResolutionKind.RESTORED
|
||||
assert result.context == context()
|
||||
assert result.linked_parent is None
|
||||
|
||||
|
||||
def test_resolve_required_restores_message_context_without_destination_lookup() -> None:
|
||||
redis = MagicMock()
|
||||
redis.get.return_value = context().model_dump_json().encode()
|
||||
resolve_destination = MagicMock()
|
||||
subject = ParentContextCoordinator(redis, resolve_destination)
|
||||
|
||||
result = subject.resolve_required(
|
||||
"message-1",
|
||||
expected_provider="langsmith",
|
||||
expected_scope="scope-a",
|
||||
)
|
||||
|
||||
assert result.kind is ParentResolutionKind.RESTORED
|
||||
assert result.context == context()
|
||||
redis.get.assert_called_once_with("trace:unified:parent:message-1")
|
||||
resolve_destination.assert_not_called()
|
||||
|
||||
|
||||
def test_missing_required_message_context_is_retryable() -> None:
|
||||
redis = MagicMock()
|
||||
redis.get.return_value = None
|
||||
subject = coordinator(redis, None)
|
||||
|
||||
with pytest.raises(PendingTraceParentContextError):
|
||||
subject.resolve_required(
|
||||
"message-1",
|
||||
expected_provider="langsmith",
|
||||
expected_scope="scope-a",
|
||||
)
|
||||
|
||||
|
||||
def test_missing_compatible_context_is_retryable() -> None:
|
||||
redis = MagicMock()
|
||||
redis.get.return_value = None
|
||||
subject = coordinator(redis, ParentDestination(provider="langsmith", scope="scope-a", unified=True))
|
||||
|
||||
with pytest.raises(PendingTraceParentContextError):
|
||||
subject.resolve(parent(), expected_provider="langsmith", expected_scope="scope-a")
|
||||
|
||||
|
||||
def test_non_unified_or_incompatible_parent_becomes_linked_root() -> None:
|
||||
redis = MagicMock()
|
||||
destinations = [
|
||||
None,
|
||||
ParentDestination(provider="langsmith", scope="scope-a", unified=False),
|
||||
ParentDestination(provider="phoenix", scope="scope-a", unified=True),
|
||||
ParentDestination(provider="langsmith", scope="scope-b", unified=True),
|
||||
]
|
||||
|
||||
for destination in destinations:
|
||||
result = coordinator(redis, destination).resolve(
|
||||
parent(), expected_provider="langsmith", expected_scope="scope-a"
|
||||
)
|
||||
assert result.kind is ParentResolutionKind.LINKED_ROOT
|
||||
assert result.linked_parent == parent()
|
||||
|
||||
redis.get.assert_not_called()
|
||||
|
||||
|
||||
def test_malformed_or_stale_context_is_terminal() -> None:
|
||||
redis = MagicMock()
|
||||
subject = coordinator(redis, ParentDestination(provider="langsmith", scope="scope-a", unified=True))
|
||||
|
||||
for payload in (b"not-json", b'{"version": 2}', context(scope="scope-b").model_dump_json().encode()):
|
||||
redis.get.return_value = payload
|
||||
with pytest.raises(InvalidTraceParentContextError):
|
||||
subject.resolve(parent(), expected_provider="langsmith", expected_scope="scope-a")
|
||||
|
||||
|
||||
def test_redis_read_and_write_failures_are_retryable() -> None:
|
||||
redis = MagicMock()
|
||||
redis.get.side_effect = ConnectionError("down")
|
||||
redis.setex.side_effect = ConnectionError("down")
|
||||
subject = coordinator(redis, ParentDestination(provider="langsmith", scope="scope-a", unified=True))
|
||||
|
||||
with pytest.raises(TraceParentContextAccessError):
|
||||
subject.resolve(parent(), expected_provider="langsmith", expected_scope="scope-a")
|
||||
with pytest.raises(TraceParentContextAccessError):
|
||||
subject.publish("outer-tool", context())
|
||||
@@ -0,0 +1,120 @@
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from core.helper.trace_id_helper import ParentTraceContext
|
||||
from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace
|
||||
from core.ops.unified_trace.parent_context import ParentResolution
|
||||
from core.ops.unified_trace.provider import UnifiedTraceInstance
|
||||
|
||||
|
||||
def canonical_trace(*, nested: bool = False, required_parent_context_id: str | None = None) -> CanonicalTrace:
|
||||
parent = (
|
||||
ParentTraceContext(parent_workflow_run_id="outer-run", parent_node_execution_id="outer-tool")
|
||||
if nested
|
||||
else None
|
||||
)
|
||||
return CanonicalTrace(
|
||||
trace_id="trace-1",
|
||||
session_id="session-1",
|
||||
root_span_id="root-1",
|
||||
external_parent=parent,
|
||||
required_parent_context_id=required_parent_context_id,
|
||||
spans=(
|
||||
CanonicalSpan(
|
||||
id="root-1",
|
||||
parent_id=None,
|
||||
name="root",
|
||||
kind=CanonicalSpanKind.CHAIN,
|
||||
start_time=datetime(2025, 1, 1),
|
||||
end_time=datetime(2025, 1, 1, 0, 0, 1),
|
||||
status=CanonicalSpanStatus.OK,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def make_runtime(trace: CanonicalTrace) -> tuple[UnifiedTraceInstance, MagicMock, MagicMock, MagicMock]:
|
||||
builder = MagicMock()
|
||||
builder.build.return_value = trace
|
||||
adapter = MagicMock(provider_name="langsmith", scope="scope-a")
|
||||
adapter.provider_name = "langsmith"
|
||||
adapter.scope = "scope-a"
|
||||
coordinator = MagicMock()
|
||||
runtime = UnifiedTraceInstance(MagicMock(), builder=builder, adapter=adapter, coordinator=coordinator)
|
||||
return runtime, builder, adapter, coordinator
|
||||
|
||||
|
||||
def test_runtime_passes_core_publisher_to_adapter() -> None:
|
||||
runtime, _, adapter, coordinator = make_runtime(canonical_trace())
|
||||
provider_context = MagicMock()
|
||||
|
||||
def emit(
|
||||
_trace: CanonicalTrace,
|
||||
_parent: ParentResolution | None,
|
||||
publish_parent_context: Callable[[str, object], None],
|
||||
) -> None:
|
||||
publish_parent_context("tool-exec", provider_context)
|
||||
|
||||
adapter.emit.side_effect = emit
|
||||
|
||||
runtime.trace(MagicMock())
|
||||
|
||||
coordinator.publish.assert_called_once_with("tool-exec", provider_context)
|
||||
|
||||
|
||||
def test_runtime_resolves_nested_parent_before_emission() -> None:
|
||||
runtime, _, adapter, coordinator = make_runtime(canonical_trace(nested=True))
|
||||
resolution = ParentResolution.restored(MagicMock())
|
||||
coordinator.resolve.return_value = resolution
|
||||
|
||||
runtime.trace(MagicMock())
|
||||
|
||||
coordinator.resolve.assert_called_once()
|
||||
assert coordinator.resolve.call_args.kwargs == {
|
||||
"expected_provider": "langsmith",
|
||||
"expected_scope": "scope-a",
|
||||
}
|
||||
adapter.emit.assert_called_once()
|
||||
assert adapter.emit.call_args.args[1] is resolution
|
||||
|
||||
|
||||
def test_runtime_resolves_required_message_parent_before_emission() -> None:
|
||||
runtime, _, adapter, coordinator = make_runtime(canonical_trace(required_parent_context_id="message-1"))
|
||||
resolution = ParentResolution.restored(MagicMock())
|
||||
events: list[str] = []
|
||||
coordinator.resolve_required.side_effect = lambda *_args, **_kwargs: events.append("resolve") or resolution
|
||||
adapter.emit.side_effect = lambda *_args: events.append("emit")
|
||||
|
||||
runtime.trace(MagicMock())
|
||||
|
||||
assert events == ["resolve", "emit"]
|
||||
coordinator.resolve_required.assert_called_once_with(
|
||||
"message-1",
|
||||
expected_provider="langsmith",
|
||||
expected_scope="scope-a",
|
||||
)
|
||||
coordinator.resolve.assert_not_called()
|
||||
assert adapter.emit.call_args.args[1] is resolution
|
||||
|
||||
|
||||
def test_runtime_does_not_publish_when_adapter_fails_before_callback() -> None:
|
||||
runtime, _, adapter, coordinator = make_runtime(canonical_trace())
|
||||
adapter.emit.side_effect = RuntimeError("provider rejected run")
|
||||
|
||||
with pytest.raises(RuntimeError, match="provider rejected run"):
|
||||
runtime.trace(MagicMock())
|
||||
|
||||
coordinator.publish.assert_not_called()
|
||||
|
||||
|
||||
def test_runtime_has_no_legacy_fallback() -> None:
|
||||
runtime, _, adapter, _ = make_runtime(canonical_trace())
|
||||
adapter.emit.side_effect = RuntimeError("terminal")
|
||||
|
||||
with pytest.raises(RuntimeError, match="terminal"):
|
||||
runtime.trace(MagicMock())
|
||||
|
||||
assert not hasattr(runtime, "legacy_provider")
|
||||
@@ -0,0 +1,14 @@
|
||||
import pytest
|
||||
|
||||
from core.ops.entities.config_entity import TracingProviderEnum
|
||||
from core.ops.unified_trace.registry import unified_provider_config_map
|
||||
|
||||
|
||||
def test_registry_exposes_only_implemented_providers() -> None:
|
||||
phoenix = unified_provider_config_map[TracingProviderEnum.PHOENIX]
|
||||
langsmith = unified_provider_config_map[TracingProviderEnum.LANGSMITH]
|
||||
|
||||
assert phoenix["trace_instance"].__name__ == "UnifiedPhoenixTrace"
|
||||
assert langsmith["trace_instance"].__name__ == "UnifiedLangSmithTrace"
|
||||
with pytest.raises(KeyError):
|
||||
unified_provider_config_map[TracingProviderEnum.LANGFUSE]
|
||||
@@ -0,0 +1,536 @@
|
||||
from copy import deepcopy
|
||||
from datetime import datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import cast
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
|
||||
from core.ops.entities.trace_entity import (
|
||||
BaseTraceInfo,
|
||||
DatasetRetrievalTraceInfo,
|
||||
GenerateNameTraceInfo,
|
||||
MessageTraceInfo,
|
||||
ModerationTraceInfo,
|
||||
SuggestedQuestionTraceInfo,
|
||||
ToolTraceInfo,
|
||||
WorkflowTraceInfo,
|
||||
)
|
||||
from core.ops.unified_trace.entities import CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace
|
||||
from core.ops.unified_trace.trace_builder import (
|
||||
CanonicalTraceBuilder,
|
||||
RepositoryWorkflowExecutionLoader,
|
||||
WorkflowExecutionLike,
|
||||
resolve_session_id,
|
||||
)
|
||||
|
||||
|
||||
def make_workflow_trace_info(**overrides: object) -> WorkflowTraceInfo:
|
||||
values: dict[str, object] = {
|
||||
"workflow_data": SimpleNamespace(),
|
||||
"conversation_id": None,
|
||||
"workflow_id": "workflow-1",
|
||||
"tenant_id": "tenant-1",
|
||||
"workflow_run_id": "run-1",
|
||||
"workflow_run_elapsed_time": 1.0,
|
||||
"workflow_run_status": "succeeded",
|
||||
"workflow_run_inputs": {},
|
||||
"workflow_run_outputs": {},
|
||||
"workflow_run_version": "1",
|
||||
"total_tokens": 0,
|
||||
"file_list": [],
|
||||
"query": "",
|
||||
"metadata": {},
|
||||
}
|
||||
values.update(overrides)
|
||||
return WorkflowTraceInfo.model_validate(values)
|
||||
|
||||
|
||||
def test_custom_session_id_wins_over_conversation_id() -> None:
|
||||
info = make_workflow_trace_info(
|
||||
conversation_id="conversation-1",
|
||||
metadata={"trace_session_id": "customer-session"},
|
||||
)
|
||||
|
||||
assert resolve_session_id(info) == "customer-session"
|
||||
|
||||
|
||||
def test_workflow_session_falls_back_to_conversation_then_run() -> None:
|
||||
assert resolve_session_id(make_workflow_trace_info(conversation_id="conversation-1")) == "conversation-1"
|
||||
assert resolve_session_id(make_workflow_trace_info()) == "run-1"
|
||||
|
||||
|
||||
def test_nested_workflow_session_falls_back_to_parent_workflow() -> None:
|
||||
info = make_workflow_trace_info(
|
||||
metadata={
|
||||
"parent_trace_context": {
|
||||
"parent_workflow_run_id": "parent-run",
|
||||
"parent_node_execution_id": "parent-node-execution",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
assert resolve_session_id(info) == "parent-run"
|
||||
|
||||
|
||||
def test_message_session_falls_back_to_message_conversation() -> None:
|
||||
info = MessageTraceInfo(
|
||||
conversation_model="chat",
|
||||
message_tokens=0,
|
||||
answer_tokens=0,
|
||||
total_tokens=0,
|
||||
conversation_mode="chat",
|
||||
message_data=SimpleNamespace(conversation_id="conversation-1", created_at=datetime(2025, 1, 1)),
|
||||
metadata={},
|
||||
)
|
||||
|
||||
assert resolve_session_id(info) == "conversation-1"
|
||||
|
||||
|
||||
def node_execution(**overrides: object) -> WorkflowExecutionLike:
|
||||
values: dict[str, object] = {
|
||||
"id": "node-exec",
|
||||
"node_execution_id": None,
|
||||
"node_id": "node",
|
||||
"title": "Node",
|
||||
"node_type": "tool",
|
||||
"predecessor_node_id": None,
|
||||
"iteration_id": None,
|
||||
"iteration_index": None,
|
||||
"loop_id": None,
|
||||
"loop_index": None,
|
||||
"created_at": datetime(2025, 1, 1),
|
||||
"elapsed_time": 1.0,
|
||||
"status": "succeeded",
|
||||
"error": None,
|
||||
"inputs": {"input": "value"},
|
||||
"outputs": {"output": "value"},
|
||||
"process_data": {},
|
||||
"metadata": {},
|
||||
}
|
||||
values.update(overrides)
|
||||
return cast(WorkflowExecutionLike, SimpleNamespace(**values))
|
||||
|
||||
|
||||
def workflow_info(**overrides: object) -> WorkflowTraceInfo:
|
||||
workflow_data = SimpleNamespace(
|
||||
created_at=datetime(2025, 1, 1),
|
||||
finished_at=datetime(2025, 1, 1) + timedelta(seconds=5),
|
||||
graph_dict={"nodes": []},
|
||||
)
|
||||
values: dict[str, object] = {
|
||||
"workflow_data": workflow_data,
|
||||
"start_time": workflow_data.created_at,
|
||||
"end_time": workflow_data.finished_at,
|
||||
"metadata": {"app_id": "app-1"},
|
||||
}
|
||||
values.update(overrides)
|
||||
return make_workflow_trace_info(**values)
|
||||
|
||||
|
||||
CHATFLOW_INPUTS: dict[str, object] = {
|
||||
"sys.app_id": "19a6d372-b8bc-4ad4-9b83-7e6e7138de31",
|
||||
"sys.dialogue_count": 1,
|
||||
"sys.files": [],
|
||||
"sys.query": "hi",
|
||||
"sys.user_id": "ca877a63-4d75-4ba3-a417-3edffe5e545c",
|
||||
"sys.workflow_id": "8b81be9e-d7c1-4fa7-b90f-03791fa015ba",
|
||||
"sys.workflow_run_id": "4eec02ea-4ed5-47cc-87fd-7c3821dd935d",
|
||||
}
|
||||
|
||||
|
||||
def test_chatflow_message_uses_query_while_workflow_keeps_complete_inputs() -> None:
|
||||
builder = CanonicalTraceBuilder(lambda _info: [])
|
||||
|
||||
trace = builder.build(
|
||||
workflow_info(
|
||||
message_id="message-1",
|
||||
query="hi",
|
||||
workflow_run_inputs=CHATFLOW_INPUTS,
|
||||
)
|
||||
)
|
||||
|
||||
assert trace is not None
|
||||
spans = {span.id: span for span in trace.spans}
|
||||
assert trace.root_span_id == "message-1"
|
||||
assert [span.id for span in trace.spans[:2]] == ["message-1", "run-1"]
|
||||
assert spans["message-1"].name == "chatflow_run-1"
|
||||
assert spans["message-1"].inputs == "hi"
|
||||
assert spans["message-1"].metadata["trace_entity_type"] == "message"
|
||||
assert spans["message-1"].publishes_parent_context is True
|
||||
assert spans["run-1"].name == "workflow_run-1"
|
||||
assert spans["run-1"].parent_id == "message-1"
|
||||
assert spans["run-1"].inputs == CHATFLOW_INPUTS
|
||||
|
||||
|
||||
def test_chatflow_message_falls_back_to_complete_inputs_for_empty_query() -> None:
|
||||
builder = CanonicalTraceBuilder(lambda _info: [])
|
||||
|
||||
trace = builder.build(
|
||||
workflow_info(
|
||||
message_id="message-1",
|
||||
query="",
|
||||
workflow_run_inputs=CHATFLOW_INPUTS,
|
||||
)
|
||||
)
|
||||
|
||||
assert trace is not None
|
||||
assert trace.spans[0].inputs == CHATFLOW_INPUTS
|
||||
|
||||
|
||||
def test_workflow_without_message_keeps_complete_inputs_on_root() -> None:
|
||||
builder = CanonicalTraceBuilder(lambda _info: [])
|
||||
|
||||
trace = builder.build(workflow_info(query="hi", workflow_run_inputs=CHATFLOW_INPUTS))
|
||||
|
||||
assert trace is not None
|
||||
assert trace.root_span_id == "run-1"
|
||||
assert trace.spans[0].inputs == CHATFLOW_INPUTS
|
||||
|
||||
|
||||
def test_build_workflow_trace_is_parent_first_and_uses_wrappers() -> None:
|
||||
container = node_execution(id="iteration-exec", node_id="iteration", node_type="iteration", title="Items")
|
||||
child = node_execution(
|
||||
id="llm-exec",
|
||||
node_id="llm",
|
||||
node_type="llm",
|
||||
title="Summarize",
|
||||
iteration_id="iteration",
|
||||
iteration_index=0,
|
||||
process_data={
|
||||
"prompts": [{"role": "user", "text": "hello"}],
|
||||
"model_provider": "openai",
|
||||
"model_name": "gpt-4",
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
||||
},
|
||||
)
|
||||
loader = MagicMock(return_value=[child, container])
|
||||
builder = CanonicalTraceBuilder(loader)
|
||||
|
||||
trace = builder.build(workflow_info())
|
||||
|
||||
assert trace is not None
|
||||
assert [span.id for span in trace.spans] == [
|
||||
"run-1",
|
||||
"iteration-exec",
|
||||
"iteration:iteration-exec:0",
|
||||
"llm-exec",
|
||||
]
|
||||
spans = {span.id: span for span in trace.spans}
|
||||
assert spans["llm-exec"].parent_id == "iteration:iteration-exec:0"
|
||||
assert spans["llm-exec"].kind is CanonicalSpanKind.LLM
|
||||
assert spans["llm-exec"].metadata["total_tokens"] == 30
|
||||
loader.assert_called_once()
|
||||
|
||||
|
||||
def test_workflow_tool_is_marked_as_nested_workflow_parent() -> None:
|
||||
builder = CanonicalTraceBuilder(lambda _info: [node_execution(id="tool-exec")])
|
||||
|
||||
trace = builder.build(workflow_info())
|
||||
|
||||
assert trace is not None
|
||||
assert trace.spans[-1].can_parent_workflow is True
|
||||
|
||||
|
||||
def retry_attempt(retry_index: object, **overrides: object) -> dict[str, object]:
|
||||
values: dict[str, object] = {
|
||||
"retry_index": retry_index,
|
||||
"inputs": {"attempt": retry_index},
|
||||
"process_data": {"request": f"attempt-{retry_index}"},
|
||||
"outputs": {"status_code": 500},
|
||||
"error": f"attempt {retry_index} failed",
|
||||
"elapsed_time": float(retry_index) if isinstance(retry_index, int) else 0.0,
|
||||
"execution_metadata": {"internal": True},
|
||||
"created_at": 1_700_000_000,
|
||||
"finished_at": 1_700_000_001,
|
||||
}
|
||||
values.update(overrides)
|
||||
return values
|
||||
|
||||
|
||||
def test_node_metadata_contains_compact_retry_summary_and_skips_malformed_entries() -> None:
|
||||
history = [
|
||||
retry_attempt(1),
|
||||
"malformed",
|
||||
retry_attempt(True),
|
||||
retry_attempt(2, error="attempt 2 timed out", elapsed_time=2.5),
|
||||
retry_attempt(3),
|
||||
]
|
||||
builder = CanonicalTraceBuilder(lambda _info: [node_execution(process_data={"__dify_retry_history": history})])
|
||||
|
||||
trace = builder.build(workflow_info())
|
||||
|
||||
assert trace is not None
|
||||
metadata = trace.spans[-1].metadata
|
||||
assert metadata["retry_count"] == 3
|
||||
assert metadata["retry_attempts"] == [
|
||||
{
|
||||
"retry_index": 1,
|
||||
"error": "attempt 1 failed",
|
||||
"elapsed_time": 1.0,
|
||||
"created_at": 1_700_000_000,
|
||||
"finished_at": 1_700_000_001,
|
||||
},
|
||||
{
|
||||
"retry_index": 2,
|
||||
"error": "attempt 2 timed out",
|
||||
"elapsed_time": 2.5,
|
||||
"created_at": 1_700_000_000,
|
||||
"finished_at": 1_700_000_001,
|
||||
},
|
||||
{
|
||||
"retry_index": 3,
|
||||
"error": "attempt 3 failed",
|
||||
"elapsed_time": 3.0,
|
||||
"created_at": 1_700_000_000,
|
||||
"finished_at": 1_700_000_001,
|
||||
},
|
||||
]
|
||||
assert all("inputs" not in attempt for attempt in metadata["retry_attempts"])
|
||||
assert all("process_data" not in attempt for attempt in metadata["retry_attempts"])
|
||||
assert all("outputs" not in attempt for attempt in metadata["retry_attempts"])
|
||||
assert all("execution_metadata" not in attempt for attempt in metadata["retry_attempts"])
|
||||
|
||||
|
||||
def test_node_without_retry_history_has_no_retry_metadata() -> None:
|
||||
builder = CanonicalTraceBuilder(lambda _info: [node_execution()])
|
||||
|
||||
trace = builder.build(workflow_info())
|
||||
|
||||
assert trace is not None
|
||||
assert "retry_count" not in trace.spans[-1].metadata
|
||||
assert "retry_attempts" not in trace.spans[-1].metadata
|
||||
|
||||
|
||||
def test_failed_node_preserves_error_without_mutating_trace_metadata() -> None:
|
||||
metadata = {"app_id": "app-1", "custom": {"nested": True}}
|
||||
original = deepcopy(metadata)
|
||||
builder = CanonicalTraceBuilder(
|
||||
lambda _info: [node_execution(status="failed", error="boom", metadata={"attempt": 1})]
|
||||
)
|
||||
|
||||
trace = builder.build(workflow_info(metadata=metadata))
|
||||
|
||||
assert trace is not None
|
||||
assert trace.spans[-1].status is CanonicalSpanStatus.ERROR
|
||||
assert trace.spans[-1].error == "boom"
|
||||
assert metadata == original
|
||||
|
||||
|
||||
def test_nested_workflow_exposes_typed_external_parent() -> None:
|
||||
builder = CanonicalTraceBuilder(lambda _info: [])
|
||||
info = workflow_info(
|
||||
metadata={
|
||||
"app_id": "app-1",
|
||||
"parent_trace_context": {
|
||||
"parent_workflow_run_id": "outer-run",
|
||||
"parent_node_execution_id": "outer-tool",
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
trace = builder.build(info)
|
||||
|
||||
assert trace is not None
|
||||
assert trace.external_parent is not None
|
||||
assert trace.external_parent.parent_node_execution_id == "outer-tool"
|
||||
|
||||
|
||||
def test_repository_loader_scopes_repository_to_trace(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
repository = MagicMock()
|
||||
repository.get_by_workflow_execution.return_value = [node_execution()]
|
||||
factory = MagicMock()
|
||||
factory.create_workflow_node_execution_repository.return_value = repository
|
||||
monkeypatch.setattr("core.ops.unified_trace.trace_builder.DifyCoreRepositoryFactory", factory)
|
||||
monkeypatch.setattr("core.ops.unified_trace.trace_builder.db", MagicMock(engine="engine"))
|
||||
account = MagicMock()
|
||||
get_account = MagicMock(return_value=account)
|
||||
loader = RepositoryWorkflowExecutionLoader(get_account)
|
||||
info = workflow_info()
|
||||
|
||||
result = loader(info)
|
||||
|
||||
assert len(result) == 1
|
||||
get_account.assert_called_once_with("app-1")
|
||||
factory.create_workflow_node_execution_repository.assert_called_once()
|
||||
call = factory.create_workflow_node_execution_repository.call_args.kwargs
|
||||
assert call["tenant_id"] == "tenant-1"
|
||||
assert call["app_id"] == "app-1"
|
||||
repository.get_by_workflow_execution.assert_called_once_with(workflow_execution_id="run-1")
|
||||
|
||||
|
||||
def test_message_trace_does_not_load_workflow_executions() -> None:
|
||||
loader = MagicMock()
|
||||
builder = CanonicalTraceBuilder(loader)
|
||||
info = MessageTraceInfo(
|
||||
conversation_model="chat",
|
||||
message_tokens=2,
|
||||
answer_tokens=3,
|
||||
total_tokens=5,
|
||||
conversation_mode="chat",
|
||||
message_id="message-1",
|
||||
message_data=SimpleNamespace(
|
||||
id="message-1",
|
||||
conversation_id="conversation-1",
|
||||
answer="hello",
|
||||
created_at=datetime(2025, 1, 1),
|
||||
updated_at=datetime(2025, 1, 1, 0, 0, 1),
|
||||
),
|
||||
inputs="hi",
|
||||
metadata={},
|
||||
)
|
||||
|
||||
trace = builder.build(info)
|
||||
|
||||
assert trace is not None
|
||||
assert trace.root_span_id == "message-1"
|
||||
assert [span.name for span in trace.spans] == ["message", "llm"]
|
||||
assert trace.spans[0].outputs == "hello"
|
||||
assert trace.spans[0].metadata["trace_entity_type"] == "message"
|
||||
assert trace.spans[0].publishes_parent_context is True
|
||||
assert trace.spans[1].parent_id == "message-1"
|
||||
assert trace.spans[1].kind is CanonicalSpanKind.LLM
|
||||
assert trace.spans[1].metadata["total_tokens"] == 5
|
||||
loader.assert_not_called()
|
||||
|
||||
|
||||
def test_generate_name_uses_message_parent_and_conversation_session() -> None:
|
||||
builder = CanonicalTraceBuilder(lambda _info: [])
|
||||
info = GenerateNameTraceInfo(
|
||||
tenant_id="tenant-1",
|
||||
conversation_id="conversation-1",
|
||||
message_id="message-1",
|
||||
inputs="title prompt",
|
||||
outputs="title",
|
||||
metadata={},
|
||||
)
|
||||
|
||||
trace = builder.build(info)
|
||||
|
||||
assert trace is not None
|
||||
assert trace.trace_id == "message-1"
|
||||
assert trace.session_id == "conversation-1"
|
||||
assert trace.required_parent_context_id == "message-1"
|
||||
assert trace.spans[0].parent_id is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"info",
|
||||
[
|
||||
ModerationTraceInfo(
|
||||
message_id="message-1",
|
||||
message_data=SimpleNamespace(id="message-1"),
|
||||
flagged=False,
|
||||
action="direct_output",
|
||||
preset_response="",
|
||||
query="hello",
|
||||
metadata={},
|
||||
),
|
||||
SuggestedQuestionTraceInfo(
|
||||
message_id="message-1",
|
||||
message_data=SimpleNamespace(id="message-1"),
|
||||
total_tokens=1,
|
||||
suggested_question=["next"],
|
||||
level="info",
|
||||
metadata={},
|
||||
),
|
||||
DatasetRetrievalTraceInfo(
|
||||
message_id="message-1",
|
||||
message_data=SimpleNamespace(id="message-1"),
|
||||
documents=[],
|
||||
metadata={},
|
||||
),
|
||||
ToolTraceInfo(
|
||||
message_id="message-1",
|
||||
tool_name="search",
|
||||
tool_inputs={},
|
||||
tool_outputs="done",
|
||||
tool_config={},
|
||||
time_cost=0.1,
|
||||
tool_parameters={},
|
||||
metadata={},
|
||||
),
|
||||
GenerateNameTraceInfo(
|
||||
tenant_id="tenant-1",
|
||||
conversation_id="conversation-1",
|
||||
message_id="message-1",
|
||||
metadata={},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_standalone_message_child_uses_explicit_required_parent(info: BaseTraceInfo) -> None:
|
||||
trace = CanonicalTraceBuilder(lambda _info: []).build(info)
|
||||
|
||||
assert trace is not None
|
||||
assert trace.spans[0].parent_id is None
|
||||
assert trace.required_parent_context_id == "message-1"
|
||||
assert_fragment_is_parent_first(trace)
|
||||
|
||||
|
||||
def assert_fragment_is_parent_first(trace: CanonicalTrace) -> None:
|
||||
seen: set[str] = set()
|
||||
ids = [span.id for span in trace.spans]
|
||||
assert len(ids) == len(set(ids))
|
||||
assert trace.root_span_id in ids
|
||||
for span in trace.spans:
|
||||
if span.id == trace.root_span_id:
|
||||
assert span.parent_id is None
|
||||
elif span.parent_id is not None:
|
||||
assert span.parent_id in seen
|
||||
seen.add(span.id)
|
||||
|
||||
|
||||
def test_generate_name_without_message_remains_root_in_conversation_session() -> None:
|
||||
builder = CanonicalTraceBuilder(lambda _info: [])
|
||||
info = GenerateNameTraceInfo(
|
||||
tenant_id="tenant-1",
|
||||
conversation_id="conversation-1",
|
||||
inputs="title prompt",
|
||||
outputs="title",
|
||||
metadata={},
|
||||
)
|
||||
|
||||
trace = builder.build(info)
|
||||
|
||||
assert trace is not None
|
||||
assert trace.session_id == "conversation-1"
|
||||
assert trace.required_parent_context_id is None
|
||||
assert trace.spans[0].parent_id is None
|
||||
|
||||
|
||||
def test_standalone_trace_reuses_persisted_operation_id() -> None:
|
||||
builder = CanonicalTraceBuilder(lambda _info: [])
|
||||
info = GenerateNameTraceInfo(
|
||||
tenant_id="tenant-1",
|
||||
conversation_id="conversation-1",
|
||||
inputs="title prompt",
|
||||
outputs="title",
|
||||
operation_id="00000000-0000-0000-0000-000000000099",
|
||||
metadata={},
|
||||
)
|
||||
|
||||
first = builder.build(info)
|
||||
second = builder.build(info)
|
||||
|
||||
assert first is not None
|
||||
assert second is not None
|
||||
assert first.root_span_id == "00000000-0000-0000-0000-000000000099"
|
||||
assert second.root_span_id == first.root_span_id
|
||||
|
||||
|
||||
def test_standalone_trace_accepts_legacy_payload_without_operation_id() -> None:
|
||||
info = GenerateNameTraceInfo.model_validate(
|
||||
{
|
||||
"tenant_id": "tenant-1",
|
||||
"conversation_id": "conversation-1",
|
||||
"inputs": "title prompt",
|
||||
"outputs": "title",
|
||||
"metadata": {},
|
||||
}
|
||||
)
|
||||
|
||||
trace = CanonicalTraceBuilder(lambda _info: []).build(info)
|
||||
|
||||
assert trace is not None
|
||||
assert UUID(trace.root_span_id)
|
||||
@@ -0,0 +1,258 @@
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import cast
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from dify_trace_arize_phoenix.config import PhoenixConfig
|
||||
from dify_trace_arize_phoenix.unified_trace import UnifiedPhoenixAdapter
|
||||
from openinference.semconv.trace import SpanAttributes
|
||||
from opentelemetry.sdk.trace.export import SpanExportResult
|
||||
from opentelemetry.trace import StatusCode
|
||||
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
|
||||
|
||||
from core.ops.exceptions import InvalidTraceParentContextError, RetryableTraceDispatchError
|
||||
from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace
|
||||
from core.ops.unified_trace.parent_context import ParentResolution, ProviderParentContext, destination_scope
|
||||
|
||||
VALID_TRACEPARENT = "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01"
|
||||
PhoenixAdapterFixture = tuple[UnifiedPhoenixAdapter, MagicMock, list[MagicMock]]
|
||||
|
||||
|
||||
def _inject_traceparent(carrier: dict[str, str], context: object) -> None:
|
||||
assert context is not None
|
||||
carrier["traceparent"] = VALID_TRACEPARENT
|
||||
|
||||
|
||||
def span(**overrides: object) -> CanonicalSpan:
|
||||
values: dict[str, object] = {
|
||||
"id": "root",
|
||||
"parent_id": None,
|
||||
"name": "root",
|
||||
"kind": CanonicalSpanKind.CHAIN,
|
||||
"start_time": datetime(2025, 1, 1),
|
||||
"end_time": datetime(2025, 1, 1, 0, 0, 1),
|
||||
"inputs": {"input": "value"},
|
||||
"outputs": {"output": "value"},
|
||||
"status": CanonicalSpanStatus.OK,
|
||||
"metadata": {},
|
||||
}
|
||||
values.update(overrides)
|
||||
return CanonicalSpan.model_validate(values)
|
||||
|
||||
|
||||
def trace(*spans: CanonicalSpan, session_id: str = "session-1") -> CanonicalTrace:
|
||||
values = spans or (span(),)
|
||||
return CanonicalTrace(
|
||||
trace_id="trace-1",
|
||||
session_id=session_id,
|
||||
root_span_id=values[0].id,
|
||||
spans=values,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def adapter(monkeypatch: pytest.MonkeyPatch) -> PhoenixAdapterFixture:
|
||||
tracer = MagicMock()
|
||||
otel_spans = [MagicMock() for _ in range(8)]
|
||||
tracer.start_span.side_effect = otel_spans
|
||||
exporter = MagicMock()
|
||||
exporter.export.return_value = SpanExportResult.SUCCESS
|
||||
monkeypatch.setattr(
|
||||
"dify_trace_arize_phoenix.unified_trace.setup_unified_tracer",
|
||||
lambda _config: (tracer, exporter),
|
||||
)
|
||||
value = UnifiedPhoenixAdapter(
|
||||
PhoenixConfig(api_key="secret", project="project-a", endpoint="https://phoenix.example")
|
||||
)
|
||||
value._propagator = MagicMock()
|
||||
return value, tracer, otel_spans
|
||||
|
||||
|
||||
def test_emit_creates_parent_before_child_and_maps_session(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, tracer, _ = adapter
|
||||
root = span()
|
||||
child = span(id="child", parent_id="root", name="llm", kind=CanonicalSpanKind.LLM)
|
||||
|
||||
subject.emit(trace(root, child, session_id="customer-session"), None, MagicMock())
|
||||
|
||||
assert [call.kwargs["name"] for call in tracer.start_span.call_args_list] == ["root", "llm"]
|
||||
root_attributes = tracer.start_span.call_args_list[0].kwargs["attributes"]
|
||||
assert root_attributes[SpanAttributes.SESSION_ID] == "customer-session"
|
||||
assert root_attributes[SpanAttributes.OPENINFERENCE_SPAN_KIND] == "CHAIN"
|
||||
assert tracer.start_span.call_args_list[1].kwargs["context"] is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", list(CanonicalSpanKind))
|
||||
def test_emit_maps_every_canonical_kind(adapter: PhoenixAdapterFixture, kind: CanonicalSpanKind) -> None:
|
||||
expected = {
|
||||
CanonicalSpanKind.CHAIN: "CHAIN",
|
||||
CanonicalSpanKind.LLM: "LLM",
|
||||
CanonicalSpanKind.RETRIEVER: "RETRIEVER",
|
||||
CanonicalSpanKind.TOOL: "TOOL",
|
||||
CanonicalSpanKind.AGENT: "AGENT",
|
||||
}[kind]
|
||||
subject, tracer, _ = adapter
|
||||
|
||||
subject.emit(trace(span(kind=kind)), None, MagicMock())
|
||||
|
||||
attributes = tracer.start_span.call_args.kwargs["attributes"]
|
||||
assert attributes[SpanAttributes.OPENINFERENCE_SPAN_KIND] == expected
|
||||
metadata = json.loads(attributes[SpanAttributes.METADATA])
|
||||
assert metadata["dify.span.kind"] == kind.value
|
||||
|
||||
|
||||
def test_emit_preserves_logical_links_and_overrides_reserved_metadata(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, tracer, _ = adapter
|
||||
linked_span = span(
|
||||
metadata={"dify.span.kind": "forged", "dify.span.links": ["forged"]},
|
||||
links=("message-a",),
|
||||
)
|
||||
|
||||
subject.emit(trace(linked_span), None, MagicMock())
|
||||
|
||||
attributes = tracer.start_span.call_args.kwargs["attributes"]
|
||||
metadata = json.loads(attributes[SpanAttributes.METADATA])
|
||||
assert metadata["dify.span.kind"] == "chain"
|
||||
assert metadata["dify.span.links"] == ["message-a"]
|
||||
|
||||
|
||||
def test_emit_restores_w3c_parent_context(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, _, _ = adapter
|
||||
subject._propagator = MagicMock(wraps=TraceContextTextMapPropagator())
|
||||
provider_context = ProviderParentContext(
|
||||
provider="phoenix",
|
||||
scope=subject.scope,
|
||||
trace_id="outer-trace",
|
||||
parent_id="outer-tool",
|
||||
provider_context={"traceparent": VALID_TRACEPARENT},
|
||||
)
|
||||
|
||||
subject.emit(trace(), ParentResolution.restored(provider_context), MagicMock())
|
||||
|
||||
subject._propagator.extract.assert_called_once_with(carrier={"traceparent": VALID_TRACEPARENT})
|
||||
|
||||
|
||||
def test_emit_rejects_restored_context_without_traceparent(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, _, _ = adapter
|
||||
provider_context = ProviderParentContext(
|
||||
provider="phoenix",
|
||||
scope=subject.scope,
|
||||
trace_id="outer-trace",
|
||||
parent_id="outer-tool",
|
||||
provider_context={},
|
||||
)
|
||||
|
||||
with pytest.raises(InvalidTraceParentContextError):
|
||||
subject.emit(trace(), ParentResolution.restored(provider_context), MagicMock())
|
||||
|
||||
|
||||
def test_emit_rejects_malformed_traceparent(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, _, _ = adapter
|
||||
subject._propagator = TraceContextTextMapPropagator()
|
||||
provider_context = ProviderParentContext(
|
||||
provider="phoenix",
|
||||
scope=subject.scope,
|
||||
trace_id="outer-trace",
|
||||
parent_id="outer-tool",
|
||||
provider_context={"traceparent": "malformed"},
|
||||
)
|
||||
|
||||
with pytest.raises(InvalidTraceParentContextError):
|
||||
subject.emit(trace(), ParentResolution.restored(provider_context), MagicMock())
|
||||
|
||||
|
||||
def test_emit_publishes_tool_context_after_span_export(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, tracer, otel_spans = adapter
|
||||
cast(MagicMock, subject._propagator.inject).side_effect = _inject_traceparent
|
||||
events: list[str] = []
|
||||
otel_spans[0].end.side_effect = lambda **_kwargs: events.append("end")
|
||||
publish = MagicMock(side_effect=lambda *_args: events.append("publish"))
|
||||
tool = span(id="tool-exec", kind=CanonicalSpanKind.TOOL, can_parent_workflow=True)
|
||||
|
||||
subject.emit(trace(tool), None, publish)
|
||||
|
||||
assert tracer.start_span.called
|
||||
node_execution_id, context = publish.call_args.args
|
||||
assert node_execution_id == "tool-exec"
|
||||
assert context.provider == "phoenix"
|
||||
assert context.provider_context == {"traceparent": VALID_TRACEPARENT}
|
||||
assert events == ["end", "publish"]
|
||||
|
||||
|
||||
def test_emit_does_not_publish_parent_context_when_export_fails(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, _, _ = adapter
|
||||
cast(MagicMock, subject._exporter.export).return_value = SpanExportResult.FAILURE
|
||||
publish = MagicMock()
|
||||
tool = span(id="tool-exec", kind=CanonicalSpanKind.TOOL, can_parent_workflow=True)
|
||||
|
||||
with pytest.raises(RetryableTraceDispatchError, match="Phoenix span export failed"):
|
||||
subject.emit(trace(tool), None, publish)
|
||||
|
||||
publish.assert_not_called()
|
||||
|
||||
|
||||
def test_emit_maps_exporter_exception_to_retryable_failure(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, _, _ = adapter
|
||||
cast(MagicMock, subject._exporter.export).side_effect = ConnectionError("network unavailable")
|
||||
publish = MagicMock()
|
||||
|
||||
with pytest.raises(RetryableTraceDispatchError, match="Phoenix span export failed"):
|
||||
subject.emit(trace(), None, publish)
|
||||
|
||||
publish.assert_not_called()
|
||||
|
||||
|
||||
def test_emit_publishes_message_context(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, _, _ = adapter
|
||||
cast(MagicMock, subject._propagator.inject).side_effect = _inject_traceparent
|
||||
publish = MagicMock()
|
||||
message = span(id="message-1", name="message", publishes_parent_context=True)
|
||||
|
||||
subject.emit(trace(message), None, publish)
|
||||
|
||||
parent_id, context = publish.call_args.args
|
||||
assert parent_id == "message-1"
|
||||
assert context.provider_context == {"traceparent": VALID_TRACEPARENT}
|
||||
|
||||
|
||||
def test_emit_records_error_status(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, _, otel_spans = adapter
|
||||
failed = span(status=CanonicalSpanStatus.ERROR, error="boom")
|
||||
|
||||
subject.emit(trace(failed), None, MagicMock())
|
||||
|
||||
status = otel_spans[0].set_status.call_args.args[0]
|
||||
assert status.status_code is StatusCode.ERROR
|
||||
otel_spans[0].record_exception.assert_called_once()
|
||||
|
||||
|
||||
def test_retry_metadata_is_serialized_for_phoenix(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, tracer, _ = adapter
|
||||
retry_metadata = {
|
||||
"retry_count": 1,
|
||||
"retry_attempts": [
|
||||
{
|
||||
"retry_index": 1,
|
||||
"error": "HTTP 500",
|
||||
"elapsed_time": 1.2,
|
||||
"created_at": 1_700_000_000,
|
||||
"finished_at": 1_700_000_001,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
subject.emit(trace(span(metadata=retry_metadata)), None, MagicMock())
|
||||
|
||||
attributes = tracer.start_span.call_args.kwargs["attributes"]
|
||||
assert json.loads(attributes[SpanAttributes.METADATA]) == {
|
||||
**retry_metadata,
|
||||
"dify.span.kind": "chain",
|
||||
}
|
||||
|
||||
|
||||
def test_scope_does_not_include_api_key(adapter: PhoenixAdapterFixture) -> None:
|
||||
subject, _, _ = adapter
|
||||
|
||||
assert subject.scope == destination_scope("phoenix", "https://phoenix.example", "project-a")
|
||||
assert "secret" not in subject.scope
|
||||
@@ -0,0 +1,328 @@
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from dify_trace_langsmith.config import LangSmithConfig
|
||||
from dify_trace_langsmith.unified_trace import UnifiedLangSmithAdapter
|
||||
from langsmith.utils import (
|
||||
LangSmithAPIError,
|
||||
LangSmithAuthError,
|
||||
LangSmithConnectionError,
|
||||
LangSmithError,
|
||||
LangSmithRateLimitError,
|
||||
LangSmithRequestTimeout,
|
||||
LangSmithUserError,
|
||||
)
|
||||
|
||||
from core.ops.exceptions import InvalidTraceParentContextError, RetryableTraceDispatchError
|
||||
from core.ops.unified_trace.entities import CanonicalSpan, CanonicalSpanKind, CanonicalSpanStatus, CanonicalTrace
|
||||
from core.ops.unified_trace.parent_context import ParentResolution, ProviderParentContext, destination_scope
|
||||
|
||||
ROOT_ID = "00000000-0000-0000-0000-000000000001"
|
||||
CHILD_ID = "00000000-0000-0000-0000-000000000002"
|
||||
TOOL_ID = "00000000-0000-0000-0000-000000000003"
|
||||
LangSmithAdapterFixture = tuple[UnifiedLangSmithAdapter, MagicMock]
|
||||
|
||||
|
||||
def span(**overrides: object) -> CanonicalSpan:
|
||||
values: dict[str, object] = {
|
||||
"id": ROOT_ID,
|
||||
"parent_id": None,
|
||||
"name": "root",
|
||||
"kind": CanonicalSpanKind.CHAIN,
|
||||
"start_time": datetime(2025, 1, 1),
|
||||
"end_time": datetime(2025, 1, 1, 0, 0, 1),
|
||||
"inputs": {"input": "value"},
|
||||
"outputs": {"output": "value"},
|
||||
"status": CanonicalSpanStatus.OK,
|
||||
"metadata": {"external_trace_id": "customer-trace"},
|
||||
}
|
||||
values.update(overrides)
|
||||
return CanonicalSpan.model_validate(values)
|
||||
|
||||
|
||||
def trace(*spans: CanonicalSpan, session_id: str = "session-1") -> CanonicalTrace:
|
||||
values = spans or (span(),)
|
||||
return CanonicalTrace(
|
||||
trace_id="customer-trace",
|
||||
session_id=session_id,
|
||||
root_span_id=values[0].id,
|
||||
spans=values,
|
||||
)
|
||||
|
||||
|
||||
def test_client_disables_async_batching_before_parent_coordination(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
client_class = MagicMock()
|
||||
monkeypatch.setattr("dify_trace_langsmith.unified_trace.Client", client_class)
|
||||
config = LangSmithConfig(api_key="secret", project="project-a", endpoint="https://smith.example")
|
||||
|
||||
UnifiedLangSmithAdapter(config)
|
||||
|
||||
client_class.assert_called_once_with(
|
||||
api_key="secret",
|
||||
api_url="https://smith.example",
|
||||
auto_batch_tracing=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def adapter(monkeypatch: pytest.MonkeyPatch) -> LangSmithAdapterFixture:
|
||||
client = MagicMock()
|
||||
monkeypatch.setattr("dify_trace_langsmith.unified_trace.Client", lambda **_kwargs: client)
|
||||
subject = UnifiedLangSmithAdapter(
|
||||
LangSmithConfig(api_key="secret", project="project-a", endpoint="https://smith.example")
|
||||
)
|
||||
return subject, client
|
||||
|
||||
|
||||
def test_root_trace_id_equals_root_run_id_and_sets_thread_session(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
|
||||
subject.emit(trace(session_id="customer-session"), None, MagicMock())
|
||||
|
||||
root = client.create_run.call_args.kwargs
|
||||
assert root["id"] == ROOT_ID
|
||||
assert root["trace_id"] == ROOT_ID
|
||||
assert root["extra"]["metadata"]["session_id"] == "customer-session"
|
||||
assert root["extra"]["metadata"]["external_trace_id"] == "customer-trace"
|
||||
|
||||
|
||||
def test_message_span_uses_explicit_langsmith_human_message_schema(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
message = span(
|
||||
name="message",
|
||||
inputs="hi",
|
||||
metadata={"trace_entity_type": "message"},
|
||||
)
|
||||
|
||||
subject.emit(trace(message), None, MagicMock())
|
||||
|
||||
assert client.create_run.call_args.kwargs["inputs"] == {"messages": [{"role": "user", "content": "hi"}]}
|
||||
|
||||
|
||||
def test_mapping_inputs_remain_unchanged(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
raw_inputs: dict[str, object] = {"sys.app_id": "app-1", "sys.files": []}
|
||||
message = span(
|
||||
name="message",
|
||||
inputs=raw_inputs,
|
||||
metadata={"trace_entity_type": "message"},
|
||||
)
|
||||
|
||||
subject.emit(trace(message), None, MagicMock())
|
||||
|
||||
assert client.create_run.call_args.kwargs["inputs"] == raw_inputs
|
||||
|
||||
|
||||
def test_empty_session_is_not_written_to_root_metadata(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
|
||||
subject.emit(trace(session_id=""), None, MagicMock())
|
||||
|
||||
metadata = client.create_run.call_args.kwargs["extra"]["metadata"]
|
||||
assert "session_id" not in metadata
|
||||
|
||||
|
||||
def test_child_uses_actual_parent_run_and_dotted_order(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
root = span()
|
||||
child = span(id=CHILD_ID, parent_id=ROOT_ID, name="llm", kind=CanonicalSpanKind.LLM)
|
||||
|
||||
subject.emit(trace(root, child), None, MagicMock())
|
||||
|
||||
root_run = client.create_run.call_args_list[0].kwargs
|
||||
child_run = client.create_run.call_args_list[1].kwargs
|
||||
assert child_run["parent_run_id"] == ROOT_ID
|
||||
assert child_run["trace_id"] == ROOT_ID
|
||||
assert child_run["dotted_order"].startswith(f"{root_run['dotted_order']}.")
|
||||
assert child_run["run_type"] == "llm"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", list(CanonicalSpanKind))
|
||||
def test_emit_maps_every_canonical_kind(adapter: LangSmithAdapterFixture, kind: CanonicalSpanKind) -> None:
|
||||
expected = {
|
||||
CanonicalSpanKind.CHAIN: "chain",
|
||||
CanonicalSpanKind.LLM: "llm",
|
||||
CanonicalSpanKind.RETRIEVER: "retriever",
|
||||
CanonicalSpanKind.TOOL: "tool",
|
||||
CanonicalSpanKind.AGENT: "chain",
|
||||
}[kind]
|
||||
subject, client = adapter
|
||||
|
||||
subject.emit(trace(span(kind=kind)), None, MagicMock())
|
||||
|
||||
run = client.create_run.call_args.kwargs
|
||||
assert run["run_type"] == expected
|
||||
assert run["extra"]["metadata"]["dify.span.kind"] == kind.value
|
||||
|
||||
|
||||
def test_emit_preserves_logical_links_and_overrides_reserved_metadata(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
linked_span = span(
|
||||
metadata={"dify.span.kind": "forged", "dify.span.links": ["forged"]},
|
||||
links=("message-a",),
|
||||
)
|
||||
|
||||
subject.emit(trace(linked_span), None, MagicMock())
|
||||
|
||||
metadata = client.create_run.call_args.kwargs["extra"]["metadata"]
|
||||
assert metadata["dify.span.kind"] == "chain"
|
||||
assert metadata["dify.span.links"] == ["message-a"]
|
||||
|
||||
|
||||
def test_synthetic_ids_are_mapped_consistently(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
root = span()
|
||||
wrapper = span(id="iteration:container:0", parent_id=ROOT_ID, name="iteration[0]")
|
||||
child = span(id=CHILD_ID, parent_id=wrapper.id)
|
||||
|
||||
subject.emit(trace(root, wrapper, child), None, MagicMock())
|
||||
|
||||
wrapper_run = client.create_run.call_args_list[1].kwargs
|
||||
child_run = client.create_run.call_args_list[2].kwargs
|
||||
assert wrapper_run["id"] != wrapper.id
|
||||
assert child_run["parent_run_id"] == wrapper_run["id"]
|
||||
|
||||
|
||||
def test_nested_workflow_restores_parent_trace_and_order(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
parent = ProviderParentContext(
|
||||
provider="langsmith",
|
||||
scope=subject.scope,
|
||||
trace_id="00000000-0000-0000-0000-000000000010",
|
||||
parent_id="00000000-0000-0000-0000-000000000011",
|
||||
provider_context={"dotted_order": "parent.order"},
|
||||
)
|
||||
|
||||
subject.emit(trace(), ParentResolution.restored(parent), MagicMock())
|
||||
|
||||
root = client.create_run.call_args.kwargs
|
||||
assert root["trace_id"] == parent.trace_id
|
||||
assert root["parent_run_id"] == parent.parent_id
|
||||
assert root["dotted_order"].startswith("parent.order.")
|
||||
|
||||
|
||||
def test_nested_workflow_rejects_parent_without_dotted_order(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, _ = adapter
|
||||
parent = ProviderParentContext(
|
||||
provider="langsmith",
|
||||
scope=subject.scope,
|
||||
trace_id="00000000-0000-0000-0000-000000000010",
|
||||
parent_id="00000000-0000-0000-0000-000000000011",
|
||||
provider_context={},
|
||||
)
|
||||
|
||||
with pytest.raises(InvalidTraceParentContextError):
|
||||
subject.emit(trace(), ParentResolution.restored(parent), MagicMock())
|
||||
|
||||
|
||||
def test_tool_context_is_published_only_after_create_run_succeeds(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
publish = MagicMock()
|
||||
tool = span(id=TOOL_ID, kind=CanonicalSpanKind.TOOL, can_parent_workflow=True)
|
||||
|
||||
subject.emit(trace(tool), None, publish)
|
||||
|
||||
assert client.create_run.called
|
||||
node_execution_id, context = publish.call_args.args
|
||||
assert node_execution_id == TOOL_ID
|
||||
assert context.parent_id == TOOL_ID
|
||||
assert context.provider_context["dotted_order"]
|
||||
|
||||
client.reset_mock()
|
||||
publish.reset_mock()
|
||||
client.create_run.side_effect = RuntimeError("rejected")
|
||||
with pytest.raises(RuntimeError, match="rejected"):
|
||||
subject.emit(trace(tool), None, publish)
|
||||
publish.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error",
|
||||
[
|
||||
LangSmithConnectionError("down"),
|
||||
LangSmithRequestTimeout("slow"),
|
||||
LangSmithRateLimitError("limited"),
|
||||
LangSmithAPIError("server"),
|
||||
],
|
||||
)
|
||||
def test_emit_maps_recoverable_sdk_errors_to_retryable_failure(
|
||||
adapter: LangSmithAdapterFixture, error: LangSmithError
|
||||
) -> None:
|
||||
subject, client = adapter
|
||||
client.create_run.side_effect = error
|
||||
|
||||
with pytest.raises(RetryableTraceDispatchError) as exc_info:
|
||||
subject.emit(trace(), None, MagicMock())
|
||||
assert exc_info.value.__cause__ is error
|
||||
|
||||
|
||||
@pytest.mark.parametrize("error", [LangSmithAuthError("unauthorized"), LangSmithUserError("invalid")])
|
||||
def test_emit_maps_known_terminal_sdk_errors_to_runtime_failure(
|
||||
adapter: LangSmithAdapterFixture, error: LangSmithError
|
||||
) -> None:
|
||||
subject, client = adapter
|
||||
client.create_run.side_effect = error
|
||||
|
||||
with pytest.raises(RuntimeError, match="LangSmith run export rejected") as exc_info:
|
||||
subject.emit(trace(), None, MagicMock())
|
||||
assert not isinstance(exc_info.value, RetryableTraceDispatchError)
|
||||
assert exc_info.value.__cause__ is error
|
||||
|
||||
|
||||
def test_emit_maps_unknown_sdk_error_to_terminal_runtime_failure(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
error = LangSmithError("unknown")
|
||||
client.create_run.side_effect = error
|
||||
|
||||
with pytest.raises(RuntimeError, match="LangSmith run export failed") as exc_info:
|
||||
subject.emit(trace(), None, MagicMock())
|
||||
assert not isinstance(exc_info.value, RetryableTraceDispatchError)
|
||||
assert exc_info.value.__cause__ is error
|
||||
|
||||
|
||||
def test_retry_metadata_is_forwarded_to_langsmith(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
retry_metadata = {
|
||||
"retry_count": 1,
|
||||
"retry_attempts": [
|
||||
{
|
||||
"retry_index": 1,
|
||||
"error": "HTTP 500",
|
||||
"elapsed_time": 1.2,
|
||||
"created_at": 1_700_000_000,
|
||||
"finished_at": 1_700_000_001,
|
||||
}
|
||||
],
|
||||
}
|
||||
node = span(metadata=retry_metadata)
|
||||
|
||||
subject.emit(trace(node), None, MagicMock())
|
||||
|
||||
assert client.create_run.call_args.kwargs["extra"]["metadata"] == {
|
||||
**retry_metadata,
|
||||
"session_id": "session-1",
|
||||
"external_trace_id": "customer-trace",
|
||||
"dify.span.kind": "chain",
|
||||
}
|
||||
|
||||
|
||||
def test_message_context_is_published_after_create_run(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, client = adapter
|
||||
publish = MagicMock()
|
||||
message = span(name="message", publishes_parent_context=True)
|
||||
|
||||
subject.emit(trace(message), None, publish)
|
||||
|
||||
assert client.create_run.called
|
||||
parent_id, context = publish.call_args.args
|
||||
assert parent_id == ROOT_ID
|
||||
assert context.parent_id == ROOT_ID
|
||||
assert context.provider_context["dotted_order"]
|
||||
|
||||
|
||||
def test_scope_does_not_include_api_key(adapter: LangSmithAdapterFixture) -> None:
|
||||
subject, _ = adapter
|
||||
|
||||
assert subject.scope == destination_scope("langsmith", "https://smith.example", "project-a")
|
||||
assert "secret" not in subject.scope
|
||||
@@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch
|
||||
import pytest
|
||||
from celery.exceptions import Retry
|
||||
|
||||
from configs import dify_config
|
||||
from core.ops.entities.config_entity import OPS_TRACE_FAILED_KEY
|
||||
from core.ops.exceptions import RetryableTraceDispatchError
|
||||
from tasks.ops_trace_task import process_trace_tasks
|
||||
@@ -215,8 +216,11 @@ def test_process_trace_tasks_skips_enterprise_trace_when_retry_payload_was_alrea
|
||||
mock_incr.assert_not_called()
|
||||
|
||||
|
||||
def test_process_trace_tasks_default_retry_window_covers_parent_span_context_ttl():
|
||||
assert process_trace_tasks.max_retries * process_trace_tasks.default_retry_delay >= 300
|
||||
def test_process_trace_tasks_default_retry_window_covers_workflow_and_export_grace_period():
|
||||
assert (
|
||||
process_trace_tasks.max_retries * process_trace_tasks.default_retry_delay
|
||||
>= dify_config.WORKFLOW_MAX_EXECUTION_TIME + 300
|
||||
)
|
||||
|
||||
|
||||
def test_process_trace_tasks_deletes_payload_on_success():
|
||||
|
||||
Reference in new issue
Block a user