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.
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/.gitignore +7 -0
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/PKG-INFO +1 -1
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/pyproject.toml +1 -1
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/src/by_framework_langgraph/_utils.py +0 -11
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/src/by_framework_langgraph/adapter.py +245 -7
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/src/by_framework_langgraph/tools.py +37 -1
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/tests/test_adapter.py +145 -5
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/tests/test_tools.py +51 -1
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/README.md +0 -0
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/src/by_framework_langgraph/__init__.py +0 -0
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/src/by_framework_langgraph/worker.py +0 -0
- {by_framework_langgraph-0.0.2 → by_framework_langgraph-0.0.3.dev1}/tests/test_utils.py +0 -0
|
@@ -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
|
|
@@ -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
|
|
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
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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."""
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|