mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(tracing): tag graph runs with the effective model (#5902)
This commit is contained in:
parent
fa2afa6a39
commit
e2c5027b0c
@ -982,7 +982,7 @@ If you are using a self-hosted Langfuse instance, set `LANGFUSE_BASE_URL` to you
|
||||
- `session_id` = LangGraph `thread_id` — groups every trace of the same conversation
|
||||
- `user_id` = effective user from `get_effective_user_id()` (falls back to `default` in no-auth mode)
|
||||
- `trace_name` = assistant id (defaults to `lead-agent`)
|
||||
- `tags` = `[env:<DEER_FLOW_ENV>, model:<model_name>]` (omitted when not set)
|
||||
- `tags` = `[env:<DEER_FLOW_ENV>, model:<model_name>]` (built-in Gateway lead-agent and embedded client runs tag the selected model after default or fallback resolution; custom graph factories may only provide the requested model; the environment tag is omitted when unset)
|
||||
- `metadata.deerflow_trace_id` = DeerFlow request correlation id, matching `X-Trace-Id` when request trace correlation is enabled
|
||||
|
||||
These are injected into `RunnableConfig.metadata` at the graph invocation root for both the gateway path (`runtime/runs/worker.py::run_agent`) and the embedded path (`client.py::DeerFlowClient.stream`), so any LangChain-compatible callback can read them. Set `DEER_FLOW_ENV` (or `ENVIRONMENT`) to tag traces by deployment environment.
|
||||
|
||||
@ -9,7 +9,9 @@ same resulting prompt. Empty overlays preserve bytes.
|
||||
- `make_lead_agent(config: RunnableConfig)` is the published `langgraph.json`
|
||||
entry point; preserve its signature and bare-graph return type.
|
||||
- Gateway calls `assemble_lead_agent(config, *, app_config=None)` for a
|
||||
`LeadAgentAssembly(graph, descriptor)`; `make_lead_agent` returns `.graph`.
|
||||
`LeadAgentAssembly(graph, descriptor, effective_model)`; `make_lead_agent` returns `.graph`.
|
||||
`effective_model` is available even when no observer requests a descriptor,
|
||||
so graph-root tracing can tag the selected default or fallback model.
|
||||
`assembly_descriptor.py::build_assembly_descriptor()` records the model after
|
||||
runtime overrides, rendered prompt hash, authorized tools and middleware order.
|
||||
Consumers must unwrap via `runtime/runs/worker.py::_agent_graph` to also accept
|
||||
|
||||
@ -100,6 +100,7 @@ class LeadAgentAssembly:
|
||||
|
||||
graph: Any
|
||||
descriptor: Any
|
||||
effective_model: str | None = None
|
||||
|
||||
|
||||
def unwrap_agent_graph(agent_result: Any) -> Any:
|
||||
@ -897,7 +898,7 @@ def _complete_assembly(
|
||||
|
||||
resolved_extensions = get_agent_build_extensions()
|
||||
if not resolved_extensions.has_agent_assembly_observers:
|
||||
return LeadAgentAssembly(graph=graph, descriptor=None)
|
||||
return LeadAgentAssembly(graph=graph, descriptor=None, effective_model=effective_model)
|
||||
|
||||
from deerflow.agents.assembly_descriptor import build_assembly_descriptor
|
||||
from deerflow.extensions.notify import notify_agent_assembled
|
||||
@ -924,7 +925,7 @@ def _complete_assembly(
|
||||
effective_policies=resolved_policies,
|
||||
)
|
||||
notify_agent_assembled(descriptor, resolved_extensions)
|
||||
return LeadAgentAssembly(graph=graph, descriptor=descriptor)
|
||||
return LeadAgentAssembly(graph=graph, descriptor=descriptor, effective_model=effective_model)
|
||||
|
||||
|
||||
def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> LeadAgentAssembly:
|
||||
|
||||
@ -247,6 +247,7 @@ class DeerFlowClient:
|
||||
# Lazy agent — created on first call, recreated when config changes.
|
||||
self._agent = None
|
||||
self._agent_config_key: tuple | None = None
|
||||
self._effective_model_name: str | None = None
|
||||
self._loaded_agent_config_key: tuple[str, str] | None = None
|
||||
self._loaded_agent_config = None
|
||||
|
||||
@ -259,6 +260,7 @@ class DeerFlowClient:
|
||||
"""
|
||||
self._agent = None
|
||||
self._agent_config_key = None
|
||||
self._effective_model_name = None
|
||||
self._loaded_agent_config_key = None
|
||||
self._loaded_agent_config = None
|
||||
|
||||
@ -487,6 +489,7 @@ class DeerFlowClient:
|
||||
|
||||
self._agent = create_agent(**kwargs)
|
||||
self._agent_config_key = key
|
||||
self._effective_model_name = model_name
|
||||
logger.info("Agent created: agent_name=%s, model=%s, thinking=%s", self._agent_name, model_name, thinking_enabled)
|
||||
|
||||
@staticmethod
|
||||
@ -966,7 +969,6 @@ class DeerFlowClient:
|
||||
if key in kwargs:
|
||||
context[key] = kwargs[key]
|
||||
|
||||
configurable = config.get("configurable") or {}
|
||||
deerflow_trace_id = ensure_trace_id()
|
||||
effective_user_id = context.get("user_id") or get_effective_user_id()
|
||||
# Materialize the storage owner in runtime context in every auth mode.
|
||||
@ -974,18 +976,19 @@ class DeerFlowClient:
|
||||
# survives worker/isolated-loop boundaries and matches the identity
|
||||
# used by prompt assembly and the agent cache.
|
||||
context["user_id"] = effective_user_id
|
||||
self._ensure_agent(config, context=context)
|
||||
configurable = config.get("configurable") or {}
|
||||
effective_model_name = getattr(self, "_effective_model_name", None)
|
||||
inject_langfuse_metadata(
|
||||
config,
|
||||
thread_id=thread_id,
|
||||
user_id=effective_user_id,
|
||||
assistant_id=self._agent_name or "lead-agent",
|
||||
model_name=configurable.get("model_name") or self._model_name,
|
||||
model_name=effective_model_name or configurable.get("model_name") or self._model_name,
|
||||
environment=self._environment or os.environ.get("DEER_FLOW_ENV") or os.environ.get("ENVIRONMENT"),
|
||||
deerflow_trace_id=deerflow_trace_id,
|
||||
)
|
||||
|
||||
self._ensure_agent(config, context=context)
|
||||
|
||||
state: dict[str, Any] = {"messages": [HumanMessage(content=message, additional_kwargs={"run_id": run_id})]}
|
||||
context[DEERFLOW_TRACE_METADATA_KEY] = deerflow_trace_id
|
||||
if self._agent_name:
|
||||
|
||||
@ -709,6 +709,17 @@ def _agent_graph(agent_result: Any) -> Any:
|
||||
return unwrap_agent_graph(agent_result)
|
||||
|
||||
|
||||
def _assembled_model_name(agent_result: Any) -> str | None:
|
||||
"""Return the selected model only for the trusted lead assembly result."""
|
||||
try:
|
||||
from deerflow.agents.lead_agent.agent import LeadAgentAssembly
|
||||
except Exception:
|
||||
return None
|
||||
if isinstance(agent_result, LeadAgentAssembly):
|
||||
return agent_result.effective_model
|
||||
return None
|
||||
|
||||
|
||||
class _SubagentEventBuffer:
|
||||
"""Buffer subagent ``task_*`` step events and flush them in one locked batch (#3779).
|
||||
|
||||
@ -1145,20 +1156,6 @@ async def run_agent(
|
||||
if journal is not None:
|
||||
config.setdefault("callbacks", []).append(journal)
|
||||
|
||||
# Inject Langfuse trace-attribute metadata so the langchain CallbackHandler
|
||||
# can lift session_id / user_id / trace_name / tags onto the root trace.
|
||||
# Shared helper with ``DeerFlowClient.stream`` so both entry points stay
|
||||
# in sync; caller-provided metadata wins via setdefault inside the helper.
|
||||
inject_langfuse_metadata(
|
||||
config,
|
||||
thread_id=thread_id,
|
||||
user_id=resolve_runtime_user_id(runtime),
|
||||
assistant_id=record.assistant_id,
|
||||
model_name=record.model_name,
|
||||
environment=os.environ.get("DEER_FLOW_ENV") or os.environ.get("ENVIRONMENT"),
|
||||
deerflow_trace_id=deerflow_trace_id,
|
||||
)
|
||||
|
||||
# Resolve after runtime context installation so context/configurable reflect
|
||||
# the agent name that this run will actually execute.
|
||||
config.setdefault("run_name", resolve_root_run_name(config, record.assistant_id))
|
||||
@ -1186,7 +1183,22 @@ async def run_agent(
|
||||
# get_available_tools(), which may block on MCP cache
|
||||
# initialization — it must not stall the calling event loop
|
||||
# (issue #5172).
|
||||
agent = _agent_graph(await run_assembly(agent_factory, **agent_factory_kwargs))
|
||||
agent_result = await run_assembly(agent_factory, **agent_factory_kwargs)
|
||||
agent = _agent_graph(agent_result)
|
||||
|
||||
# Assembly resolves request, agent, and authorization fallbacks. Trace the
|
||||
# model that will run, rather than the model name originally requested.
|
||||
effective_model = _assembled_model_name(agent_result) or record.model_name
|
||||
for trace_config in (config, initial_runnable_config):
|
||||
inject_langfuse_metadata(
|
||||
trace_config,
|
||||
thread_id=thread_id,
|
||||
user_id=resolve_runtime_user_id(runtime),
|
||||
assistant_id=record.assistant_id,
|
||||
model_name=effective_model,
|
||||
environment=os.environ.get("DEER_FLOW_ENV") or os.environ.get("ENVIRONMENT"),
|
||||
deerflow_trace_id=deerflow_trace_id,
|
||||
)
|
||||
|
||||
accessor = CheckpointStateAccessor.bind(
|
||||
agent,
|
||||
@ -1773,6 +1785,7 @@ async def run_agent(
|
||||
# Drop graph and per-run payload references before the terminal
|
||||
# worker task itself becomes collectable.
|
||||
agent = None
|
||||
agent_result = None
|
||||
accessor = None
|
||||
runtime = None
|
||||
runtime_ctx = None
|
||||
|
||||
@ -5,14 +5,14 @@ LangSmith and Langfuse are both supported. The wiring lives in two layers:
|
||||
- `factory.py::build_tracing_callbacks()` — returns the LangChain `CallbackHandler` list for the providers currently enabled via env vars (`LANGSMITH_TRACING`, `LANGFUSE_TRACING`, etc.). The handlers are attached at the **graph invocation root** for in-graph runs (`make_lead_agent` and `DeerFlowClient.stream` both append them to `config["callbacks"]` before invoking the graph) so a single run produces one trace with all node / LLM / tool calls as child spans. Standalone callers — anything that invokes a model outside such a graph (e.g. `MemoryUpdater`) — keep `create_chat_model`'s default `attach_tracing=True`, which falls back to model-level callback attachment.
|
||||
- `metadata.py::build_langfuse_trace_metadata()` — builds the Langfuse-reserved trace attributes for `RunnableConfig.metadata`. The Langfuse v4 `langchain.CallbackHandler` lifts these onto the root trace (see its `_parse_langfuse_trace_attributes`), but only when it sees `on_chain_start(parent_run_id=None)` — which is why the callbacks have to live at the graph root, not the model.
|
||||
|
||||
**Trace-attribute injection points**: both `runtime/runs/worker.py::run_agent` (gateway path) and `client.py::DeerFlowClient.stream` (embedded path) merge the metadata into `config["metadata"]` right before constructing the graph. `subagents/executor.py::_aexecute` does the same for every subagent run so subagent traces group under the parent thread's session card (carrying the parent `thread_id` → `langfuse_session_id`, the user_id captured at `task_tool` → `langfuse_user_id`, and a `subagent:<normalized-name>` trace name). Caller-supplied keys win via `setdefault`, so an external `session_id` override is preserved. Field mapping:
|
||||
**Trace-attribute injection points**: both `runtime/runs/worker.py::run_agent` (gateway path) and `client.py::DeerFlowClient.stream` (embedded path) merge the metadata into `config["metadata"]` after model selection, before invoking the graph. `subagents/executor.py::_aexecute` does the same for every subagent run so subagent traces group under the parent thread's session card (carrying the parent `thread_id` → `langfuse_session_id`, the user_id captured at `task_tool` → `langfuse_user_id`, and a `subagent:<normalized-name>` trace name). Caller-supplied keys win via `setdefault`, so an external `session_id` override is preserved. Field mapping:
|
||||
|
||||
| Langfuse field | Source |
|
||||
|-----------------------|----------------------------------------------|
|
||||
| `langfuse_session_id` | LangGraph `thread_id` |
|
||||
| `langfuse_user_id` | `get_effective_user_id()` (`default` in no-auth); for subagents, captured from `runtime.context` at `task_tool` time via `resolve_runtime_user_id()` |
|
||||
| `langfuse_trace_name` | `RunRecord.assistant_id` / client `agent_name` (defaults to `lead-agent`); for subagents, `subagent:<name>` (lowercased, `_` → `-`) |
|
||||
| `langfuse_tags` | `env:<DEER_FLOW_ENV>` + `model:<model_name>` |
|
||||
| `langfuse_tags` | `env:<DEER_FLOW_ENV>` + `model:<effective model name>` for built-in Gateway lead-agent and embedded client runs (including default and fallback selections); custom bare graph factories fall back to the requested model |
|
||||
| `deerflow_trace_id` | Current entry trace id from `deerflow.trace_context`; always written, and always equal to the `X-Trace-Id` the same Gateway request returned. Not gated by config — see `packages/harness/deerflow/AGENTS.md` |
|
||||
|
||||
Returns `{}` when Langfuse is not in the enabled providers — LangSmith-only deployments are unaffected. Set `DEER_FLOW_ENV` (or `ENVIRONMENT`) to tag traces by deployment environment. Tests live in `tests/test_tracing_factory.py`, `tests/test_tracing_metadata.py`, `tests/test_worker_langfuse_metadata.py`, `tests/test_client_langfuse_metadata.py`, and `tests/test_subagent_executor.py::TestSubagentTracingWiring`.
|
||||
|
||||
@ -241,6 +241,7 @@ class TestLeadAgentAssembly:
|
||||
assert assembly.graph is not None
|
||||
assert assembly.graph["context_schema"] is dict
|
||||
assert assembly.descriptor.effective_model
|
||||
assert assembly.effective_model == assembly.descriptor.effective_model
|
||||
assert assembly.descriptor.fingerprint
|
||||
|
||||
def test_descriptor_hashes_the_same_scoped_prompt_passed_to_the_graph(self, monkeypatch):
|
||||
@ -338,6 +339,7 @@ class TestLeadAgentAssembly:
|
||||
self._isolate_from_the_ambient_config(monkeypatch)
|
||||
assembly = assemble_lead_agent({"configurable": {"thread_id": "t-3"}})
|
||||
assert assembly.descriptor is None
|
||||
assert assembly.effective_model
|
||||
|
||||
|
||||
class TestFactoryConsumersUnwrapTheGraph:
|
||||
|
||||
@ -117,6 +117,28 @@ def test_stream_injects_langfuse_metadata_when_enabled(monkeypatch):
|
||||
assert sentinel in callbacks
|
||||
|
||||
|
||||
def test_stream_tags_the_effective_default_model(monkeypatch):
|
||||
monkeypatch.setenv("LANGFUSE_TRACING", "true")
|
||||
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-test")
|
||||
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-test")
|
||||
from deerflow.config.tracing_config import reset_tracing_config
|
||||
|
||||
reset_tracing_config()
|
||||
fake_agent = _FakeAgent()
|
||||
|
||||
def ensure_agent(self, config, **kwargs):
|
||||
self._effective_model_name = "stub-model"
|
||||
self._agent = fake_agent
|
||||
|
||||
monkeypatch.setattr(DeerFlowClient, "_ensure_agent", ensure_agent)
|
||||
client = _make_client(monkeypatch)
|
||||
client._model_name = None
|
||||
|
||||
list(client.stream("hi", thread_id="thread-default-model"))
|
||||
|
||||
assert "model:stub-model" in fake_agent.captured_config["metadata"]["langfuse_tags"]
|
||||
|
||||
|
||||
def test_stream_is_inert_when_langfuse_disabled(monkeypatch):
|
||||
monkeypatch.setattr("deerflow.client.build_tracing_callbacks", lambda: [])
|
||||
|
||||
|
||||
@ -619,6 +619,7 @@ def test_client_ensure_agent_enforces_model_use_when_authorized(monkeypatch):
|
||||
|
||||
# Denied ``gpt-4`` was swapped for the authorized fallback ``claude-3``.
|
||||
assert captured_name["name"] == "claude-3"
|
||||
assert client._effective_model_name == "claude-3"
|
||||
|
||||
|
||||
def test_client_ensure_agent_resolves_none_default_before_authorization(monkeypatch):
|
||||
@ -647,6 +648,7 @@ def test_client_ensure_agent_resolves_none_default_before_authorization(monkeypa
|
||||
client._ensure_agent(config)
|
||||
|
||||
assert captured_name["name"] == "claude-3"
|
||||
assert client._effective_model_name == "claude-3"
|
||||
|
||||
|
||||
def test_client_ensure_agent_noop_when_authorization_disabled(monkeypatch):
|
||||
|
||||
@ -144,6 +144,68 @@ async def test_run_agent_injects_langfuse_metadata(monkeypatch):
|
||||
assert "model:gpt-4o" in tags
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("requested_model", [None, "missing-model", "actual-model"])
|
||||
async def test_run_agent_tags_the_effective_model(monkeypatch, requested_model):
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
|
||||
from deerflow.agents.lead_agent import agent as lead_agent_module
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.config.model_config import ModelConfig
|
||||
from deerflow.config.sandbox_config import SandboxConfig
|
||||
|
||||
monkeypatch.setenv("LANGFUSE_TRACING", "true")
|
||||
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-test")
|
||||
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-test")
|
||||
from deerflow.config.tracing_config import reset_tracing_config
|
||||
|
||||
reset_tracing_config()
|
||||
fake_agent = _FakeAgent()
|
||||
app_config = AppConfig(
|
||||
models=[ModelConfig(name="actual-model", model="actual-model", use="langchain_openai:ChatOpenAI")],
|
||||
sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider"),
|
||||
)
|
||||
selected_models = []
|
||||
callback = BaseCallbackHandler()
|
||||
|
||||
def create_model(**kwargs):
|
||||
selected_models.append(kwargs["name"])
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(lead_agent_module, "create_chat_model", create_model)
|
||||
monkeypatch.setattr(lead_agent_module, "create_agent", lambda **kwargs: fake_agent)
|
||||
monkeypatch.setattr(lead_agent_module, "build_tracing_callbacks", lambda: [callback])
|
||||
|
||||
def agent_factory(config):
|
||||
return lead_agent_module.assemble_lead_agent(config, app_config=app_config)
|
||||
|
||||
record = RunRecord(
|
||||
run_id="run-effective-model",
|
||||
thread_id="thread-effective-model",
|
||||
assistant_id="lead-agent",
|
||||
status=RunStatus.pending,
|
||||
on_disconnect=DisconnectMode.cancel,
|
||||
model_name=requested_model,
|
||||
)
|
||||
record.abort_event = asyncio.Event()
|
||||
|
||||
await run_agent(
|
||||
_FakeBridge(),
|
||||
_FakeRunManager(),
|
||||
record,
|
||||
ctx=RunContext(checkpointer=None, app_config=app_config),
|
||||
agent_factory=agent_factory,
|
||||
graph_input={"messages": []},
|
||||
config={"configurable": {"thread_id": record.thread_id, "model_name": requested_model}},
|
||||
)
|
||||
|
||||
assert selected_models == ["actual-model"]
|
||||
assert callback in fake_agent.captured_config["callbacks"]
|
||||
tags = fake_agent.captured_config["metadata"].get("langfuse_tags", [])
|
||||
assert "model:actual-model" in tags
|
||||
assert "model:missing-model" not in tags
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_agent_uses_context_user_id_over_contextvar(monkeypatch):
|
||||
"""A run carrying ``context.user_id`` traces to that user, not the contextvar.
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user