by-framework-langgraph 0.0.2__tar.gz → 0.0.3.dev1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -3,6 +3,7 @@ __pycache__/
3
3
  *.py[cod]
4
4
  *$py.class
5
5
  *.so
6
+ .ruff_cache/
6
7
  .Python
7
8
  build/
8
9
  develop-eggs/
@@ -66,3 +67,9 @@ Thumbs.db
66
67
  *gateway-sdk.log*
67
68
  .claude/settings.local.json
68
69
  .claude/settings.json
70
+
71
+ core
72
+
73
+ .coderfleet-uploads/
74
+
75
+ .claude/skills
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: by-framework-langgraph
3
- Version: 0.0.2
3
+ Version: 0.0.3.dev1
4
4
  Summary: LangGraph integration for by-framework
5
5
  Requires-Python: >=3.12
6
6
  Requires-Dist: by-framework>=0.2.1
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "by-framework-langgraph"
3
- version = "0.0.2"
3
+ version = "0.0.3.dev1"
4
4
  description = "LangGraph integration for by-framework"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.12"
@@ -2,7 +2,6 @@
2
2
 
3
3
  from __future__ import annotations
4
4
 
5
- import hashlib
6
5
  from typing import Any
7
6
 
8
7
  from by_framework.core.protocol.commands import ResumeCommand
@@ -52,13 +51,3 @@ def extract_resume_data(command: ResumeCommand) -> str:
52
51
  return extract_content_text(command.content)
53
52
 
54
53
  return ""
55
-
56
-
57
- def str_to_uint128(s: str) -> int:
58
- """Convert a string to a 128-bit integer (for OTEL TraceId)."""
59
- return int(hashlib.md5(s.encode()).hexdigest(), 16)
60
-
61
-
62
- def str_to_uint64(s: str) -> int:
63
- """Convert a string to a 64-bit integer (for OTEL SpanId)."""
64
- return int(hashlib.md5(s.encode()).hexdigest()[:16], 16)
@@ -16,15 +16,12 @@ from by_framework.common.logger import logger
16
16
  from by_framework.core.protocol.agent_state import AgentState
17
17
  from by_framework.core.protocol.commands import ResumeCommand
18
18
  from by_framework.core.protocol.events import StreamChunkEvent
19
+ from by_framework.trace.span_recorder import str_to_uint64, str_to_uint128
20
+ from langchain_core.callbacks import BaseCallbackHandler
19
21
  from langchain_core.messages import HumanMessage
20
22
  from langgraph.types import Command
21
23
 
22
- from ._utils import (
23
- extract_content_text,
24
- extract_resume_data,
25
- str_to_uint64,
26
- str_to_uint128,
27
- )
24
+ from ._utils import extract_content_text, extract_resume_data
28
25
 
29
26
  if TYPE_CHECKING:
30
27
  from by_framework.core.protocol.commands import GatewayCommand
@@ -35,6 +32,172 @@ if TYPE_CHECKING:
35
32
  LANGFUSE_OBSERVATION_ATTR = "_langfuse_observation"
36
33
 
37
34
 
