generated from kgod/ai-review-template
提交
This commit is contained in:
kgod
@@ -0,0 +1,84 @@
|
||||
"""LangSmith integration for Claude Agent SDK.
|
||||
|
||||
This module provides automatic tracing for the Claude Agent SDK by instrumenting
|
||||
`ClaudeSDKClient` and injecting hooks to trace all tool calls.
|
||||
|
||||
Instrumentation is applied **in place** on the original ``ClaudeSDKClient`` class
|
||||
so that callers who imported the class *before* ``configure_claude_agent_sdk()``
|
||||
was called still get traced.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from ._client import instrument_claude_client, instrument_sdk_mcp_tool
|
||||
from ._config import set_tracing_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = ["configure_claude_agent_sdk"]
|
||||
|
||||
|
||||
def configure_claude_agent_sdk(
|
||||
name: Optional[str] = None,
|
||||
project_name: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
tags: Optional[list[str]] = None,
|
||||
) -> bool:
|
||||
"""Enable LangSmith tracing for the Claude Agent SDK by patching entry points.
|
||||
|
||||
This function instruments the Claude Agent SDK to automatically trace:
|
||||
- Chain runs for each conversation stream (via `ClaudeSDKClient`)
|
||||
- Model runs for each assistant turn
|
||||
- All tool calls including built-in tools, external MCP tools, and SDK MCP tools
|
||||
|
||||
Tool tracing is implemented via `PreToolUse` and `PostToolUse` hooks.
|
||||
|
||||
The class is patched **in place**, so references obtained via
|
||||
``from claude_agent_sdk import ClaudeSDKClient`` before this call
|
||||
will still be instrumented.
|
||||
|
||||
Args:
|
||||
name: Name of the root trace.
|
||||
project_name: LangSmith project to trace to.
|
||||
metadata: Metadata to associate with all traces.
|
||||
tags: Tags to associate with all traces.
|
||||
|
||||
Returns:
|
||||
`True` if configuration was successful, `False` otherwise.
|
||||
|
||||
Example:
|
||||
>>> from langsmith.integrations.claude_agent_sdk import (
|
||||
... configure_claude_agent_sdk,
|
||||
... )
|
||||
>>> configure_claude_agent_sdk(
|
||||
... project_name="my-project", tags=["production"]
|
||||
... ) # doctest: +SKIP
|
||||
>>> # Now use claude_agent_sdk as normal - tracing is automatic
|
||||
"""
|
||||
try:
|
||||
import claude_agent_sdk # type: ignore[import-not-found]
|
||||
except ImportError:
|
||||
logger.warning("Claude Agent SDK not installed.")
|
||||
return False
|
||||
|
||||
if not hasattr(claude_agent_sdk, "ClaudeSDKClient"):
|
||||
logger.warning("Claude Agent SDK missing ClaudeSDKClient.")
|
||||
return False
|
||||
|
||||
set_tracing_config(
|
||||
name=name,
|
||||
project_name=project_name,
|
||||
metadata=metadata,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
instrument_claude_client(claude_agent_sdk.ClaudeSDKClient)
|
||||
|
||||
# Patch SdkMcpTool so that tool handlers are lazily wrapped with
|
||||
# run-context propagation, regardless of import order.
|
||||
sdk_mcp_tool_cls = getattr(claude_agent_sdk, "SdkMcpTool", None)
|
||||
if sdk_mcp_tool_cls:
|
||||
instrument_sdk_mcp_tool(sdk_mcp_tool_cls)
|
||||
|
||||
return True
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,709 @@
|
||||
"""Client instrumentation for Claude Agent SDK."""
|
||||
|
||||
import logging
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import AsyncGenerator, AsyncIterable
|
||||
from datetime import datetime, timezone
|
||||
from functools import cache
|
||||
from typing import Any, Optional
|
||||
|
||||
from langsmith._internal import _context
|
||||
from langsmith.run_helpers import get_current_run_tree, trace
|
||||
|
||||
from ._config import get_tracing_config
|
||||
from ._hooks import (
|
||||
SessionState,
|
||||
_current_session,
|
||||
_register_session,
|
||||
_set_session_root,
|
||||
_unregister_session,
|
||||
clear_active_tool_runs,
|
||||
get_subagent_run_by_tool_id,
|
||||
post_tool_use_failure_hook,
|
||||
post_tool_use_hook,
|
||||
pre_tool_use_hook,
|
||||
subagent_start_hook,
|
||||
subagent_stop_hook,
|
||||
)
|
||||
from ._messages import (
|
||||
build_llm_input,
|
||||
flatten_content_blocks,
|
||||
unwrap_message_dicts,
|
||||
)
|
||||
from ._tools import (
|
||||
clear_parent_run_tree,
|
||||
get_parent_run_tree,
|
||||
set_parent_run_tree,
|
||||
)
|
||||
from ._transcripts import LLM_RUN_NAME, reconcile_from_transcripts
|
||||
from ._usage import extract_usage_metadata
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TRACE_CHAIN_NAME = "claude.conversation"
|
||||
|
||||
|
||||
@cache
|
||||
def _get_package_version(package_name: str) -> str | None:
|
||||
try:
|
||||
from importlib.metadata import version
|
||||
|
||||
return version(package_name)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
class TurnLifecycle:
|
||||
"""Track ongoing model runs so consecutive messages are recorded correctly.
|
||||
|
||||
The Claude Agent SDK may deliver a single assistant turn as multiple
|
||||
``AssistantMessage`` events (e.g. one with ``ThinkingBlock``, another
|
||||
with ``TextBlock``/``ToolUseBlock``). Messages that share the same
|
||||
``message_id`` are accumulated into a single LLM run.
|
||||
"""
|
||||
|
||||
def __init__(self, query_start_time: Optional[float] = None):
|
||||
self.current_run: Optional[Any] = None
|
||||
self.current_message_id: Optional[str] = None
|
||||
self.next_start_time: Optional[float] = query_start_time
|
||||
# message_id → RunTree for all LLM runs created this conversation.
|
||||
# Used to retroactively set usage from transcripts.
|
||||
self.llm_runs_by_message_id: dict[str, Any] = {}
|
||||
# Runs that have been end()ed but not yet patch()ed.
|
||||
# Deferred so transcript usage can be set before the single patch().
|
||||
self._pending_patch: list[Any] = []
|
||||
|
||||
def start_llm_run(
|
||||
self,
|
||||
message: Any,
|
||||
prompt: Any,
|
||||
history: list[dict[str, Any]],
|
||||
parent: Optional[Any] = None,
|
||||
) -> Optional[dict[str, Any]]:
|
||||
"""Begin or continue a model run for *message*.
|
||||
|
||||
If *message* has the same ``message_id`` as the current run the
|
||||
output is appended; otherwise a new run is started (ending any
|
||||
previous one first).
|
||||
"""
|
||||
message_id = getattr(message, "message_id", None)
|
||||
start = self.next_start_time or time.time()
|
||||
|
||||
# Same turn – just accumulate the output blocks and update usage.
|
||||
# Return None so the caller does NOT append a duplicate history
|
||||
# entry; the original entry in ``history`` is updated in place.
|
||||
if message_id and message_id == self.current_message_id and self.current_run:
|
||||
content = flatten_content_blocks(getattr(message, "content", None))
|
||||
if content and self.current_run.outputs:
|
||||
prev = self.current_run.outputs.get("content", [])
|
||||
if isinstance(prev, list) and isinstance(content, list):
|
||||
merged = prev + content
|
||||
self.current_run.outputs["content"] = merged
|
||||
# Update the existing history entry in place so
|
||||
# subsequent LLM runs see a single merged message.
|
||||
for entry in reversed(history):
|
||||
if entry.get("role") == "assistant":
|
||||
entry["content"] = merged
|
||||
break
|
||||
elif isinstance(content, list):
|
||||
self.current_run.outputs["content"] = content
|
||||
self._set_usage_from_message(message, self.current_run)
|
||||
return None
|
||||
|
||||
# Different turn – end previous but defer patch() until
|
||||
# transcript usage is available.
|
||||
if self.current_run:
|
||||
self.current_run.end()
|
||||
self._pending_patch.append(self.current_run)
|
||||
|
||||
final_output, run = begin_llm_run_from_assistant_messages(
|
||||
[message], prompt, history, start_time=start, parent=parent
|
||||
)
|
||||
self.current_run = run
|
||||
self.current_message_id = message_id
|
||||
self.next_start_time = None
|
||||
|
||||
if run:
|
||||
if message_id:
|
||||
self.llm_runs_by_message_id[message_id] = run
|
||||
self._set_usage_from_message(message, run)
|
||||
|
||||
return final_output
|
||||
|
||||
@staticmethod
|
||||
def _set_usage_from_message(message: Any, run: Any) -> None:
|
||||
"""Set usage metadata on a run from a live AssistantMessage.
|
||||
|
||||
Always overwrites — later chunks in the same turn have more
|
||||
accurate counts. Transcript-based usage will overwrite again
|
||||
if available.
|
||||
"""
|
||||
raw_usage = getattr(message, "usage", None)
|
||||
if not raw_usage:
|
||||
return
|
||||
usage_meta = extract_usage_metadata(raw_usage)
|
||||
if usage_meta:
|
||||
meta = run.extra.setdefault("metadata", {})
|
||||
meta["usage_metadata"] = usage_meta
|
||||
|
||||
def mark_next_start(self) -> None:
|
||||
"""Mark when the next assistant message will start."""
|
||||
self.next_start_time = time.time()
|
||||
|
||||
def close(self) -> None:
|
||||
"""End any open run and add to pending patch list."""
|
||||
if self.current_run:
|
||||
self.current_run.end()
|
||||
self._pending_patch.append(self.current_run)
|
||||
self.current_run = None
|
||||
|
||||
def flush(self) -> None:
|
||||
"""Patch all deferred LLM runs. Call after usage has been set."""
|
||||
for run in self._pending_patch:
|
||||
try:
|
||||
run.patch()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to patch LLM run: {e}")
|
||||
self._pending_patch.clear()
|
||||
|
||||
|
||||
def begin_llm_run_from_assistant_messages(
|
||||
messages: list[Any],
|
||||
prompt: Any,
|
||||
history: list[dict[str, Any]],
|
||||
start_time: Optional[float] = None,
|
||||
parent: Optional[Any] = None,
|
||||
) -> tuple[Optional[dict[str, Any]], Optional[Any]]:
|
||||
"""Create a traced model run from assistant messages."""
|
||||
if not messages or type(messages[-1]).__name__ != "AssistantMessage":
|
||||
return None, None
|
||||
|
||||
last_msg = messages[-1]
|
||||
model = getattr(last_msg, "model", None)
|
||||
if parent is None:
|
||||
parent = get_parent_run_tree() or get_current_run_tree()
|
||||
if not parent:
|
||||
return None, None
|
||||
|
||||
inputs = build_llm_input(prompt, history)
|
||||
outputs = [
|
||||
{"content": flatten_content_blocks(m.content), "role": "assistant"}
|
||||
for m in messages
|
||||
if hasattr(m, "content")
|
||||
]
|
||||
|
||||
llm_metadata: dict[str, Any] = {"ls_provider": "anthropic"}
|
||||
if model:
|
||||
llm_metadata["ls_model_name"] = model
|
||||
|
||||
llm_run = parent.create_child(
|
||||
name=LLM_RUN_NAME,
|
||||
run_type="llm",
|
||||
inputs={"messages": inputs} if inputs else {},
|
||||
extra={"metadata": llm_metadata},
|
||||
start_time=datetime.fromtimestamp(start_time, tz=timezone.utc)
|
||||
if start_time
|
||||
else None,
|
||||
)
|
||||
|
||||
try:
|
||||
llm_run.post()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to post LLM run: {e}")
|
||||
|
||||
# Set outputs after posting so they are sent with end_time on the patch.
|
||||
llm_run.outputs = outputs[-1] if len(outputs) == 1 else {"content": outputs}
|
||||
|
||||
final_content = (
|
||||
{"content": flatten_content_blocks(last_msg.content), "role": "assistant"}
|
||||
if hasattr(last_msg, "content")
|
||||
else None
|
||||
)
|
||||
return final_content, llm_run
|
||||
|
||||
|
||||
def _bind_hook_to_session(hook: Any, session: Optional[SessionState]) -> Any:
|
||||
"""Return a hook callable that runs with *session* bound, if provided."""
|
||||
if session is None:
|
||||
return hook
|
||||
|
||||
async def _bound(input_data: Any, tool_use_id: Any, context: Any) -> Any:
|
||||
token = _current_session.set(session)
|
||||
try:
|
||||
return await hook(input_data, tool_use_id, context)
|
||||
finally:
|
||||
_current_session.reset(token)
|
||||
|
||||
return _bound
|
||||
|
||||
|
||||
def _inject_tracing_hooks(options: Any, session: Optional[SessionState] = None) -> None:
|
||||
"""Inject LangSmith tracing hooks into ClaudeAgentOptions.
|
||||
|
||||
If *session* is provided, injected hook callables bind that session around
|
||||
each hook invocation. This is important because the Claude SDK may execute
|
||||
hooks in async contexts that do not inherit the ``receive_response``
|
||||
ContextVar; binding at hook injection time keeps each client isolated.
|
||||
"""
|
||||
if not hasattr(options, "hooks"):
|
||||
return
|
||||
|
||||
# Initialize hooks dict if not present
|
||||
if options.hooks is None:
|
||||
options.hooks = {}
|
||||
|
||||
for event in (
|
||||
"PreToolUse",
|
||||
"PostToolUse",
|
||||
"PostToolUseFailure",
|
||||
"SubagentStart",
|
||||
"SubagentStop",
|
||||
):
|
||||
if event not in options.hooks:
|
||||
options.hooks[event] = []
|
||||
|
||||
try:
|
||||
from claude_agent_sdk import HookMatcher # type: ignore[import-not-found]
|
||||
|
||||
langsmith_pre_matcher = HookMatcher(
|
||||
matcher=None, hooks=[_bind_hook_to_session(pre_tool_use_hook, session)]
|
||||
)
|
||||
langsmith_post_matcher = HookMatcher(
|
||||
matcher=None, hooks=[_bind_hook_to_session(post_tool_use_hook, session)]
|
||||
)
|
||||
langsmith_failure_matcher = HookMatcher(
|
||||
matcher=None,
|
||||
hooks=[_bind_hook_to_session(post_tool_use_failure_hook, session)],
|
||||
)
|
||||
langsmith_subagent_start_matcher = HookMatcher(
|
||||
matcher=None, hooks=[_bind_hook_to_session(subagent_start_hook, session)]
|
||||
)
|
||||
langsmith_subagent_stop_matcher = HookMatcher(
|
||||
matcher=None, hooks=[_bind_hook_to_session(subagent_stop_hook, session)]
|
||||
)
|
||||
|
||||
options.hooks["PreToolUse"].insert(0, langsmith_pre_matcher)
|
||||
options.hooks["PostToolUse"].insert(0, langsmith_post_matcher)
|
||||
options.hooks["PostToolUseFailure"].insert(0, langsmith_failure_matcher)
|
||||
options.hooks["SubagentStart"].insert(0, langsmith_subagent_start_matcher)
|
||||
options.hooks["SubagentStop"].insert(0, langsmith_subagent_stop_matcher)
|
||||
|
||||
logger.debug("Injected LangSmith tracing hooks into ClaudeAgentOptions")
|
||||
except ImportError:
|
||||
logger.warning("Failed to import HookMatcher from claude_agent_sdk")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to inject tracing hooks: {e}")
|
||||
|
||||
|
||||
def _wrap_tool_handler(
|
||||
original_handler: Any,
|
||||
session: Optional[SessionState] = None,
|
||||
tool_name: Optional[str] = None,
|
||||
) -> Any:
|
||||
"""Wrap an MCP tool handler to propagate LangSmith run context.
|
||||
|
||||
The Claude SDK runs hooks and tool handlers in different async task
|
||||
contexts, so contextvars set in ``PreToolUse`` are invisible to the
|
||||
handler. This wrapper copies the active tool run into the contextvar before
|
||||
calling the original handler, so ``@traceable`` calls inside the handler
|
||||
nest correctly.
|
||||
"""
|
||||
|
||||
async def _wrapped(args: Any) -> Any:
|
||||
# The most recently added active tool run is the one PreToolUse just
|
||||
# created for this invocation. Prefer an explicitly bound client
|
||||
# session because tool handlers may run in an async context that did
|
||||
# not inherit _current_session.
|
||||
tool_run = _get_last_active_tool_run(session, args=args, tool_name=tool_name)
|
||||
if tool_run:
|
||||
token = _context._PARENT_RUN_TREE_REF.set(weakref.ref(tool_run))
|
||||
session_token = (
|
||||
_current_session.set(session) if session is not None else None
|
||||
)
|
||||
try:
|
||||
return await original_handler(args)
|
||||
finally:
|
||||
if session_token is not None:
|
||||
_current_session.reset(session_token)
|
||||
_context._PARENT_RUN_TREE_REF.reset(token)
|
||||
return await original_handler(args)
|
||||
|
||||
_wrapped._langsmith_wrapped = True # type: ignore[attr-defined]
|
||||
_wrapped._langsmith_original_handler = original_handler # type: ignore[attr-defined]
|
||||
_wrapped._langsmith_session = session # type: ignore[attr-defined]
|
||||
_wrapped._langsmith_tool_name = tool_name # type: ignore[attr-defined]
|
||||
return _wrapped
|
||||
|
||||
|
||||
def _tool_run_matches(run: Any, args: Any, tool_name: Optional[str]) -> bool:
|
||||
"""Return whether *run* appears to be for this SDK MCP handler call.
|
||||
|
||||
Matching is intentionally strict: we require both the tool name and the
|
||||
handler args to line up with what the ``PreToolUse`` hook recorded. This
|
||||
avoids cross-attributing a handler invocation to the wrong client's active
|
||||
tool run under concurrency.
|
||||
"""
|
||||
if not tool_name:
|
||||
return False
|
||||
run_name = str(getattr(run, "name", ""))
|
||||
# SDK MCP tools show up in hook data as e.g. ``mcp__weather__get_weather``
|
||||
# while the handler only knows its short name ``get_weather``.
|
||||
name_matches = (
|
||||
tool_name == run_name or tool_name in run_name or run_name in tool_name
|
||||
)
|
||||
if not name_matches:
|
||||
return False
|
||||
inputs = getattr(run, "inputs", None)
|
||||
if not isinstance(inputs, dict):
|
||||
return False
|
||||
# PreToolUse stores {} when the tool had no inputs, otherwise
|
||||
# {"input": <tool_input>}. Normalise both sides before comparing.
|
||||
recorded = inputs.get("input", {}) if inputs else {}
|
||||
return recorded == (args or {})
|
||||
|
||||
|
||||
def _newest_matching_tool_run(
|
||||
sessions: list[SessionState], args: Any, tool_name: Optional[str]
|
||||
) -> Any:
|
||||
"""Return the most recently created active tool run that matches."""
|
||||
candidates: list[tuple[float, Any]] = []
|
||||
for candidate_session in sessions:
|
||||
for run, start_time in candidate_session.active_tool_runs.values():
|
||||
if _tool_run_matches(run, args, tool_name):
|
||||
candidates.append((start_time, run))
|
||||
if not candidates:
|
||||
return None
|
||||
return max(candidates, key=lambda item: item[0])[1]
|
||||
|
||||
|
||||
def _get_last_active_tool_run(
|
||||
session: Optional[SessionState] = None,
|
||||
*,
|
||||
args: Any = None,
|
||||
tool_name: Optional[str] = None,
|
||||
) -> Any:
|
||||
"""Return the active tool run for an SDK MCP handler, or None.
|
||||
|
||||
Lookup order:
|
||||
|
||||
1. The session explicitly bound to the handler (if any).
|
||||
2. The session bound to the current ContextVar.
|
||||
3. The module-level default session (unit tests / unbound callers).
|
||||
4. Any live client session — only used when the handler is unbound and the
|
||||
SDK invoked it in a detached async context. Requires an exact tool
|
||||
name + args match to avoid cross-attribution across clients.
|
||||
"""
|
||||
from ._hooks import (
|
||||
_current_session,
|
||||
_current_session_or_default,
|
||||
_registered_sessions,
|
||||
)
|
||||
|
||||
# If we have a specific session (explicitly bound, current-context, or the
|
||||
# test default), just return its newest active tool run. There is no
|
||||
# cross-client ambiguity at that point.
|
||||
def _newest_in(s: SessionState) -> Any:
|
||||
if not s.active_tool_runs:
|
||||
return None
|
||||
latest_id = max(
|
||||
s.active_tool_runs,
|
||||
key=lambda tid: s.active_tool_runs[tid][1],
|
||||
)
|
||||
return s.active_tool_runs[latest_id][0]
|
||||
|
||||
if session is not None:
|
||||
return _newest_in(session)
|
||||
|
||||
current_session = _current_session.get()
|
||||
if current_session is not None:
|
||||
run = _newest_in(current_session)
|
||||
if run is not None:
|
||||
return run
|
||||
|
||||
default_session = _current_session_or_default()
|
||||
if default_session is not current_session:
|
||||
run = _newest_in(default_session)
|
||||
if run is not None:
|
||||
return run
|
||||
|
||||
# Last resort: the SDK invoked this handler in a detached async context and
|
||||
# the handler object wasn't bound to a session. Require strict tool-name +
|
||||
# args match so concurrent clients can't steal each other's attribution.
|
||||
return _newest_matching_tool_run(_registered_sessions(), args, tool_name)
|
||||
|
||||
|
||||
def instrument_claude_client(original_class: Any) -> None:
|
||||
"""Patch ``ClaudeSDKClient`` **in place** to trace calls.
|
||||
|
||||
In-place patching (rather than subclassing + reference replacement)
|
||||
ensures that callers who imported ``ClaudeSDKClient`` *before*
|
||||
``configure_claude_agent_sdk()`` was called still get instrumented.
|
||||
"""
|
||||
if getattr(original_class, "_langsmith_instrumented", False):
|
||||
return # Already wrapped, avoid double-tracing
|
||||
|
||||
# ── stash originals ──────────────────────────────────────────────
|
||||
_orig_init = original_class.__init__
|
||||
_orig_query = original_class.query
|
||||
_orig_receive_response = original_class.receive_response
|
||||
|
||||
# ── patched __init__ ─────────────────────────────────────────────
|
||||
def _traced_init(self: Any, *args: Any, **kwargs: Any) -> None:
|
||||
options = kwargs.get("options") or (args[0] if args else None)
|
||||
self._ls_session = SessionState()
|
||||
if options:
|
||||
_inject_tracing_hooks(options, self._ls_session)
|
||||
_orig_init(self, *args, **kwargs)
|
||||
self._ls_prompt = None
|
||||
self._ls_start_time = None
|
||||
self._ls_streamed_input = None
|
||||
|
||||
# ── patched query ────────────────────────────────────────────────
|
||||
async def _traced_query(self: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
self._ls_start_time = time.time()
|
||||
self._ls_streamed_input = None
|
||||
prompt = args[0] if args else kwargs.get("prompt")
|
||||
|
||||
if prompt is None:
|
||||
pass
|
||||
elif isinstance(prompt, str):
|
||||
self._ls_prompt = prompt
|
||||
elif isinstance(prompt, AsyncIterable):
|
||||
collector: list[dict[str, Any]] = []
|
||||
self._ls_streamed_input = collector
|
||||
self._ls_prompt = None
|
||||
|
||||
async def _gen_wrapper() -> AsyncGenerator[dict[str, Any], None]:
|
||||
async for msg in prompt:
|
||||
collector.append(msg)
|
||||
yield msg
|
||||
|
||||
if args:
|
||||
args = (_gen_wrapper(),) + args[1:]
|
||||
else:
|
||||
kwargs["prompt"] = _gen_wrapper()
|
||||
else:
|
||||
self._ls_prompt = str(prompt)
|
||||
|
||||
return await _orig_query(self, *args, **kwargs)
|
||||
|
||||
# ── patched receive_response ─────────────────────────────────────
|
||||
async def _traced_receive_response(self: Any) -> AsyncGenerator[Any, None]:
|
||||
messages = _orig_receive_response(self)
|
||||
|
||||
trace_inputs: dict[str, Any] = {}
|
||||
trace_metadata: dict[str, Any] = {
|
||||
"ls_integration": "claude-agent-sdk",
|
||||
"ls_integration_version": _get_package_version("claude_agent_sdk"),
|
||||
}
|
||||
|
||||
awaiting_streamed_input = self._ls_streamed_input is not None
|
||||
|
||||
if self._ls_prompt:
|
||||
trace_inputs["prompt"] = self._ls_prompt
|
||||
|
||||
if hasattr(self, "options") and self.options:
|
||||
if hasattr(self.options, "system_prompt") and self.options.system_prompt:
|
||||
system_prompt = self.options.system_prompt
|
||||
if isinstance(system_prompt, str):
|
||||
trace_inputs["system"] = system_prompt
|
||||
elif isinstance(system_prompt, dict):
|
||||
if system_prompt.get("type") == "preset":
|
||||
preset_text = (
|
||||
f"preset: {system_prompt.get('preset', 'claude_code')}"
|
||||
)
|
||||
if "append" in system_prompt:
|
||||
preset_text += f"\nappend: {system_prompt['append']}"
|
||||
trace_inputs["system"] = preset_text
|
||||
else:
|
||||
trace_inputs["system"] = system_prompt
|
||||
|
||||
for attr in ["model", "permission_mode", "max_turns"]:
|
||||
if hasattr(self.options, attr):
|
||||
val = getattr(self.options, attr)
|
||||
if val is not None:
|
||||
trace_metadata[attr] = val
|
||||
|
||||
config = get_tracing_config()
|
||||
user_metadata = config.get("metadata") or {}
|
||||
|
||||
trace_kwargs: dict[str, Any] = {
|
||||
"name": config.get("name") or TRACE_CHAIN_NAME,
|
||||
"run_type": "chain",
|
||||
"inputs": trace_inputs,
|
||||
"metadata": {
|
||||
**trace_metadata,
|
||||
**user_metadata,
|
||||
"ls_agent_type": "root",
|
||||
},
|
||||
}
|
||||
if config.get("project_name"):
|
||||
trace_kwargs["project_name"] = config["project_name"]
|
||||
if config.get("tags"):
|
||||
trace_kwargs["tags"] = config["tags"]
|
||||
|
||||
async with trace(**trace_kwargs) as run:
|
||||
# Bind this client's state container to the ContextVar so stream
|
||||
# helpers on this SDK event loop pick it up (see
|
||||
# _hooks.SessionState). This keeps concurrent ClaudeSDKClient
|
||||
# instances — eval runs, FastAPI handlers, Celery workers,
|
||||
# asyncio.gather — from corrupting each other's correlation state.
|
||||
session = getattr(self, "_ls_session", None)
|
||||
if session is None:
|
||||
session = self._ls_session = SessionState()
|
||||
session_token = _register_session(session)
|
||||
_set_session_root(session, run)
|
||||
parent_token = set_parent_run_tree(run)
|
||||
tracker = TurnLifecycle(self._ls_start_time)
|
||||
collected_by_ctx: dict[Optional[str], list[dict[str, Any]]] = {None: []}
|
||||
|
||||
prompt_for_llm: Any = self._ls_prompt
|
||||
|
||||
try:
|
||||
async for msg in messages:
|
||||
if awaiting_streamed_input and self._ls_streamed_input:
|
||||
unwrapped_messages = unwrap_message_dicts(
|
||||
self._ls_streamed_input
|
||||
)
|
||||
if unwrapped_messages:
|
||||
run.inputs["messages"] = unwrapped_messages
|
||||
prompt_for_llm = self._ls_streamed_input
|
||||
awaiting_streamed_input = False
|
||||
|
||||
msg_type = type(msg).__name__
|
||||
|
||||
if msg_type == "AssistantMessage":
|
||||
parent_tool_use_id = getattr(msg, "parent_tool_use_id", None)
|
||||
llm_parent = (
|
||||
get_subagent_run_by_tool_id(parent_tool_use_id)
|
||||
if parent_tool_use_id
|
||||
else None
|
||||
)
|
||||
|
||||
ctx_key = parent_tool_use_id
|
||||
ctx_history = collected_by_ctx.setdefault(ctx_key, [])
|
||||
|
||||
content = tracker.start_llm_run(
|
||||
msg,
|
||||
prompt_for_llm if parent_tool_use_id is None else None,
|
||||
ctx_history,
|
||||
parent=llm_parent,
|
||||
)
|
||||
if content:
|
||||
ctx_history.append(content)
|
||||
|
||||
elif msg_type == "UserMessage":
|
||||
parent_tool_use_id = getattr(msg, "parent_tool_use_id", None)
|
||||
ctx_key = parent_tool_use_id
|
||||
ctx_history = collected_by_ctx.setdefault(ctx_key, [])
|
||||
|
||||
if hasattr(msg, "content"):
|
||||
flattened = flatten_content_blocks(msg.content)
|
||||
if (
|
||||
isinstance(flattened, list)
|
||||
and flattened
|
||||
and isinstance(flattened[0], dict)
|
||||
and flattened[0].get("type") == "tool_result"
|
||||
):
|
||||
for block in flattened:
|
||||
tool_use_id = block.get("tool_use_id")
|
||||
ctx_history.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"content": block.get("content", ""),
|
||||
"tool_call_id": tool_use_id,
|
||||
}
|
||||
)
|
||||
if (
|
||||
tool_use_id
|
||||
and tool_use_id in session.active_tool_runs
|
||||
):
|
||||
tool_run, _ = session.active_tool_runs.pop(
|
||||
tool_use_id
|
||||
)
|
||||
result_content = block.get("content", "")
|
||||
is_error = block.get("is_error", False)
|
||||
tool_run.end(
|
||||
outputs={"output": result_content},
|
||||
error=str(result_content)
|
||||
if is_error
|
||||
else None,
|
||||
)
|
||||
try:
|
||||
tool_run.patch()
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to patch"
|
||||
f" orphaned tool run: {e}"
|
||||
)
|
||||
else:
|
||||
ctx_history.append(
|
||||
{
|
||||
"content": flattened,
|
||||
"role": "user",
|
||||
}
|
||||
)
|
||||
tracker.mark_next_start()
|
||||
elif msg_type == "ResultMessage":
|
||||
meta = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"num_turns": getattr(msg, "num_turns", None),
|
||||
"session_id": getattr(msg, "session_id", None),
|
||||
"duration_ms": getattr(msg, "duration_ms", None),
|
||||
"duration_api_ms": getattr(
|
||||
msg, "duration_api_ms", None
|
||||
),
|
||||
"is_error": getattr(msg, "is_error", None),
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
if meta:
|
||||
run.metadata.update(meta)
|
||||
|
||||
yield msg
|
||||
main_collected = collected_by_ctx.get(None, [])
|
||||
run.end(outputs=main_collected[-1] if main_collected else None)
|
||||
except Exception:
|
||||
logger.exception("Error while tracing Claude Agent stream")
|
||||
finally:
|
||||
tracker.close()
|
||||
reconcile_from_transcripts(tracker, session=session)
|
||||
tracker.flush()
|
||||
clear_parent_run_tree(parent_token)
|
||||
try:
|
||||
clear_active_tool_runs(session)
|
||||
finally:
|
||||
_unregister_session(session, session_token)
|
||||
|
||||
# ── apply patches to the class itself ────────────────────────────
|
||||
original_class.__init__ = _traced_init
|
||||
original_class.query = _traced_query
|
||||
original_class.receive_response = _traced_receive_response
|
||||
original_class._langsmith_instrumented = True
|
||||
|
||||
|
||||
def instrument_sdk_mcp_tool(tool_class: Any) -> None:
|
||||
"""Patch ``SdkMcpTool.__init__`` to auto-wrap handlers.
|
||||
|
||||
Wrapping happens at construction time so that any tool created
|
||||
*after* ``configure_claude_agent_sdk()`` automatically gets
|
||||
run-context propagation, regardless of how ``tool`` or
|
||||
``create_sdk_mcp_server`` were imported.
|
||||
"""
|
||||
if getattr(tool_class, "_langsmith_handler_patched", False):
|
||||
return
|
||||
|
||||
_orig_init = tool_class.__init__
|
||||
|
||||
def _patched_init(self: Any, *args: Any, **kwargs: Any) -> None:
|
||||
_orig_init(self, *args, **kwargs)
|
||||
handler = self.handler
|
||||
if callable(handler) and not getattr(handler, "_langsmith_wrapped", False):
|
||||
self.handler = _wrap_tool_handler(
|
||||
handler, tool_name=getattr(self, "name", None)
|
||||
)
|
||||
|
||||
tool_class.__init__ = _patched_init
|
||||
tool_class._langsmith_handler_patched = True
|
||||
@@ -0,0 +1,39 @@
|
||||
"""Configuration management for Claude Agent SDK tracing."""
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
# Global configuration for tracing
|
||||
_tracing_config: dict[str, Any] = {
|
||||
"name": None,
|
||||
"project_name": None,
|
||||
"metadata": None,
|
||||
"tags": None,
|
||||
}
|
||||
|
||||
|
||||
def set_tracing_config(
|
||||
name: Optional[str] = None,
|
||||
project_name: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
tags: Optional[list[str]] = None,
|
||||
) -> None:
|
||||
"""Set the global tracing configuration for Claude Agent SDK.
|
||||
|
||||
Args:
|
||||
name: Name of the root trace.
|
||||
project_name: LangSmith project to trace to.
|
||||
metadata: Metadata to associate with all traces.
|
||||
tags: Tags to associate with all traces.
|
||||
"""
|
||||
global _tracing_config
|
||||
_tracing_config = {
|
||||
"name": name,
|
||||
"project_name": project_name,
|
||||
"metadata": metadata,
|
||||
"tags": tags,
|
||||
}
|
||||
|
||||
|
||||
def get_tracing_config() -> dict[str, Any]:
|
||||
"""Get the current tracing configuration."""
|
||||
return _tracing_config.copy()
|
||||
@@ -0,0 +1,564 @@
|
||||
"""Hook-based tool tracing for Claude Agent SDK.
|
||||
|
||||
Correlation state is scoped **per client session** via a
|
||||
:class:`contextvars.ContextVar`. Each instrumented ``ClaudeSDKClient`` owns a
|
||||
:class:`SessionState`; ``receive_response()`` binds it while processing the
|
||||
stream so helper functions can look up the right state regardless of how many
|
||||
clients are concurrently active in the process.
|
||||
|
||||
Hooks injected by ``_client.py`` are also bound to their owning
|
||||
``SessionState`` so hook callbacks use the correct state even if the SDK runs
|
||||
them in an async context that did not inherit ``receive_response``'s
|
||||
ContextVar.
|
||||
|
||||
When no ContextVar is active, hooks use a module-level default session. This is
|
||||
primarily for direct unit tests; real traffic under ``receive_response`` uses a
|
||||
client-bound session.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from langsmith.run_helpers import get_current_run_tree
|
||||
from langsmith.run_trees import RunTree
|
||||
|
||||
from ._tools import get_parent_run_tree
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from claude_agent_sdk import (
|
||||
HookContext,
|
||||
HookInput,
|
||||
HookJSONOutput,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ── Per-session state ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionState:
|
||||
"""All mutable correlation state for a single conversation.
|
||||
|
||||
One instance is created per instrumented ``ClaudeSDKClient`` and bound to
|
||||
the ``_current_session`` ContextVar while that client is active.
|
||||
"""
|
||||
|
||||
# Key: tool_use_id → (run_tree, start_time)
|
||||
active_tool_runs: dict[str, tuple[Any, float]] = field(default_factory=dict)
|
||||
|
||||
# Key: agent_id → RunTree for the subagent chain.
|
||||
# Populated by SubagentStart, consumed by SubagentStop.
|
||||
subagent_runs: dict[str, RunTree] = field(default_factory=dict)
|
||||
|
||||
# Key: tool_use_id → tool_input dict.
|
||||
# When PreToolUse fires for an "Agent" tool, it stashes here.
|
||||
# SubagentStart pops it to find the matching Agent tool run.
|
||||
pending_agent_tools: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||||
|
||||
# Key: agent_id → Agent tool_use_id.
|
||||
# Maps a subagent back to the Agent tool that spawned it.
|
||||
agent_to_tool_mapping: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
# Key: Agent tool_use_id → RunTree.
|
||||
# SubagentStop moves the run here; PostToolUse sets outputs on it;
|
||||
# clear_active_tool_runs() ends + patches it.
|
||||
ended_subagent_runs: dict[str, RunTree] = field(default_factory=dict)
|
||||
|
||||
# (transcript_path, subagent_RunTree) captured from SubagentStop.
|
||||
# Used for usage extraction and creating missing LLM runs.
|
||||
subagent_transcript_paths: list[tuple[str, RunTree]] = field(default_factory=list)
|
||||
|
||||
# Main session transcript path, captured from BaseHookInput.transcript_path
|
||||
# on the first hook that fires (every hook inherits this field).
|
||||
main_transcript_path: Optional[str] = None
|
||||
|
||||
# Root LangSmith run used for parenting root-level hook spans.
|
||||
root_run: Optional[RunTree] = None
|
||||
|
||||
|
||||
# Module-level *default* session. Used when no ContextVar is set (e.g. tests
|
||||
# that poke hooks directly, or hooks firing outside a traced conversation).
|
||||
_default_session: SessionState = SessionState()
|
||||
|
||||
# ContextVar holding the active session for a conversation. Injected hook
|
||||
# callables bind this explicitly before calling the shared hook function.
|
||||
_current_session: ContextVar[Optional[SessionState]] = ContextVar(
|
||||
"langsmith_claude_agent_session", default=None
|
||||
)
|
||||
|
||||
# Live sessions are only used by SDK MCP tool handlers when the SDK invokes the
|
||||
# handler in a detached async context that did not inherit _current_session.
|
||||
# Store weak values so this fallback registry never owns session lifetime.
|
||||
_live_sessions_lock = threading.Lock()
|
||||
_live_sessions: weakref.WeakValueDictionary[int, SessionState] = (
|
||||
weakref.WeakValueDictionary()
|
||||
)
|
||||
|
||||
|
||||
def _current_session_or_default() -> SessionState:
|
||||
"""Return the session bound to the current context, or the default."""
|
||||
session = _current_session.get()
|
||||
if session is not None:
|
||||
return session
|
||||
return _default_session
|
||||
|
||||
|
||||
def _session_for_hook() -> SessionState:
|
||||
"""Resolve the session that owns the current hook invocation.
|
||||
|
||||
Real Claude SDK hook invocations are wrapped by ``_bind_hook_to_session``
|
||||
in ``_client.py``, so the ContextVar should be set. The default session is
|
||||
only for tests or direct, unbound hook calls.
|
||||
"""
|
||||
return _current_session_or_default()
|
||||
|
||||
|
||||
def _register_session(session: SessionState) -> object:
|
||||
"""Bind *session* to the ContextVar and return a reset token.
|
||||
|
||||
The caller must pass the returned token to ``_unregister_session`` when
|
||||
the conversation ends.
|
||||
"""
|
||||
with _live_sessions_lock:
|
||||
_live_sessions[id(session)] = session
|
||||
return _current_session.set(session)
|
||||
|
||||
|
||||
def _set_session_root(session: SessionState, run_tree: RunTree) -> None:
|
||||
"""Store the root LangSmith run for *session*."""
|
||||
session.root_run = run_tree
|
||||
|
||||
|
||||
def _unregister_session(session: SessionState, token: Any) -> None:
|
||||
"""Reset the ContextVar for the current session and drop the live entry."""
|
||||
try:
|
||||
_current_session.reset(token)
|
||||
except ValueError:
|
||||
# Token was created in a different context. Don't clobber an unrelated
|
||||
# current value — just log and continue. The live-sessions registry
|
||||
# below is still cleaned up so matching won't find a stale session.
|
||||
logger.debug("Could not reset _current_session with token from another context")
|
||||
finally:
|
||||
with _live_sessions_lock:
|
||||
_live_sessions.pop(id(session), None)
|
||||
|
||||
|
||||
def _registered_sessions() -> list[SessionState]:
|
||||
"""Return currently active client sessions."""
|
||||
with _live_sessions_lock:
|
||||
return list(_live_sessions.values())
|
||||
|
||||
|
||||
# ── Public helpers (used by _client.py) ───────────────────────────────────────
|
||||
|
||||
|
||||
def get_subagent_run_by_tool_id(tool_use_id: str) -> Optional[RunTree]:
|
||||
"""Get a subagent run by the Agent tool's tool_use_id.
|
||||
|
||||
Checks both active subagent runs and ended-but-not-finalised runs,
|
||||
because the SDK fires ``SubagentStop`` before the subagent's messages
|
||||
reach the client.
|
||||
"""
|
||||
session = _current_session_or_default()
|
||||
# Check active subagents first
|
||||
for aid, tid in session.agent_to_tool_mapping.items():
|
||||
if tid == tool_use_id:
|
||||
return session.subagent_runs.get(aid)
|
||||
# Fall back to ended-but-not-finalised subagents
|
||||
return session.ended_subagent_runs.get(tool_use_id)
|
||||
|
||||
|
||||
# ── Hook functions ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def pre_tool_use_hook(
|
||||
input_data: "HookInput",
|
||||
tool_use_id: Optional[str],
|
||||
context: "HookContext",
|
||||
) -> "HookJSONOutput":
|
||||
"""Trace tool execution before it starts.
|
||||
|
||||
Args:
|
||||
input_data: Contains `tool_name`, `tool_input`, `session_id`, `agent_id`
|
||||
tool_use_id: Unique identifier for this tool invocation
|
||||
context: Hook context (currently contains only signal)
|
||||
|
||||
Returns:
|
||||
Hook output (empty dict allows execution to proceed)
|
||||
"""
|
||||
if not tool_use_id:
|
||||
return {}
|
||||
|
||||
data: dict[str, Any] = dict(input_data) # flatten TypedDict union
|
||||
tool_name: str = str(data.get("tool_name", "unknown_tool"))
|
||||
tool_input: dict[str, Any] = dict(data.get("tool_input") or {})
|
||||
agent_id: Optional[str] = str(data["agent_id"]) if data.get("agent_id") else None
|
||||
session = _session_for_hook()
|
||||
|
||||
# Capture main session transcript path from BaseHookInput
|
||||
if session.main_transcript_path is None and data.get("transcript_path"):
|
||||
session.main_transcript_path = str(data["transcript_path"])
|
||||
|
||||
# If this is an Agent tool call, record it so SubagentStart can find it
|
||||
if tool_name == "Agent":
|
||||
session.pending_agent_tools[tool_use_id] = tool_input
|
||||
|
||||
try:
|
||||
# Determine parent: subagent chain > root chain.
|
||||
# Tool runs are siblings of LLM runs, not children.
|
||||
parent: Optional[RunTree] = None
|
||||
if agent_id and agent_id in session.subagent_runs:
|
||||
parent = session.subagent_runs[agent_id]
|
||||
else:
|
||||
parent = (
|
||||
session.root_run
|
||||
if session is not _default_session
|
||||
else get_parent_run_tree()
|
||||
) or get_current_run_tree()
|
||||
|
||||
if not parent:
|
||||
return {}
|
||||
|
||||
start_time = time.time()
|
||||
tool_run = parent.create_child(
|
||||
name=tool_name,
|
||||
run_type="tool",
|
||||
inputs={"input": tool_input} if tool_input else {},
|
||||
start_time=datetime.fromtimestamp(start_time, tz=timezone.utc),
|
||||
)
|
||||
|
||||
try:
|
||||
tool_run.post()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to post tool run for {tool_name}: {e}")
|
||||
|
||||
session.active_tool_runs[tool_use_id] = (tool_run, start_time)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Error in PreToolUse hook for {tool_name}: {e}", exc_info=True)
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
async def post_tool_use_hook(
|
||||
input_data: "HookInput",
|
||||
tool_use_id: Optional[str],
|
||||
context: "HookContext",
|
||||
) -> "HookJSONOutput":
|
||||
"""Trace tool execution after it completes.
|
||||
|
||||
Args:
|
||||
input_data: Contains `tool_name`, `tool_input`, `tool_response`,
|
||||
`session_id`, etc.
|
||||
tool_use_id: Unique identifier for this tool invocation
|
||||
context: Hook context (currently contains only signal)
|
||||
|
||||
Returns:
|
||||
Hook output (empty `dict` by default)
|
||||
"""
|
||||
if not tool_use_id:
|
||||
return {}
|
||||
|
||||
tool_name: str = str(input_data.get("tool_name", "unknown_tool"))
|
||||
tool_response = input_data.get("tool_response")
|
||||
session = _session_for_hook()
|
||||
|
||||
try:
|
||||
run_info = session.active_tool_runs.pop(tool_use_id, None)
|
||||
if not run_info:
|
||||
return {}
|
||||
|
||||
tool_run, _ = run_info
|
||||
|
||||
if isinstance(tool_response, dict):
|
||||
outputs = tool_response
|
||||
elif isinstance(tool_response, list):
|
||||
outputs = {"content": tool_response}
|
||||
else:
|
||||
outputs = {"output": str(tool_response)} if tool_response else {}
|
||||
|
||||
# Check if the tool execution was an error
|
||||
is_error = False
|
||||
if isinstance(tool_response, dict):
|
||||
is_error = tool_response.get("is_error", False)
|
||||
|
||||
tool_run.end(
|
||||
outputs=outputs,
|
||||
error=outputs.get("output") if is_error else None,
|
||||
)
|
||||
|
||||
try:
|
||||
tool_run.patch()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to patch tool run for {tool_name}: {e}")
|
||||
|
||||
# If this is an Agent tool, also set outputs on the stashed
|
||||
# subagent run. We don't end/patch the subagent here because
|
||||
# its AssistantMessages may not have been yielded to
|
||||
# receive_response() yet. clear_active_tool_runs() will
|
||||
# finalise it at the end of the conversation.
|
||||
subagent_run = session.ended_subagent_runs.get(tool_use_id)
|
||||
if subagent_run:
|
||||
try:
|
||||
subagent_run.outputs = outputs
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to set subagent run outputs: {e}")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Error in PostToolUse hook for {tool_name}: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
async def post_tool_use_failure_hook(
|
||||
input_data: "HookInput",
|
||||
tool_use_id: Optional[str],
|
||||
context: "HookContext",
|
||||
) -> "HookJSONOutput":
|
||||
"""Trace tool execution when it fails.
|
||||
|
||||
This hook fires for built-in tool failures (Bash, Read, Write, etc.)
|
||||
and is mutually exclusive with :func:`post_tool_use_hook` — when a
|
||||
built-in tool fails, only ``PostToolUseFailure`` fires.
|
||||
|
||||
Args:
|
||||
input_data: Contains ``tool_name``, ``tool_input``, ``error``,
|
||||
and optionally ``is_interrupt``.
|
||||
tool_use_id: Unique identifier for this tool invocation
|
||||
context: Hook context (currently contains only signal)
|
||||
|
||||
Returns:
|
||||
Hook output (empty dict)
|
||||
"""
|
||||
if not tool_use_id:
|
||||
return {}
|
||||
|
||||
tool_name: str = str(input_data.get("tool_name", "unknown_tool"))
|
||||
error: str = str(input_data.get("error", "Unknown error"))
|
||||
session = _session_for_hook()
|
||||
|
||||
try:
|
||||
run_info = session.active_tool_runs.pop(tool_use_id, None)
|
||||
if not run_info:
|
||||
return {}
|
||||
|
||||
tool_run, _ = run_info
|
||||
|
||||
tool_run.end(
|
||||
outputs={"error": error},
|
||||
error=error,
|
||||
)
|
||||
|
||||
try:
|
||||
tool_run.patch()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to patch failed tool run for {tool_name}: {e}")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Error in PostToolUseFailure hook for {tool_name}: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
async def subagent_start_hook(
|
||||
input_data: "HookInput",
|
||||
tool_use_id: Optional[str],
|
||||
context: "HookContext",
|
||||
) -> "HookJSONOutput":
|
||||
"""Create a chain run when a subagent starts.
|
||||
|
||||
The subagent chain is nested under the Agent tool run that spawned it.
|
||||
Since the SDK passes a different ``tool_use_id`` to this hook than the
|
||||
one from ``PreToolUse`` for the Agent tool, we match them via the
|
||||
``_pending_agent_tools`` queue.
|
||||
|
||||
Args:
|
||||
input_data: Contains ``agent_id``, ``agent_type``, ``session_id``
|
||||
tool_use_id: SDK-internal session id (not the Agent tool's
|
||||
tool_use_id)
|
||||
context: Hook context
|
||||
|
||||
Returns:
|
||||
Hook output (empty dict)
|
||||
"""
|
||||
data: dict[str, Any] = dict(input_data)
|
||||
agent_id: Optional[str] = str(data["agent_id"]) if data.get("agent_id") else None
|
||||
agent_type: str = str(data.get("agent_type") or "subagent")
|
||||
session = _session_for_hook()
|
||||
|
||||
if not agent_id:
|
||||
return {}
|
||||
|
||||
try:
|
||||
# Find the Agent tool run that triggered this subagent.
|
||||
# pending_agent_tools is populated by pre_tool_use_hook when
|
||||
# tool_name == "Agent". Pop the most recent one.
|
||||
agent_tool_use_id: Optional[str] = None
|
||||
agent_tool_input: dict[str, Any] = {}
|
||||
parent: Optional[RunTree] = None
|
||||
|
||||
if session.pending_agent_tools:
|
||||
agent_tool_use_id, agent_tool_input = session.pending_agent_tools.popitem()
|
||||
|
||||
if agent_tool_use_id in session.active_tool_runs:
|
||||
agent_tool_run, _ = session.active_tool_runs[agent_tool_use_id]
|
||||
parent = agent_tool_run
|
||||
|
||||
if parent is None:
|
||||
parent = (
|
||||
session.root_run
|
||||
if session is not _default_session
|
||||
else get_parent_run_tree()
|
||||
) or get_current_run_tree()
|
||||
|
||||
if not parent:
|
||||
return {}
|
||||
|
||||
start_time = time.time()
|
||||
subagent_run = parent.create_child(
|
||||
name=agent_type,
|
||||
run_type="chain",
|
||||
inputs=agent_tool_input if agent_tool_input else {},
|
||||
start_time=datetime.fromtimestamp(start_time, tz=timezone.utc),
|
||||
)
|
||||
subagent_run.extra["metadata"] = {
|
||||
**subagent_run.extra.get("metadata", {}),
|
||||
"ls_agent_type": "subagent",
|
||||
}
|
||||
|
||||
try:
|
||||
subagent_run.post()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to post subagent run: {e}")
|
||||
|
||||
# Store by agent_id so tool hooks and LLM run lookup can find it
|
||||
session.subagent_runs[agent_id] = subagent_run
|
||||
|
||||
# Remember which Agent tool_use_id spawned this agent_id
|
||||
if agent_tool_use_id:
|
||||
session.agent_to_tool_mapping[agent_id] = agent_tool_use_id
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Error in SubagentStart hook: {e}", exc_info=True)
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
async def subagent_stop_hook(
|
||||
input_data: "HookInput",
|
||||
tool_use_id: Optional[str],
|
||||
context: "HookContext",
|
||||
) -> "HookJSONOutput":
|
||||
"""Move the subagent run to ended state when it finishes.
|
||||
|
||||
Does NOT end/patch the run — ``PostToolUse`` for the Agent tool will
|
||||
set outputs, and ``clear_active_tool_runs()`` will finalise it at the
|
||||
end of the conversation.
|
||||
|
||||
Args:
|
||||
input_data: Contains ``agent_id``, ``agent_type``, ``session_id``,
|
||||
``agent_transcript_path``
|
||||
tool_use_id: SDK-internal session id
|
||||
context: Hook context
|
||||
|
||||
Returns:
|
||||
Hook output (empty dict)
|
||||
"""
|
||||
data: dict[str, Any] = dict(input_data)
|
||||
agent_id: Optional[str] = str(data["agent_id"]) if data.get("agent_id") else None
|
||||
transcript_path: Optional[str] = (
|
||||
str(data["agent_transcript_path"])
|
||||
if data.get("agent_transcript_path")
|
||||
else None
|
||||
)
|
||||
session = _session_for_hook()
|
||||
|
||||
if not agent_id:
|
||||
return {}
|
||||
|
||||
try:
|
||||
subagent_run = session.subagent_runs.pop(agent_id, None)
|
||||
if not subagent_run:
|
||||
return {}
|
||||
|
||||
if transcript_path:
|
||||
session.subagent_transcript_paths.append((transcript_path, subagent_run))
|
||||
|
||||
# Move to ended state so PostToolUse can set outputs.
|
||||
agent_tool_id = session.agent_to_tool_mapping.pop(agent_id, None)
|
||||
if agent_tool_id:
|
||||
session.ended_subagent_runs[agent_tool_id] = subagent_run
|
||||
else:
|
||||
# No matching Agent tool — just end it now
|
||||
subagent_run.end()
|
||||
try:
|
||||
subagent_run.patch()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to patch subagent run: {e}")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Error in SubagentStop hook: {e}", exc_info=True)
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
# ── Cleanup ───────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def clear_active_tool_runs(session: Optional[SessionState] = None) -> None:
|
||||
"""Finalise all runs and clear state for *session*.
|
||||
|
||||
If *session* is omitted the current ContextVar-bound session is used
|
||||
(falling back to the module-level default session). ``receive_response``
|
||||
passes the per-call session explicitly.
|
||||
"""
|
||||
if session is None:
|
||||
session = _current_session_or_default()
|
||||
|
||||
# 1. End orphaned subagent runs (SubagentStop never fired)
|
||||
for agent_id, subagent_run in session.subagent_runs.items():
|
||||
try:
|
||||
subagent_run.end(error="Subagent run not completed (conversation ended)")
|
||||
subagent_run.patch()
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to clean up orphaned subagent run {agent_id}: {e}")
|
||||
|
||||
# 2. Finalise ended subagent runs (outputs already set by PostToolUse)
|
||||
for tool_use_id, subagent_run in session.ended_subagent_runs.items():
|
||||
try:
|
||||
subagent_run.end()
|
||||
subagent_run.patch()
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to finalise ended subagent run {tool_use_id}: {e}")
|
||||
|
||||
# 3. End orphaned tool runs
|
||||
for tool_use_id, (tool_run, _) in session.active_tool_runs.items():
|
||||
try:
|
||||
tool_run.end(error="Tool run not completed (conversation ended)")
|
||||
tool_run.patch()
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to clean up orphaned tool run {tool_use_id}: {e}")
|
||||
|
||||
# 4. Reset session state
|
||||
session.active_tool_runs.clear()
|
||||
session.subagent_runs.clear()
|
||||
session.pending_agent_tools.clear()
|
||||
session.agent_to_tool_mapping.clear()
|
||||
session.ended_subagent_runs.clear()
|
||||
session.subagent_transcript_paths.clear()
|
||||
session.main_transcript_path = None
|
||||
session.root_run = None
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Message processing and content serialization for Claude Agent SDK."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _extract_tool_result_text(content: Any) -> str:
|
||||
"""Extract text content from tool result content blocks."""
|
||||
if content is None:
|
||||
return ""
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
texts = []
|
||||
for item in content:
|
||||
if isinstance(item, dict):
|
||||
if item.get("type") == "text":
|
||||
texts.append(item.get("text", ""))
|
||||
elif hasattr(item, "text"):
|
||||
texts.append(getattr(item, "text", ""))
|
||||
return "\n".join(texts) if texts else str(content)
|
||||
return str(content)
|
||||
|
||||
|
||||
def flatten_content_blocks(content: Any) -> Any:
|
||||
"""Convert SDK content blocks into serializable dicts using explicit type checks."""
|
||||
if not isinstance(content, list):
|
||||
return content
|
||||
|
||||
result = []
|
||||
for block in content:
|
||||
block_type = type(block).__name__
|
||||
|
||||
# Handle known Claude SDK block types
|
||||
if block_type == "TextBlock":
|
||||
result.append(
|
||||
{
|
||||
"type": "text",
|
||||
"text": getattr(block, "text", ""),
|
||||
}
|
||||
)
|
||||
elif block_type == "ThinkingBlock":
|
||||
result.append(
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": getattr(block, "thinking", ""),
|
||||
"signature": getattr(block, "signature", ""),
|
||||
}
|
||||
)
|
||||
elif block_type == "ToolUseBlock":
|
||||
result.append(
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": getattr(block, "id", None),
|
||||
"name": getattr(block, "name", None),
|
||||
"input": getattr(block, "input", None),
|
||||
}
|
||||
)
|
||||
elif block_type == "ToolResultBlock":
|
||||
# Extract text from nested content for tool results
|
||||
tool_content = getattr(block, "content", None)
|
||||
content_text = _extract_tool_result_text(tool_content)
|
||||
result.append(
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": getattr(block, "tool_use_id", None),
|
||||
"content": content_text,
|
||||
"is_error": getattr(block, "is_error", False),
|
||||
}
|
||||
)
|
||||
else:
|
||||
result.append(block)
|
||||
return result
|
||||
|
||||
|
||||
def unwrap_message_dicts(messages: list[Any]) -> list[dict[str, Any]]:
|
||||
"""Normalize SDK message dicts into ``{role, content}`` form.
|
||||
|
||||
The Claude SDK wraps messages in ``{"message": {"role": ..., "content": ...}}``
|
||||
envelopes. This function unwraps them into a flat list.
|
||||
"""
|
||||
result: list[dict[str, Any]] = []
|
||||
for msg in messages:
|
||||
if not isinstance(msg, dict):
|
||||
result.append(msg)
|
||||
continue
|
||||
if "message" in msg:
|
||||
inner = msg["message"]
|
||||
if isinstance(inner, dict):
|
||||
result.append(
|
||||
{
|
||||
"role": inner.get("role", "user"),
|
||||
"content": inner.get("content", ""),
|
||||
}
|
||||
)
|
||||
else:
|
||||
result.append(msg)
|
||||
else:
|
||||
result.append(msg)
|
||||
return result
|
||||
|
||||
|
||||
def build_llm_input(prompt: Any, history: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Construct a combined prompt + history message list."""
|
||||
if isinstance(prompt, str):
|
||||
entry = {"content": prompt, "role": "user"}
|
||||
return [entry, *history] if history else [entry]
|
||||
|
||||
if isinstance(prompt, list):
|
||||
formatted = unwrap_message_dicts(prompt)
|
||||
return [*formatted, *history] if history else formatted
|
||||
|
||||
return list(history) if history else []
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Context-var storage utilities for Claude Agent SDK tracing.
|
||||
|
||||
This module stores the *parent run tree* — the root chain span opened by
|
||||
``_traced_receive_response`` — for direct/default-session hook calls. Real
|
||||
instrumented client sessions primarily parent hook spans via
|
||||
``SessionState.root_run``.
|
||||
|
||||
A :class:`contextvars.ContextVar` is used so concurrent conversations on the
|
||||
same thread (e.g. multiple ``ClaudeSDKClient`` instances driven by
|
||||
``asyncio.gather``) each see their own parent run tree.
|
||||
"""
|
||||
|
||||
from contextvars import ContextVar
|
||||
from typing import Any, Optional
|
||||
|
||||
_parent_run_tree: ContextVar[Optional[Any]] = ContextVar(
|
||||
"langsmith_claude_agent_parent_run_tree", default=None
|
||||
)
|
||||
|
||||
|
||||
def set_parent_run_tree(run_tree: Any) -> Any:
|
||||
"""Bind *run_tree* to the current context and return a reset token."""
|
||||
return _parent_run_tree.set(run_tree)
|
||||
|
||||
|
||||
def clear_parent_run_tree(token: Any = None) -> None:
|
||||
"""Reset the parent run tree in the current context.
|
||||
|
||||
If a *token* from :func:`set_parent_run_tree` is provided, it is used
|
||||
to restore the previous value; otherwise the context is cleared.
|
||||
"""
|
||||
if token is not None:
|
||||
try:
|
||||
_parent_run_tree.reset(token)
|
||||
except ValueError:
|
||||
_parent_run_tree.set(None)
|
||||
else:
|
||||
_parent_run_tree.set(None)
|
||||
|
||||
|
||||
def get_parent_run_tree() -> Any:
|
||||
"""Return the parent run tree bound to the current context."""
|
||||
return _parent_run_tree.get()
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Post-conversation transcript reconciliation.
|
||||
|
||||
After a conversation ends this module:
|
||||
|
||||
1. Creates LLM runs for subagent turns that were not relayed through
|
||||
the parent stream (the SDK only streams the first assistant message
|
||||
per subagent; subsequent turns are folded into the Agent tool result).
|
||||
|
||||
2. Patches accurate token usage onto all LLM runs from the JSONL
|
||||
transcripts (the live stream only has partial streaming counts).
|
||||
|
||||
.. note::
|
||||
|
||||
The transcript JSONL format is **not a contracted API** of the Claude
|
||||
Agent SDK. Changes to the format could silently degrade trace
|
||||
fidelity. If the SDK begins relaying all subagent messages through
|
||||
the stream, step 1 becomes a no-op (the dedup guard skips already-
|
||||
seen ``message_id`` values).
|
||||
"""
|
||||
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from ._hooks import SessionState, _current_session_or_default
|
||||
from ._usage import (
|
||||
extract_usage_metadata,
|
||||
read_llm_turns_from_transcript,
|
||||
read_usage_from_transcript,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._client import TurnLifecycle
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
LLM_RUN_NAME = "claude.assistant.turn"
|
||||
|
||||
|
||||
def reconcile_from_transcripts(
|
||||
tracker: "TurnLifecycle",
|
||||
session: Optional[SessionState] = None,
|
||||
) -> None:
|
||||
"""Read transcripts and reconcile LLM runs.
|
||||
|
||||
This function does two things after the conversation ends:
|
||||
|
||||
1. **Missing subagent LLM runs** — creates LLM runs for subagent
|
||||
turns whose ``message_id`` is not already in
|
||||
``tracker.llm_runs_by_message_id``.
|
||||
|
||||
2. **Usage correction** — patches accurate usage from the JSONL
|
||||
transcripts onto all LLM runs (both streamed and synthetic).
|
||||
|
||||
If *session* is omitted the current ContextVar-bound session (or the
|
||||
module-level default) is used. The caller in ``receive_response``
|
||||
passes the per-conversation session explicitly to stay safe under
|
||||
concurrent tracing.
|
||||
"""
|
||||
if session is None:
|
||||
session = _current_session_or_default()
|
||||
_create_missing_subagent_llm_runs(tracker, session)
|
||||
_patch_usage_on_llm_runs(tracker, session)
|
||||
|
||||
|
||||
# ── Step 1: synthetic subagent LLM runs ─────────────────────────────
|
||||
|
||||
|
||||
def _create_missing_subagent_llm_runs(
|
||||
tracker: "TurnLifecycle",
|
||||
session: SessionState,
|
||||
) -> None:
|
||||
# Guard against the same message_id being processed twice (e.g. if
|
||||
# the same transcript path appears multiple times in the list).
|
||||
created: set[str] = set()
|
||||
for path, subagent_run in session.subagent_transcript_paths:
|
||||
try:
|
||||
turns = read_llm_turns_from_transcript(path)
|
||||
for turn in turns:
|
||||
mid = turn["message_id"]
|
||||
if mid in tracker.llm_runs_by_message_id:
|
||||
continue
|
||||
if mid in created:
|
||||
continue
|
||||
|
||||
ts = _parse_timestamp(turn.get("timestamp"))
|
||||
|
||||
input_messages = turn.get("input_messages", [])
|
||||
llm_metadata: dict[str, Any] = {
|
||||
"ls_provider": "anthropic",
|
||||
}
|
||||
if turn.get("model"):
|
||||
llm_metadata["ls_model_name"] = turn["model"]
|
||||
|
||||
llm_run = subagent_run.create_child(
|
||||
name=LLM_RUN_NAME,
|
||||
run_type="llm",
|
||||
inputs={"messages": input_messages} if input_messages else {},
|
||||
extra={"metadata": llm_metadata},
|
||||
start_time=ts,
|
||||
)
|
||||
|
||||
llm_run.outputs = {
|
||||
"content": turn.get("content", []),
|
||||
"role": "assistant",
|
||||
}
|
||||
|
||||
raw_usage = turn.get("usage")
|
||||
if raw_usage:
|
||||
usage_meta = extract_usage_metadata(raw_usage)
|
||||
if usage_meta:
|
||||
meta = llm_run.extra.setdefault("metadata", {})
|
||||
meta["usage_metadata"] = usage_meta
|
||||
|
||||
llm_run.end(end_time=ts)
|
||||
try:
|
||||
llm_run.post()
|
||||
llm_run.patch()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to post/patch subagent LLM run: {e}")
|
||||
|
||||
tracker.llm_runs_by_message_id[mid] = llm_run
|
||||
created.add(mid)
|
||||
logger.debug(f"Created missing subagent LLM run for message {mid}")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to create subagent LLM runs from {path}: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
|
||||
# ── Step 2: usage patching ──────────────────────────────────────────
|
||||
|
||||
|
||||
def _patch_usage_on_llm_runs(
|
||||
tracker: "TurnLifecycle",
|
||||
session: SessionState,
|
||||
) -> None:
|
||||
if not tracker.llm_runs_by_message_id:
|
||||
return
|
||||
|
||||
all_usage: dict[str, dict[str, Any]] = {}
|
||||
|
||||
main_path = session.main_transcript_path
|
||||
if main_path:
|
||||
all_usage.update(read_usage_from_transcript(main_path))
|
||||
|
||||
for path, _run in session.subagent_transcript_paths:
|
||||
all_usage.update(read_usage_from_transcript(path))
|
||||
|
||||
patched = 0
|
||||
for message_id, run in tracker.llm_runs_by_message_id.items():
|
||||
usage = all_usage.get(message_id)
|
||||
if usage:
|
||||
meta = run.extra.setdefault("metadata", {})
|
||||
meta["usage_metadata"] = usage
|
||||
patched += 1
|
||||
|
||||
if patched:
|
||||
logger.debug(f"Set usage on {patched} LLM run(s) from transcripts")
|
||||
|
||||
|
||||
# ── Helpers ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _parse_timestamp(value: Optional[str]) -> Optional[datetime]:
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
return datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
@@ -0,0 +1,248 @@
|
||||
"""Token usage utilities for Claude Agent SDK.
|
||||
|
||||
Normalizes raw Anthropic usage dicts into the canonical ``usage_metadata``
|
||||
format expected by LangSmith. The key Anthropic-specific behavior is that
|
||||
cache tokens (``cache_read_input_tokens`` and ``cache_creation_input_tokens``)
|
||||
are **additive** — they are *not* included in the raw ``input_tokens`` value,
|
||||
so they must be summed in.
|
||||
|
||||
The canonical shape matches the JS LangSmith SDK's ``createUsageMetadata``:
|
||||
|
||||
.. code-block:: json
|
||||
|
||||
{
|
||||
"input_tokens": 21400,
|
||||
"output_tokens": 7,
|
||||
"total_tokens": 21407,
|
||||
"input_token_details": {
|
||||
"cache_read": 21375,
|
||||
"ephemeral_5m_input_tokens": 0,
|
||||
"ephemeral_1hr_input_tokens": 0
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _to_int(value: Any) -> int:
|
||||
try:
|
||||
return int(value)
|
||||
except (ValueError, TypeError):
|
||||
return 0
|
||||
|
||||
|
||||
def extract_usage_metadata(usage: Any) -> dict[str, Any]:
|
||||
"""Normalize a raw Anthropic usage dict into canonical ``usage_metadata``.
|
||||
|
||||
Anthropic cache tokens are **additive**: ``cache_read_input_tokens`` and
|
||||
``cache_creation_input_tokens`` are not included in the raw
|
||||
``input_tokens``, so we sum them in to get the true input total.
|
||||
"""
|
||||
if not usage:
|
||||
return {}
|
||||
|
||||
get = (
|
||||
usage.get if isinstance(usage, dict) else lambda k, d=None: getattr(usage, k, d)
|
||||
)
|
||||
|
||||
raw_input = _to_int(get("input_tokens"))
|
||||
output_tokens = _to_int(get("output_tokens"))
|
||||
|
||||
# Build input_token_details from cache fields
|
||||
input_token_details: dict[str, int] = {}
|
||||
|
||||
cache_read = _to_int(get("cache_read_input_tokens"))
|
||||
if cache_read:
|
||||
input_token_details["cache_read"] = cache_read
|
||||
|
||||
# Structured cache_creation (with ephemeral breakdown) takes precedence
|
||||
# over the flat cache_creation_input_tokens field.
|
||||
cache_creation = get("cache_creation")
|
||||
if isinstance(cache_creation, dict):
|
||||
eph_5m = _to_int(cache_creation.get("ephemeral_5m_input_tokens"))
|
||||
eph_1h = _to_int(cache_creation.get("ephemeral_1h_input_tokens"))
|
||||
if eph_5m:
|
||||
input_token_details["ephemeral_5m_input_tokens"] = eph_5m
|
||||
if eph_1h:
|
||||
input_token_details["ephemeral_1hr_input_tokens"] = eph_1h
|
||||
else:
|
||||
# Flat/legacy field — assume 5-minute cache
|
||||
flat_cache_create = _to_int(get("cache_creation_input_tokens"))
|
||||
if flat_cache_create:
|
||||
input_token_details["ephemeral_5m_input_tokens"] = flat_cache_create
|
||||
|
||||
# Sum cache tokens into input_tokens (Anthropic cache tokens are additive)
|
||||
cache_token_sum = sum(input_token_details.values())
|
||||
adjusted_input = raw_input + cache_token_sum
|
||||
total_tokens = adjusted_input + output_tokens
|
||||
|
||||
meta: dict[str, Any] = {
|
||||
"input_tokens": adjusted_input,
|
||||
"output_tokens": output_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
if input_token_details:
|
||||
meta["input_token_details"] = input_token_details
|
||||
|
||||
return meta
|
||||
|
||||
|
||||
def read_usage_from_transcript(
|
||||
file_path: str,
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
"""Read a JSONL transcript and return final usage per message_id.
|
||||
|
||||
The Claude SDK streams assistant messages as multiple JSONL chunks
|
||||
with the same ``message.id``. Only the final chunk (where
|
||||
``stop_reason`` is set) has accurate ``output_tokens``.
|
||||
|
||||
Returns:
|
||||
``{message_id: usage_metadata}`` with canonical usage dicts.
|
||||
"""
|
||||
try:
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
return {}
|
||||
|
||||
# Collect the last usage seen per message_id — the final chunk
|
||||
# (with stop_reason set) overwrites earlier partials.
|
||||
raw_usage: dict[str, dict[str, Any]] = {}
|
||||
with open(path) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
data = json.loads(line)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
if data.get("type") != "assistant":
|
||||
continue
|
||||
msg = data.get("message", {})
|
||||
msg_id = msg.get("id")
|
||||
usage = msg.get("usage")
|
||||
if not msg_id or not usage:
|
||||
continue
|
||||
# Always overwrite — later chunks have better counts.
|
||||
# The final chunk (with stop_reason) is last.
|
||||
raw_usage[msg_id] = usage
|
||||
|
||||
return {mid: extract_usage_metadata(u) for mid, u in raw_usage.items() if u}
|
||||
except OSError as e:
|
||||
logger.debug(f"Could not read transcript {file_path}: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def read_llm_turns_from_transcript(
|
||||
file_path: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Read final LLM turns from a JSONL transcript.
|
||||
|
||||
Each API request produces two assistant entries in the transcript
|
||||
(a streaming partial and a final completion) sharing the same
|
||||
``message.id``. Only the **final** entry (with ``stop_reason`` set)
|
||||
is returned.
|
||||
|
||||
Returns a list of dicts ordered by appearance, each containing::
|
||||
|
||||
{
|
||||
"message_id": str,
|
||||
"model": str,
|
||||
"content": list[dict], # Anthropic content blocks
|
||||
"stop_reason": str, # "end_turn" or "tool_use"
|
||||
"usage": dict, # raw Anthropic usage dict
|
||||
"timestamp": str | None, # ISO 8601 timestamp
|
||||
"input_messages": list[dict], # preceding conversation messages
|
||||
}
|
||||
"""
|
||||
try:
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
return []
|
||||
|
||||
# Single pass: build a running conversation history so each
|
||||
# assistant turn gets the full message context the API saw.
|
||||
# Tool result blocks are formatted as role:"tool" messages.
|
||||
entries_by_id: dict[str, dict[str, Any]] = {}
|
||||
conversation: list[dict[str, Any]] = []
|
||||
with open(path) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
data = json.loads(line)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
|
||||
entry_type = data.get("type")
|
||||
|
||||
if entry_type == "user":
|
||||
msg = data.get("message", {})
|
||||
content = msg.get("content")
|
||||
if content is None:
|
||||
continue
|
||||
# Detect tool_result blocks and format as tool
|
||||
# messages, matching the live stream formatting.
|
||||
if isinstance(content, list) and content:
|
||||
is_tool_result = any(
|
||||
isinstance(b, dict) and b.get("type") == "tool_result"
|
||||
for b in content
|
||||
)
|
||||
if is_tool_result:
|
||||
for block in content:
|
||||
if (
|
||||
isinstance(block, dict)
|
||||
and block.get("type") == "tool_result"
|
||||
):
|
||||
conversation.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"content": block.get("content", ""),
|
||||
"tool_call_id": block.get("tool_use_id"),
|
||||
}
|
||||
)
|
||||
continue
|
||||
conversation.append({"role": "user", "content": content})
|
||||
continue
|
||||
|
||||
if entry_type != "assistant":
|
||||
continue
|
||||
|
||||
msg = data.get("message", {})
|
||||
msg_id = msg.get("id")
|
||||
if not msg_id:
|
||||
continue
|
||||
# Always overwrite — the final entry (with stop_reason)
|
||||
# comes last for each message_id.
|
||||
entries_by_id[msg_id] = {
|
||||
"message_id": msg_id,
|
||||
"model": msg.get("model"),
|
||||
"content": msg.get("content", []),
|
||||
"stop_reason": msg.get("stop_reason"),
|
||||
"usage": msg.get("usage"),
|
||||
"timestamp": data.get("timestamp"),
|
||||
# Full conversation history up to this turn.
|
||||
"input_messages": list(conversation),
|
||||
}
|
||||
# Add completed assistant content to conversation
|
||||
# so subsequent turns see it.
|
||||
if msg.get("stop_reason"):
|
||||
conversation.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": msg.get("content", []),
|
||||
}
|
||||
)
|
||||
|
||||
# Only return final entries (stop_reason is set).
|
||||
return [e for e in entries_by_id.values() if e.get("stop_reason")]
|
||||
except OSError as e:
|
||||
logger.debug(f"Could not read transcript {file_path}: {e}")
|
||||
return []
|
||||
Reference in New Issue
Block a user