35
+ class _TokenAccumulatingCallbackHandler(BaseCallbackHandler):
36
+ """LangChain callback handler that accumulates LLM token usage into AgentContext.
37
+
38
+ Extracts token usage from whichever location the provider populates:
39
+ - ``llm_output["token_usage"]`` (OpenAI-style via LangChain)
40
+ - ``llm_output["usage"]`` (Anthropic-style / raw provider mapping)
41
+ - ``generation.message.usage_metadata`` (LangChain >= 0.2 standard)
42
+ - ``generation.generation_info`` (some community integrations)
43
+
44
+ ``run_id`` deduplication prevents double-counting when both ``on_llm_end``
45
+ and ``on_chat_model_end`` fire for the same call (LangChain >= 0.2).
46
+ """
47
+
48
+ def __init__(self, context: Any) -> None:
49
+ super().__init__()
50
+ self._context = context
51
+ self._seen_run_ids: set = set()
52
+
53
+ # ------------------------------------------------------------------
54
+ # LangChain callback entry point
55
+ # ------------------------------------------------------------------
56
+
57
+ def on_llm_end(self, response: Any, *, run_id: Any = None, **_kwargs: Any) -> None:
58
+ # Guard: only mark run as seen when we actually extracted tokens so that
59
+ # the on_chat_model_end event path in _stream_invoke can still fire as a
60
+ # fallback when the callback found nothing (e.g. stream_options not set).
61
+ if run_id is not None and run_id in self._seen_run_ids:
62
+ return
63
+ self._handle_llm_result(response, run_id=run_id)
64
+
65
+ # ------------------------------------------------------------------
66
+ # Internal extraction logic
67
+ # ------------------------------------------------------------------
68
+
69
+ def _handle_llm_result(self, response: Any, *, run_id: Any = None) -> None:
70
+ context = self._context
71
+ if context is None:
72
+ return
73
+
74
+ prompt, completion = self._extract_tokens(response)
75
+
76
+ if not (prompt or completion):
77
+ # Log at WARNING (always visible) to help diagnose providers whose
78
+ # token format is not yet handled, or where stream_options is missing.
79
+ gens = getattr(response, "generations", []) or []
80
+ first = (gens[0] or [None])[0] if gens else None
81
+ msg_type = type(getattr(first, "message", None)).__name__
82
+ logger.warning(
83
+ "[TokenAccumulator] on_llm_end fired but extracted 0 tokens. "
84
+ "For OpenAI-compatible streaming APIs add "
85
+ "stream_options={'include_usage': True} to your ChatModel. "
86
+ "llm_output=%r first_gen_message_type=%s",
87
+ getattr(response, "llm_output", None),
88
+ msg_type,
89
+ )
90
+ # Do NOT mark run_id as seen — let the on_chat_model_end event path
91
+ # in _stream_invoke attempt extraction from the merged message.
92
+ return
93
+
94
+ if run_id is not None:
95
+ self._seen_run_ids.add(run_id)
96
+ try:
97
+ context.record_token_usage(
98
+ prompt_tokens=prompt,
99
+ completion_tokens=completion,
100
+ )
101
+ except Exception: # pylint: disable=broad-exception-caught
102
+ pass
103
+
104
+ @staticmethod
105
+ def _extract_tokens(response: Any) -> tuple[int, int]:
106
+ """Return (prompt_tokens, completion_tokens) from an LLMResult.
107
+
108
+ Checks every known location across providers:
109
+ 1. llm_output["token_usage"] — OpenAI via LangChain
110
+ 2. llm_output["usage"] — Anthropic / raw provider mapping
111
+ 3. message.usage_metadata — LangChain >= 0.2 standard
112
+ 4. message.response_metadata — some community integrations
113
+ 5. generation.generation_info — older / custom integrations
114
+ """
115
+ prompt, completion = 0, 0
116
+
117
+ llm_output = getattr(response, "llm_output", None) or {}
118
+ if isinstance(llm_output, dict):
119
+ for key in ("token_usage", "usage"):
120
+ usage = llm_output.get(key) or {}
121
+ if usage:
122
+ prompt = int(
123
+ usage.get("prompt_tokens") or usage.get("input_tokens") or 0
124
+ )
125
+ completion = int(
126
+ usage.get("completion_tokens")
127
+ or usage.get("output_tokens")
128
+ or 0
129
+ )
130
+ break
131
+
132
+ if prompt or completion:
133
+ return prompt, completion
134
+
135
+ # Iterate all generations
136
+ for gen_list in getattr(response, "generations", []) or []:
137
+ for gen in (gen_list if isinstance(gen_list, list) else [gen_list]):
138
+ msg = getattr(gen, "message", None)
139
+
140
+ # LangChain >= 0.2: message.usage_metadata
141
+ meta = getattr(msg, "usage_metadata", None)
142
+ if meta:
143
+ prompt += int(
144
+ meta.get("input_tokens") or meta.get("prompt_tokens") or 0
145
+ )
146
+ completion += int(
147
+ meta.get("output_tokens") or meta.get("completion_tokens") or 0
148
+ )
149
+ continue
150
+
151
+ # response_metadata (e.g. MiniMax, Qwen, some Chinese providers)
152
+ resp_meta = getattr(msg, "response_metadata", None) or {}
153
+ if isinstance(resp_meta, dict):
154
+ for key in ("token_usage", "usage"):
155
+ usage = resp_meta.get(key) or {}
156
+ if usage:
157
+ prompt += int(
158
+ usage.get("prompt_tokens")
159
+ or usage.get("input_tokens")
160
+ or 0
161
+ )
162
+ completion += int(
163
+ usage.get("completion_tokens")
164
+ or usage.get("output_tokens")
165
+ or 0
166
+ )
167
+ break
168
+ # Flat keys at root of response_metadata
169
+ if not (prompt or completion):
170
+ prompt += int(
171
+ resp_meta.get("prompt_tokens")
172
+ or resp_meta.get("input_tokens")
173
+ or 0
174
+ )
175
+ completion += int(
176
+ resp_meta.get("completion_tokens")
177
+ or resp_meta.get("output_tokens")
178
+ or 0
179
+ )
180
+ if prompt or completion:
181
+ continue
182
+
183
+ # generation_info fallback
184
+ info = getattr(gen, "generation_info", None) or {}
185
+ for key in ("token_usage", "usage"):
186
+ sub = info.get(key) or {}
187
+ if sub:
188
+ prompt += int(
189
+ sub.get("prompt_tokens") or sub.get("input_tokens") or 0
190
+ )
191
+ completion += int(
192
+ sub.get("completion_tokens")
193
+ or sub.get("output_tokens")
194
+ or 0
195
+ )
196
+ break
197
+
198
+ return prompt, completion
199
+
200
+
38
201
  @dataclass(frozen=True)
39
202
  class _AdapterTracingConfig:
40
203
  """Tracing-related adapter config kept separate from core graph handles."""
@@ -171,6 +334,47 @@ class LangGraphAdapter:
171
334
  await self._context.emit_chunk(
172
335
  chunk.content, content_type="text"
173
336
  )
337
+ elif kind == "on_chat_model_end":
338
+ # Fallback: capture token usage from the event's merged output
339
+ # message when the on_llm_end callback found nothing (e.g. the
340
+ # provider requires stream_options but it was not set).
341
+ # _TokenAccumulatingCallbackHandler marks run_id as seen only
342
+ # after a successful extraction, so this path fires only when
343
+ # the callback got 0 tokens.
344
+ run_id = event.get("run_id")
345
+ token_handler = next(
346
+ (
347
+ cb
348
+ for cb in scoped_callbacks
349
+ if isinstance(cb, _TokenAccumulatingCallbackHandler)
350
+ ),
351
+ None,
352
+ )
353
+ if token_handler is not None and run_id not in (
354
+ token_handler._seen_run_ids # pylint: disable=protected-access
355
+ ):
356
+ output = event["data"].get("output")
357
+ meta = getattr(output, "usage_metadata", None)
358
+ if meta:
359
+ prompt = int(
360
+ meta.get("input_tokens")
361
+ or meta.get("prompt_tokens")
362
+ or 0
363
+ )
364
+ completion = int(
365
+ meta.get("output_tokens")
366
+ or meta.get("completion_tokens")
367
+ or 0
368
+ )
369
+ if prompt or completion:
370
+ token_handler._seen_run_ids.add(run_id) # pylint: disable=protected-access
371
+ try:
372
+ self._context.record_token_usage(
373
+ prompt_tokens=prompt,
374
+ completion_tokens=completion,
375
+ )
376
+ except Exception: # pylint: disable=broad-exception-caught
377
+ pass
174
378
  elif kind == "on_tool_start":
175
379
  tool_name = event["name"]
176
380
  tool_input = event["data"].get("input")
@@ -188,7 +392,7 @@ class LangGraphAdapter:
188
392
  ),
189
393
  )
190
394
 
191
- # Use a stable logical ID from metadata if available, fallback to run_id
395
+ # Use stable metadata when available, otherwise fall back.
192
396
  stable_id = (
193
397
  event.get("metadata", {}).get("tool_call_id")
194
398
  or event.get("metadata", {}).get("checkpoint_ns")
@@ -383,6 +587,29 @@ class LangGraphAdapter:
383
587
  @contextmanager
384
588
  def _langfuse_callback_manager(self, callbacks: list[Any]) -> Iterator[None]:
385
589
  """Prepare Langfuse callback and observation for LangChain."""
590
+ # Always inject token accumulator — works regardless of Langfuse config.
591
+ callbacks.append(_TokenAccumulatingCallbackHandler(self._context))
592
+
593
+ # Prefer AgentContext's callback factory so trace and parent ids align.
594
+ # Filter out auto-generated MagicMock attributes when tests use a mock
595
+ # context — real callback objects always come from a non-test module.
596
+ langfuse_callback_value = getattr(self._context, "langfuse_callback", None)
597
+ is_real_callback = langfuse_callback_value is not None and type(
598
+ langfuse_callback_value
599
+ ).__module__ not in ("unittest.mock",)
600
+
601
+ if is_real_callback:
602
+ handler = (
603
+ langfuse_callback_value()
604
+ if callable(langfuse_callback_value)
605
+ else langfuse_callback_value
606
+ )
607
+ if handler is not None:
608
+ callbacks.append(handler)
609
+ yield
610
+ return
611
+
612
+ # Fallback to local import if context method is missing
386
613
  # pylint: disable=import-outside-toplevel
387
614
  try:
388
615
  langfuse_config = import_module(
@@ -414,6 +641,17 @@ class LangGraphAdapter:
414
641
  },
415
642
  metadata=self._default_metadata(),
416
643
  ):
644
+ # Prevent the generated OTel span from being promoted to a trace root.
645
+ # The native LangfusePlugin sets the same attribute on its own path
646
+ # (via _SdkLangfuseTracer); this covers the LangGraph fallback path.
647
+ try:
648
+ from opentelemetry import trace
649
+
650
+ current_span = trace.get_current_span()
651
+ if current_span and hasattr(current_span, "set_attribute"):
652
+ current_span.set_attribute("langfuse.internal.as_root", False)
653
+ except Exception: # pylint: disable=broad-exception-caught
654
+ pass
417
655
  yield
418
656
 
419
657
  @staticmethod
@@ -7,7 +7,7 @@ mechanism.
7
7
 
8
8
  from __future__ import annotations
9
9
 
10
- from typing import TYPE_CHECKING, Annotated
10
+ from typing import TYPE_CHECKING, Annotated, Any
11
11
 
12
12
  from langchain_core.tools import BaseTool, InjectedToolCallId, tool
13
13
  from langgraph.types import interrupt
@@ -16,6 +16,31 @@ if TYPE_CHECKING:
16
16
  from by_framework.worker.context import AgentContext
17
17
 
18
18
 
19
+ def _langfuse_observation_id_from_callbacks(callbacks: Any) -> str:
20
+ """Return the Langfuse observation id for the active LangChain tool run."""
21
+ run_id = getattr(callbacks, "run_id", None) or getattr(
22
+ callbacks,
23
+ "parent_run_id",
24
+ None,
25
+ )
26
+ if not run_id:
27
+ return ""
28
+
29
+ handlers = [
30
+ *list(getattr(callbacks, "handlers", []) or []),
31
+ *list(getattr(callbacks, "inheritable_handlers", []) or []),
32
+ ]
33
+ for handler in handlers:
34
+ runs = getattr(handler, "_runs", None)
35
+ if not isinstance(runs, dict):
36
+ continue
37
+ observation = runs.get(run_id)
38
+ observation_id = getattr(observation, "id", None)
39
+ if observation_id:
40
+ return str(observation_id)
41
+ return ""
42
+
43
+
19
44
  def make_remote_agent_tool(
20
45
  context: AgentContext,
21
46
  tool_name: str,
@@ -50,6 +75,7 @@ def make_remote_agent_tool(
50
75
  async def remote_agent_tool(
51
76
  topic: str,
52
77
  tool_call_id: Annotated[str, InjectedToolCallId],
78
+ callbacks: Any = None,
53
79
  ) -> str:
54
80
  # Idempotency guard: checkpoint restore replays tool execution,
55
81
  # but we must not re-dispatch the command.
@@ -57,9 +83,19 @@ def make_remote_agent_tool(
57
83
  is_dispatched = await context.redis.exists(redis_key)
58
84
 
59
85
  if not is_dispatched:
86
+ metadata = {}
87
+ langfuse_parent_observation_id = _langfuse_observation_id_from_callbacks(
88
+ callbacks
89
+ )
90
+ if langfuse_parent_observation_id:
91
+ metadata["langfuse_parent_observation_id"] = (
92
+ langfuse_parent_observation_id
93
+ )
94
+
60
95
  await context.call_agent(
61
96
  target_agent_type=target_agent_type,
62
97
  content=topic,
98
+ metadata=metadata,
63
99
  )
64
100
  await context.redis.set(redis_key, "1", ex=idempotency_ttl)
65
101
 
@@ -3,13 +3,16 @@
3
3
  import sys
4
4
  from contextlib import contextmanager
5
5
  from types import SimpleNamespace
6
- from unittest.mock import AsyncMock, MagicMock
6
+ from unittest.mock import AsyncMock, MagicMock, patch
7
7
 
8
8
  import pytest
9
9
  from by_framework.core.protocol.commands import AskAgentCommand, ResumeCommand
10
10
  from by_framework.core.protocol.message_header import MessageHeader
11
11
 
12
- from by_framework_langgraph.adapter import LangGraphAdapter
12
+ from by_framework_langgraph.adapter import (
13
+ LangGraphAdapter,
14
+ _TokenAccumulatingCallbackHandler,
15
+ )
13
16
  from by_framework_langgraph.worker import LangGraphWorker
14
17
 
15
18
 
@@ -251,7 +254,12 @@ class TestAdapterRun:
251
254
  assert config["metadata"]["langfuse_user_id"] == "user-1"
252
255
  assert config["metadata"]["by_framework_message_id"] == "msg-ctx"
253
256
  assert config["metadata"]["langgraph_thread_id"] == "test-session"
254
- assert config["callbacks"] == callback_instances
257
+ # Callbacks list includes the Langfuse handler(s) plus the token accumulator.
258
+ assert all(cb in config["callbacks"] for cb in callback_instances)
259
+ assert any(
260
+ type(cb).__name__ == "_TokenAccumulatingCallbackHandler"
261
+ for cb in config["callbacks"]
262
+ )
255
263
  assert len(observation_calls) == 1
256
264
  assert observation_calls[0]["trace_context"] == {
257
265
  "trace_id": "trace-ctx",
@@ -259,6 +267,54 @@ class TestAdapterRun:
259
267
  }
260
268
  assert observation_calls[0]["name"] == "planner:langgraph"
261
269
 
270
+ @pytest.mark.asyncio
271
+ async def test_uses_context_langfuse_callback_property(self):
272
+ """Verify AgentContext.langfuse_callback property is used directly."""
273
+ handler = object()
274
+
275
+ # pylint: disable=too-few-public-methods,missing-class-docstring,missing-function-docstring
276
+ class ContextWithCallbackProperty:
277
+ session_id = "test-session"
278
+ trace_id = "trace-ctx"
279
+ message_id = "msg-ctx"
280
+ parent_message_id = ""
281
+ current_agent_id = "planner"
282
+ current_command = SimpleNamespace(
283
+ header=MessageHeader(
284
+ message_id="msg-ctx",
285
+ session_id="test-session",
286
+ trace_id="trace-ctx",
287
+ target_agent_type="planner",
288
+ )
289
+ )
290
+
291
+ def __init__(self):
292
+ self.emit_chunk = AsyncMock()
293
+
294
+ @property
295
+ def langfuse_callback(self):
296
+ return handler
297
+
298
+ ctx = ContextWithCallbackProperty()
299
+ graph = MagicMock()
300
+ graph.ainvoke = AsyncMock(
301
+ return_value={"messages": [MagicMock(content="hello")]}
302
+ )
303
+ snapshot = MagicMock()
304
+ snapshot.next = ()
305
+ graph.get_state.return_value = snapshot
306
+
307
+ adapter = LangGraphAdapter(graph, ctx, stream=False)
308
+ result = await adapter.run(AskAgentCommand(header=_make_header(), content="hi"))
309
+
310
+ assert result == "hello"
311
+ _, kwargs = graph.ainvoke.call_args
312
+ callbacks = kwargs["config"]["callbacks"]
313
+ assert handler in callbacks
314
+ assert any(
315
+ type(cb).__name__ == "_TokenAccumulatingCallbackHandler" for cb in callbacks
316
+ )
317
+
262
318
  @pytest.mark.asyncio
263
319
  async def test_skips_langfuse_tracing_silently_when_not_configured(
264
320
  self, monkeypatch, caplog
@@ -296,7 +352,13 @@ class TestAdapterRun:
296
352
  assert result == "hello"
297
353
  _, kwargs = graph.ainvoke.call_args
298
354
  config = kwargs["config"]
299
- assert "callbacks" not in config
355
+ # Token accumulator is always present; only Langfuse callbacks should be absent.
356
+ non_token_callbacks = [
357
+ cb
358
+ for cb in config.get("callbacks", [])
359
+ if type(cb).__name__ != "_TokenAccumulatingCallbackHandler"
360
+ ]
361
+ assert non_token_callbacks == []
300
362
  assert "Langfuse" not in caplog.text
301
363
 
302
364
  @pytest.mark.asyncio
@@ -337,7 +399,12 @@ class TestAdapterRun:
337
399
  assert result == "hello"
338
400
  _, kwargs = graph.ainvoke.call_args
339
401
  config = kwargs["config"]
340
- assert "callbacks" not in config
402
+ non_token_callbacks = [
403
+ cb
404
+ for cb in config.get("callbacks", [])
405
+ if type(cb).__name__ != "_TokenAccumulatingCallbackHandler"
406
+ ]
407
+ assert non_token_callbacks == []
341
408
  assert "Langfuse" not in caplog.text
342
409
 
343
410
 
@@ -406,3 +473,76 @@ class TestLangGraphWorkerHooks: # pylint: disable=too-few-public-methods
406
473
  assert captured["run_name"] == "custom-run"
407
474
  assert captured["metadata"] == {"team": "alpha"}
408
475
  assert captured["callbacks"] == ["cb-1"]
476
+
477
+
478
+ class TestTokenAccumulatingCallbackHandler:
479
+
480
+ def _make_llm_result(self, prompt=10, completion=20, style="openai"):
481
+ """Build a mock LLMResult in either openai or usage_metadata style."""
482
+ result = MagicMock()
483
+ if style == "openai":
484
+ result.llm_output = {
485
+ "token_usage": {
486
+ "prompt_tokens": prompt,
487
+ "completion_tokens": completion,
488
+ }
489
+ }
490
+ result.generations = []
491
+ else:
492
+ result.llm_output = {}
493
+ gen = MagicMock()
494
+ gen.message.usage_metadata = {
495
+ "input_tokens": prompt,
496
+ "output_tokens": completion,
497
+ }
498
+ result.generations = [[gen]]
499
+ return result
500
+
501
+ def test_accumulates_openai_style(self):
502
+ ctx = MagicMock()
503
+ handler = _TokenAccumulatingCallbackHandler(ctx)
504
+ handler.on_llm_end(self._make_llm_result(10, 20, "openai"))
505
+ ctx.record_token_usage.assert_called_once_with(
506
+ prompt_tokens=10, completion_tokens=20
507
+ )
508
+
509
+ def test_accumulates_usage_metadata_style(self):
510
+ ctx = MagicMock()
511
+ handler = _TokenAccumulatingCallbackHandler(ctx)
512
+ handler.on_llm_end(self._make_llm_result(5, 15, "metadata"))
513
+ ctx.record_token_usage.assert_called_once_with(
514
+ prompt_tokens=5, completion_tokens=15
515
+ )
516
+
517
+ def test_no_call_on_zero_tokens(self):
518
+ ctx = MagicMock()
519
+ handler = _TokenAccumulatingCallbackHandler(ctx)
520
+ result = MagicMock()
521
+ result.llm_output = {}
522
+ result.generations = []
523
+ handler.on_llm_end(result)
524
+ ctx.record_token_usage.assert_not_called()
525
+
526
+ def test_none_context_does_not_raise(self):
527
+ handler = _TokenAccumulatingCallbackHandler(None)
528
+ handler.on_llm_end(self._make_llm_result())
529
+
530
+ def test_callback_injected_in_tracing_scope(self):
531
+ """_langfuse_callback_manager injects _TokenAccumulatingCallbackHandler."""
532
+ ctx = MagicMock()
533
+ ctx.langfuse_callback = None
534
+ graph = MagicMock()
535
+ adapter = LangGraphAdapter(graph=graph, context=ctx)
536
+
537
+ callbacks = []
538
+ with patch.object(
539
+ adapter,
540
+ "_langfuse_callback_manager",
541
+ wraps=adapter._langfuse_callback_manager,
542
+ ):
543
+ with adapter._langfuse_callback_manager(callbacks):
544
+ pass
545
+
546
+ assert any(
547
+ isinstance(cb, _TokenAccumulatingCallbackHandler) for cb in callbacks
548
+ )
@@ -1,8 +1,15 @@
1
1
  """Tests for tools module."""
2
2
 
3
+ from types import SimpleNamespace
3
4
  from unittest.mock import AsyncMock, MagicMock
4
5
 
5
- from by_framework_langgraph.tools import (make_ask_user_tool, make_remote_agent_tool)
6
+ import pytest
7
+
8
+ from by_framework_langgraph.tools import (
9
+ _langfuse_observation_id_from_callbacks,
10
+ make_ask_user_tool,
11
+ make_remote_agent_tool,
12
+ )
6
13
 
7
14
 
8
15
  def _make_mock_context(session_id: str = "test-session"):
@@ -32,6 +39,49 @@ class TestMakeRemoteAgentTool:
32
39
  )
33
40
  assert "Invoke the poet agent" in tool.description
34
41
 
42
+ def test_resolves_langfuse_observation_id_from_callbacks(self):
43
+ """The active tool observation can be read from Langfuse callback state."""
44
+ run_id = object()
45
+ handler = SimpleNamespace(_runs={run_id: SimpleNamespace(id="obs-tool")})
46
+ callbacks = SimpleNamespace(parent_run_id=run_id, handlers=[handler])
47
+
48
+ assert _langfuse_observation_id_from_callbacks(callbacks) == "obs-tool"
49
+
50
+ @pytest.mark.asyncio
51
+ async def test_passes_tool_observation_id_to_call_agent(self, monkeypatch):
52
+ """Remote calls are parented to the current LangGraph tool observation."""
53
+ ctx = _make_mock_context()
54
+ ctx.redis.exists.return_value = False
55
+ tool = make_remote_agent_tool(
56
+ ctx,
57
+ "query_weather",
58
+ "weather-agent",
59
+ "Query weather",
60
+ )
61
+ monkeypatch.setattr(
62
+ "by_framework_langgraph.tools.interrupt",
63
+ lambda _message: "queued",
64
+ )
65
+
66
+ run_id = object()
67
+ handler = SimpleNamespace(_runs={run_id: SimpleNamespace(id="obs-tool")})
68
+ callbacks = SimpleNamespace(parent_run_id=run_id, handlers=[handler])
69
+ run_manager = SimpleNamespace(get_child=lambda: callbacks)
70
+
71
+ result = await tool._arun(
72
+ "Beijing weather",
73
+ tool_call_id="tool-call-1",
74
+ run_manager=run_manager,
75
+ config={},
76
+ )
77
+
78
+ assert result == "queued"
79
+ ctx.call_agent.assert_awaited_once_with(
80
+ target_agent_type="weather-agent",
81
+ content="Beijing weather",
82
+ metadata={"langfuse_parent_observation_id": "obs-tool"},
83
+ )
84
+
35
85
 
36
86
  class TestMakeAskUserTool:
37
87
  """Tests for make_ask_user_tool."""