pydftracer 2.0.5.post2.dev0__tar.gz → 2.0.6.post2.dev0__tar.gz

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydftracer
-Version: 2.0.5.post2.dev0
+Version: 2.0.6.post2.dev0
 Summary: Python libs for DFTracer
 Author-email: "Hariharan Devarajan (Hari)" <hariharandev1@llnl.gov>, Ray Andrew Sinurat <raydreww@gmail.com>
 Maintainer-email: "Hariharan Devarajan (Hari)" <hariharandev1@llnl.gov>, Ray Andrew Sinurat <raydreww@gmail.com>

pydftracer-2.0.6.post2.dev0/docs/source/api/agent.rst

@@ -0,0 +1,109 @@
+Agent API
+=========
+
+This section documents the agent tracing namespace and its event model.
+
+Agent Module
+------------
+
+The ``agent`` module provides decorators and utilities for tracing agentic workflows.
+
+.. automodule:: dftracer.python.agent
+   :members:
+   :undoc-members:
+
+Agent Categories
+----------------
+
+.. autoclass:: dftracer.python.AgentCategory
+   :members:
+   :undoc-members:
+
+Agent Components
+----------------
+
+.. autoclass:: dftracer.python.Workflow
+   :members:
+   :undoc-members:
+
+.. autoclass:: dftracer.python.AgentStep
+   :members:
+   :undoc-members:
+
+.. autoclass:: dftracer.python.LLM
+   :members:
+   :undoc-members:
+
+.. autoclass:: dftracer.python.Tool
+   :members:
+   :undoc-members:
+
+.. autoclass:: dftracer.python.AgentData
+   :members:
+   :undoc-members:
+
+.. autoclass:: dftracer.python.Message
+   :members:
+   :undoc-members:
+
+.. autoclass:: dftracer.python.Judge
+   :members:
+   :undoc-members:
+
+Agent Instance
+--------------
+
+.. autoclass:: dftracer.python.Agent
+   :members:
+   :undoc-members:
+
+Convenience aliases:
+
+- ``agent`` - Main agent tracing instance
+- ``workflow`` - Workflow tracing category
+- ``agent_step`` - Agent step tracing category
+- ``llm`` - LLM tracing category
+- ``tool`` - Tool tracing category
+- ``agent_data`` - Data operation tracing category
+- ``message`` - Message/event tracing category
+- ``judge`` - Evaluation/retry tracing category
+
+Span Arguments and Aggregation
+------------------------------
+
+Spans carry their ``args`` (``trace_id``, ``step``, token counts, and so on) in the
+trace, and ``span.update_args({...})`` adds values after the work finishes, such as
+token usage from an LLM response.
+
+DFTracer can aggregate short events into counters to keep traces small. An
+aggregated event no longer carries its ``args``, so with the default aggregation a
+short ``llm.call`` or ``tool.call`` loses its arguments, token counts included. To
+keep agent spans intact, either turn aggregation off:
+
+.. code-block:: bash
+
+   DFTRACER_ENABLE_AGGREGATION=0
+
+or aggregate only the I/O events, with the ``SELECTIVE`` type and a rules file that
+includes only POSIX and STDIO:
+
+.. code-block:: bash
+
+   DFTRACER_ENABLE_AGGREGATION=1
+   DFTRACER_AGGREGATION_TYPE=SELECTIVE
+   DFTRACER_AGGREGATION_FILE=aggregation_rules.yaml
+
+.. code-block:: yaml
+
+   inclusion:
+     - cat IN {POSIX,STDIO}
+   exclusion:
+
+LangGraph Adapter
+-----------------
+
+Utilities for tracing LangGraph-style node and workflow invocation boundaries:
+
+.. autofunction:: dftracer.python.get_or_init_trace_context
+.. autofunction:: dftracer.python.traced_node
+.. autofunction:: dftracer.python.trace_invoke

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/docs/source/api/index.rst

@@ -8,6 +8,7 @@ This section contains the complete API reference for pydftracer.
 
    core
    ai
+   agent
    dynamo
    env
 
@@ -35,6 +36,7 @@ AI/ML Support
 ~~~~~~~~~~~~~
 
 - ``DFTracerAI`` - Base class for AI/ML tracing
+- ``Agent`` - Base class for agent/workflow tracing
 - ``Dynamo`` - PyTorch Dynamo integration class
 - ``dynamo`` - Pre-configured Dynamo instance for tracing
 - ``create_backend`` - Create a custom torch.compile backend with DFTracer
@@ -58,5 +60,6 @@ Modules
 
 - :doc:`core` - Core tracing functionality (logger, common utilities)
 - :doc:`ai` - AI/ML specific tracing features
+- :doc:`agent` - Agent/workflow tracing features
 - :doc:`dynamo` - PyTorch Dynamo integration
 - :doc:`env` - Environment configuration and setup

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/docs/source/index.rst

@@ -15,6 +15,7 @@ Features
 - **Type-Safe Decorators**: Fully type-checked with mypy, preserves function signatures
 - **Function Decorators**: Simple ``@dft_fn`` decorator for tracing Python functions
 - **AI/ML Support**: Specialized tracing for AI/ML workflows
+- **Agent Workflow Support**: Structured tracing for agentic workflows
 - **PyTorch Integration**: Capture PyTorch profiler events to DFTracer traces
 - **PyTorch Dynamo Integration**: Wrapper of PyTorch's Dynamo
 - **Automatic I/O Tracing**: Transparent tracing of I/O operations when enabled

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/pyproject.toml

@@ -125,7 +125,10 @@ ignore = [
 "tests/*" = ["B011", "B023"]
 "python/dftracer/python/logger.py" = ["E402", "F405"]
 "python/dftracer/python/dbg/logger.py" = ["E402", "F405"]
+"python/dftracer/python/dbg/agent.py" = ["E402"]
 "python/dftracer/python/dbg/ai.py" = ["E402"]
+"python/dftracer/python/agent.py" = ["E402"]
+"python/dftracer/python/agent_init.py" = ["F405"]
 "python/dftracer/python/ai.py" = ["E402"]
 "python/dftracer/python/ai_common.py" = ["F405"]
 "python/dftracer/python/ai_init.py" = ["F405"]

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/python/dftracer/python/__init__.py

@@ -1,5 +1,7 @@
 from importlib.metadata import PackageNotFoundError, version
 
+from dftracer.python.agent import *
+from dftracer.python.agent_langgraph import *
 from dftracer.python.ai import *
 from dftracer.python.dynamo import *
 from dftracer.python.logger import *

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/python/dftracer/python/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '2.0.5.post2.dev0'
-__version_tuple__ = version_tuple = (2, 0, 5, 'post2', 'dev0')
+__version__ = version = '2.0.6.post2.dev0'
+__version_tuple__ = version_tuple = (2, 0, 6, 'post2', 'dev0')
 
-__commit_id__ = commit_id = 'gb88b10fd4'
+__commit_id__ = commit_id = 'g68307eab6'

pydftracer-2.0.6.post2.dev0/python/dftracer/python/agent.py

@@ -0,0 +1,9 @@
+import dftracer.python.ai_common as ai_common
+from dftracer.python.logger import dft_fn as dft_fn_orig
+from dftracer.python.logger import dftracer as dftracer_orig
+
+ai_common.dft_fn = dft_fn_orig  # type: ignore[misc]
+ai_common.dftracer = dftracer_orig  # type: ignore[misc]
+
+from dftracer.python.agent_common import *
+from dftracer.python.agent_init import *

pydftracer-2.0.6.post2.dev0/python/dftracer/python/agent_common.py

@@ -0,0 +1,381 @@
+from typing import Any, Optional
+
+from dftracer.python.ai_common import DFTracerAI
+from dftracer.python.common import StringEnum, auto
+
+ROOT_NAME = "agent_root"
+ROOT_CAT = "agent_root"
+
+TRACE_ID_KEY = "trace_id"
+WORKFLOW_ID_KEY = "workflow_id"
+RUN_ID_KEY = "run_id"
+STEP_ID_KEY = "step_id"
+AGENT_ID_KEY = "agent_id"
+AGENT_ROLE_KEY = "agent_role"
+SPAN_ID_KEY = "span_id"
+PARENT_SPAN_ID_KEY = "parent_span_id"
+TOOL_CALL_ID_KEY = "tool_call_id"
+LLM_CALL_ID_KEY = "llm_call_id"
+FRAMEWORK_KEY = "framework"
+MODEL_NAME_KEY = "model_name"
+MESSAGE_TYPE_KEY = "message_type"
+LOGICAL_PATH_KEY = "logical_path"
+PHYSICAL_PATH_KEY = "physical_path"
+FORMAT_KEY = "format"
+OPERATION_KIND_KEY = "operation_kind"
+SUCCESS_KEY = "success"
+ERROR_TYPE_KEY = "error_type"
+
+# Token usage tracking keys for LLM calls
+PROMPT_TOKENS_KEY = "prompt_tokens"
+COMPLETION_TOKENS_KEY = "completion_tokens"
+TOTAL_TOKENS_KEY = "total_tokens"
+CONTEXT_SIZE_KEY = "context_size"
+CONTEXT_WINDOW_KEY = "context_window"
+CONTEXT_UTILIZATION_KEY = "context_utilization"
+
+
+class AgentCategory(StringEnum):
+    WORKFLOW = auto()
+    STEP = auto()
+    LLM = auto()
+    TOOL = auto()
+    DATA = auto()
+    MESSAGE = auto()
+    JUDGE = auto()
+
+
+class WorkflowEvent(StringEnum):
+    RUN = auto()
+    RESUME = auto()
+
+
+class AgentStepEvent(StringEnum):
+    PLAN = auto()
+    ACT = auto()
+    OBSERVE = auto()
+    REFLECT = auto()
+
+
+class LLMEvent(StringEnum):
+    CALL = auto()
+    STREAM = auto()
+    PARSE = auto()
+
+
+class ToolEvent(StringEnum):
+    CALL = auto()
+    RESULT = auto()
+
+
+class AgentDataEvent(StringEnum):
+    LOAD = auto()
+    SAVE = auto()
+
+
+class MessageEvent(StringEnum):
+    SEND = auto()
+    RECEIVE = auto()
+    STREAM = auto()
+
+
+class JudgeEvent(StringEnum):
+    EVALUATE = auto()
+    RETRY = auto()
+
+
+class Workflow(DFTracerAI):
+    run: DFTracerAI
+    resume: DFTracerAI
+
+    def __init__(
+        self,
+        epoch: Optional[int] = None,
+        step: Optional[int] = None,
+        image_idx: Optional[int] = None,
+        image_size: Optional[Any] = None,
+        enable: bool = True,
+    ):
+        super().__init__(
+            cat=AgentCategory.WORKFLOW,
+            name=AgentCategory.WORKFLOW,
+            epoch=epoch,
+            step=step,
+            image_idx=image_idx,
+            image_size=image_size,
+            enable=enable,
+        )
+        self.create_children(
+            {
+                "run": WorkflowEvent.RUN,
+                "resume": WorkflowEvent.RESUME,
+            }
+        )
+
+
+class AgentStep(DFTracerAI):
+    plan: DFTracerAI
+    act: DFTracerAI
+    observe: DFTracerAI
+    reflect: DFTracerAI
+
+    def __init__(
+        self,
+        epoch: Optional[int] = None,
+        step: Optional[int] = None,
+        image_idx: Optional[int] = None,
+        image_size: Optional[Any] = None,
+        enable: bool = True,
+    ):
+        super().__init__(
+            cat=AgentCategory.STEP,
+            name=AgentCategory.STEP,
+            epoch=epoch,
+            step=step,
+            image_idx=image_idx,
+            image_size=image_size,
+            enable=enable,
+        )
+        self.create_children(
+            {
+                "plan": AgentStepEvent.PLAN,
+                "act": AgentStepEvent.ACT,
+                "observe": AgentStepEvent.OBSERVE,
+                "reflect": AgentStepEvent.REFLECT,
+            }
+        )
+
+
+class LLM(DFTracerAI):
+    call: DFTracerAI
+    stream: DFTracerAI
+    parse: DFTracerAI
+
+    def __init__(
+        self,
+        epoch: Optional[int] = None,
+        step: Optional[int] = None,
+        image_idx: Optional[int] = None,
+        image_size: Optional[Any] = None,
+        enable: bool = True,
+    ):
+        super().__init__(
+            cat=AgentCategory.LLM,
+            name=AgentCategory.LLM,
+            epoch=epoch,
+            step=step,
+            image_idx=image_idx,
+            image_size=image_size,
+            enable=enable,
+        )
+        self.create_children(
+            {
+                "call": LLMEvent.CALL,
+                "stream": LLMEvent.STREAM,
+                "parse": LLMEvent.PARSE,
+            }
+        )
+
+
+class Tool(DFTracerAI):
+    call: DFTracerAI
+    result: DFTracerAI
+
+    def __init__(
+        self,
+        epoch: Optional[int] = None,
+        step: Optional[int] = None,
+        image_idx: Optional[int] = None,
+        image_size: Optional[Any] = None,
+        enable: bool = True,
+    ):
+        super().__init__(
+            cat=AgentCategory.TOOL,
+            name=AgentCategory.TOOL,
+            epoch=epoch,
+            step=step,
+            image_idx=image_idx,
+            image_size=image_size,
+            enable=enable,
+        )
+        self.create_children(
+            {
+                "call": ToolEvent.CALL,
+                "result": ToolEvent.RESULT,
+            }
+        )
+
+
+class AgentData(DFTracerAI):
+    load: DFTracerAI
+    save: DFTracerAI
+
+    def __init__(
+        self,
+        epoch: Optional[int] = None,
+        step: Optional[int] = None,
+        image_idx: Optional[int] = None,
+        image_size: Optional[Any] = None,
+        enable: bool = True,
+    ):
+        super().__init__(
+            cat=AgentCategory.DATA,
+            name=AgentCategory.DATA,
+            epoch=epoch,
+            step=step,
+            image_idx=image_idx,
+            image_size=image_size,
+            enable=enable,
+        )
+        self.create_children(
+            {
+                "load": AgentDataEvent.LOAD,
+                "save": AgentDataEvent.SAVE,
+            }
+        )
+
+
+class Message(DFTracerAI):
+    send: DFTracerAI
+    receive: DFTracerAI
+    stream: DFTracerAI
+
+    def __init__(
+        self,
+        epoch: Optional[int] = None,
+        step: Optional[int] = None,
+        image_idx: Optional[int] = None,
+        image_size: Optional[Any] = None,
+        enable: bool = True,
+    ):
+        super().__init__(
+            cat=AgentCategory.MESSAGE,
+            name=AgentCategory.MESSAGE,
+            epoch=epoch,
+            step=step,
+            image_idx=image_idx,
+            image_size=image_size,
+            enable=enable,
+        )
+        self.create_children(
+            {
+                "send": MessageEvent.SEND,
+                "receive": MessageEvent.RECEIVE,
+                "stream": MessageEvent.STREAM,
+            }
+        )
+
+
+class Judge(DFTracerAI):
+    evaluate: DFTracerAI
+    retry: DFTracerAI
+
+    def __init__(
+        self,
+        epoch: Optional[int] = None,
+        step: Optional[int] = None,
+        image_idx: Optional[int] = None,
+        image_size: Optional[Any] = None,
+        enable: bool = True,
+    ):
+        super().__init__(
+            cat=AgentCategory.JUDGE,
+            name=AgentCategory.JUDGE,
+            epoch=epoch,
+            step=step,
+            image_idx=image_idx,
+            image_size=image_size,
+            enable=enable,
+        )
+        self.create_children(
+            {
+                "evaluate": JudgeEvent.EVALUATE,
+                "retry": JudgeEvent.RETRY,
+            }
+        )
+
+
+# fmt: off
+class Agent(DFTracerAI):
+    workflow: Workflow
+    step: AgentStep
+    llm: LLM
+    tool: Tool
+    data: AgentData
+    message: Message
+    judge: Judge
+
+    def __init__(
+        self,
+        epoch: Optional[int] = None,
+        step: Optional[int] = None,
+        image_idx: Optional[int] = None,
+        image_size: Optional[Any] = None,
+        enable: bool = True,
+    ):
+        super().__init__(cat=ROOT_CAT, name=ROOT_NAME, epoch=epoch, step=step, image_idx=image_idx, image_size=image_size, enable=enable)
+        self.workflow = Workflow(epoch=epoch, step=step, image_idx=image_idx, image_size=image_size, enable=enable)
+        self.step = AgentStep(epoch=epoch, step=step, image_idx=image_idx, image_size=image_size, enable=enable)
+        self.llm = LLM(epoch=epoch, step=step, image_idx=image_idx, image_size=image_size, enable=enable)
+        self.tool = Tool(epoch=epoch, step=step, image_idx=image_idx, image_size=image_size, enable=enable)
+        self.data = AgentData(epoch=epoch, step=step, image_idx=image_idx, image_size=image_size, enable=enable)
+        self.message = Message(epoch=epoch, step=step, image_idx=image_idx, image_size=image_size, enable=enable)
+        self.judge = Judge(epoch=epoch, step=step, image_idx=image_idx, image_size=image_size, enable=enable)
+
+        self._children = {
+            "workflow": self.workflow,
+            "step": self.step,
+            "llm": self.llm,
+            "tool": self.tool,
+            "data": self.data,
+            "message": self.message,
+            "judge": self.judge,
+        }
+# fmt: on
+
+
+__all__ = [
+    "ROOT_NAME",
+    "ROOT_CAT",
+    "TRACE_ID_KEY",
+    "WORKFLOW_ID_KEY",
+    "RUN_ID_KEY",
+    "STEP_ID_KEY",
+    "AGENT_ID_KEY",
+    "AGENT_ROLE_KEY",
+    "SPAN_ID_KEY",
+    "PARENT_SPAN_ID_KEY",
+    "TOOL_CALL_ID_KEY",
+    "LLM_CALL_ID_KEY",
+    "FRAMEWORK_KEY",
+    "MODEL_NAME_KEY",
+    "MESSAGE_TYPE_KEY",
+    "LOGICAL_PATH_KEY",
+    "PHYSICAL_PATH_KEY",
+    "FORMAT_KEY",
+    "OPERATION_KIND_KEY",
+    "SUCCESS_KEY",
+    "ERROR_TYPE_KEY",
+    "PROMPT_TOKENS_KEY",
+    "COMPLETION_TOKENS_KEY",
+    "TOTAL_TOKENS_KEY",
+    "CONTEXT_SIZE_KEY",
+    "CONTEXT_WINDOW_KEY",
+    "CONTEXT_UTILIZATION_KEY",
+    "AgentCategory",
+    "WorkflowEvent",
+    "AgentStepEvent",
+    "LLMEvent",
+    "ToolEvent",
+    "AgentDataEvent",
+    "MessageEvent",
+    "JudgeEvent",
+    "Workflow",
+    "AgentStep",
+    "LLM",
+    "Tool",
+    "AgentData",
+    "Message",
+    "Judge",
+    "Agent",
+]

pydftracer-2.0.6.post2.dev0/python/dftracer/python/agent_init.py

@@ -0,0 +1,56 @@
+from dftracer.python.agent_common import *
+
+agent = Agent()
+workflow = agent.workflow
+agent_step = agent.step
+llm = agent.llm
+tool = agent.tool
+agent_data = agent.data
+message = agent.message
+judge = agent.judge
+
+__all__ = [
+    "ROOT_NAME",
+    "ROOT_CAT",
+    "TRACE_ID_KEY",
+    "WORKFLOW_ID_KEY",
+    "RUN_ID_KEY",
+    "STEP_ID_KEY",
+    "AGENT_ID_KEY",
+    "AGENT_ROLE_KEY",
+    "SPAN_ID_KEY",
+    "PARENT_SPAN_ID_KEY",
+    "TOOL_CALL_ID_KEY",
+    "LLM_CALL_ID_KEY",
+    "FRAMEWORK_KEY",
+    "MODEL_NAME_KEY",
+    "MESSAGE_TYPE_KEY",
+    "LOGICAL_PATH_KEY",
+    "PHYSICAL_PATH_KEY",
+    "FORMAT_KEY",
+    "OPERATION_KIND_KEY",
+    "SUCCESS_KEY",
+    "ERROR_TYPE_KEY",
+    "PROMPT_TOKENS_KEY",
+    "COMPLETION_TOKENS_KEY",
+    "TOTAL_TOKENS_KEY",
+    "CONTEXT_SIZE_KEY",
+    "CONTEXT_WINDOW_KEY",
+    "CONTEXT_UTILIZATION_KEY",
+    "AgentCategory",
+    "WorkflowEvent",
+    "AgentStepEvent",
+    "LLMEvent",
+    "ToolEvent",
+    "AgentDataEvent",
+    "MessageEvent",
+    "JudgeEvent",
+    "agent",
+    "workflow",
+    "agent_step",
+    "llm",
+    "tool",
+    "agent_data",
+    "message",
+    "judge",
+]

pydftracer-2.0.6.post2.dev0/python/dftracer/python/agent_langgraph.py

@@ -0,0 +1,211 @@
+import uuid
+from typing import Any, Callable, Dict, Mapping, MutableMapping, Optional
+
+from dftracer.python.agent_common import (
+    AGENT_ID_KEY,
+    AGENT_ROLE_KEY,
+    FRAMEWORK_KEY,
+    OPERATION_KIND_KEY,
+    PARENT_SPAN_ID_KEY,
+    RUN_ID_KEY,
+    SPAN_ID_KEY,
+    STEP_ID_KEY,
+    TRACE_ID_KEY,
+    WORKFLOW_ID_KEY,
+)
+from dftracer.python.agent_init import agent_step, workflow
+from dftracer.python.ai_common import DFTracerAI
+from dftracer.python.common import P, R
+
+DEFAULT_FRAMEWORK_NAME = "langgraph"
+DEFAULT_WORKFLOW_EVENT_NAME = "langgraph.invoke"
+DEFAULT_NODE_OPERATION_KIND = "node"
+DEFAULT_WORKFLOW_OPERATION_KIND = "workflow"
+
+
+def _new_id() -> str:
+    return uuid.uuid4().hex
+
+
+def _extract_thread_id(config: Optional[Mapping[str, Any]]) -> Optional[str]:
+    if not config:
+        return None
+    configurable = config.get("configurable")
+    if not isinstance(configurable, Mapping):
+        return None
+    thread_id = configurable.get("thread_id")
+    if thread_id is None:
+        return None
+    return str(thread_id)
+
+
+def _prune_none_values(args: Dict[str, Any]) -> Dict[str, Any]:
+    return {k: v for k, v in args.items() if v is not None}
+
+
+def _get_step_event_tracer(step_event: str) -> DFTracerAI:
+    tracer: Optional[DFTracerAI] = getattr(agent_step, step_event, None)
+    if tracer is None:
+        valid_events = ", ".join(sorted(agent_step._children.keys()))
+        raise ValueError(
+            f"Unknown step_event '{step_event}'. Valid values: {valid_events}"
+        )
+    return tracer
+
+
+def get_or_init_trace_context(
+    state: Optional[Mapping[str, Any]] = None,
+    config: Optional[Mapping[str, Any]] = None,
+) -> Dict[str, Optional[str]]:
+    state = state or {}
+
+    trace_id = state.get(TRACE_ID_KEY)
+    if trace_id is not None:
+        trace_id = str(trace_id)
+    if not trace_id:
+        trace_id = _new_id()
+
+    workflow_id = state.get(WORKFLOW_ID_KEY)
+    if workflow_id is not None:
+        workflow_id = str(workflow_id)
+    if not workflow_id:
+        workflow_id = _extract_thread_id(config) or _new_id()
+
+    parent_span_id = state.get(PARENT_SPAN_ID_KEY)
+    if parent_span_id is not None:
+        parent_span_id = str(parent_span_id)
+
+    return {
+        TRACE_ID_KEY: trace_id,
+        WORKFLOW_ID_KEY: workflow_id,
+        PARENT_SPAN_ID_KEY: parent_span_id,
+    }
+
+
+def _merge_trace_updates(
+    updates: MutableMapping[str, Any], trace_id: str, workflow_id: str, span_id: str
+) -> MutableMapping[str, Any]:
+    updates.setdefault(TRACE_ID_KEY, trace_id)
+    updates.setdefault(WORKFLOW_ID_KEY, workflow_id)
+    updates.setdefault(PARENT_SPAN_ID_KEY, span_id)
+    return updates
+
+
+def traced_node(
+    node_name: str,
+    *,
+    agent_role: str = "",
+    step_event: str = "act",
+    framework: str = DEFAULT_FRAMEWORK_NAME,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    tracer = _get_step_event_tracer(step_event)
+
+    def _decorator(fn: Callable[P, R]) -> Callable[P, R]:
+        def _wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
+            state: Optional[Mapping[str, Any]] = None
+            kw_state = kwargs.get("state")
+            if len(args) > 0 and isinstance(args[0], Mapping):
+                state = args[0]
+            elif isinstance(kw_state, Mapping):
+                state = kw_state
+
+            config: Optional[Mapping[str, Any]] = None
+            kw_config = kwargs.get("config")
+            if len(args) > 1 and isinstance(args[1], Mapping):
+                config = args[1]
+            elif isinstance(kw_config, Mapping):
+                config = kw_config
+
+            ctx = get_or_init_trace_context(state=state, config=config)
+            span_id = _new_id()
+            event_args = _prune_none_values(
+                {
+                    TRACE_ID_KEY: ctx[TRACE_ID_KEY],
+                    WORKFLOW_ID_KEY: ctx[WORKFLOW_ID_KEY],
+                    SPAN_ID_KEY: span_id,
+                    STEP_ID_KEY: span_id,
+                    PARENT_SPAN_ID_KEY: ctx[PARENT_SPAN_ID_KEY],
+                    AGENT_ID_KEY: node_name,
+                    AGENT_ROLE_KEY: agent_role or None,
+                    FRAMEWORK_KEY: framework,
+                    OPERATION_KIND_KEY: DEFAULT_NODE_OPERATION_KIND,
+                }
+            )
+
+            with tracer(args=event_args):
+                result = fn(*args, **kwargs)
+
+            if result is None:
+                updates: MutableMapping[str, Any] = {}
+                merged = _merge_trace_updates(
+                    updates=updates,
+                    trace_id=str(ctx[TRACE_ID_KEY]),
+                    workflow_id=str(ctx[WORKFLOW_ID_KEY]),
+                    span_id=span_id,
+                )
+                return merged  # type: ignore[return-value]
+
+            if isinstance(result, Mapping):
+                updates = dict(result)
+                merged = _merge_trace_updates(
+                    updates=updates,
+                    trace_id=str(ctx[TRACE_ID_KEY]),
+                    workflow_id=str(ctx[WORKFLOW_ID_KEY]),
+                    span_id=span_id,
+                )
+                return merged  # type: ignore[return-value]
+
+            return result
+
+        return _wrapped
+
+    return _decorator
+
+
+def trace_invoke(
+    app: Any,
+    state: Any,
+    config: Optional[Mapping[str, Any]] = None,
+    *,
+    workflow_name: str = DEFAULT_WORKFLOW_EVENT_NAME,
+    framework: str = DEFAULT_FRAMEWORK_NAME,
+    run_id: Optional[str] = None,
+) -> Any:
+    if not hasattr(app, "invoke"):
+        raise AttributeError("Expected app to have an 'invoke' method")
+
+    state_mapping = state if isinstance(state, Mapping) else None
+    ctx = get_or_init_trace_context(state=state_mapping, config=config)
+    run_id = run_id or _new_id()
+    span_id = _new_id()
+
+    event_args = _prune_none_values(
+        {
+            TRACE_ID_KEY: ctx[TRACE_ID_KEY],
+            WORKFLOW_ID_KEY: ctx[WORKFLOW_ID_KEY],
+            RUN_ID_KEY: run_id,
+            SPAN_ID_KEY: span_id,
+            PARENT_SPAN_ID_KEY: ctx[PARENT_SPAN_ID_KEY],
+            FRAMEWORK_KEY: framework,
+            OPERATION_KIND_KEY: DEFAULT_WORKFLOW_OPERATION_KIND,
+        }
+    )
+
+    workflow_tracer = (
+        workflow.run.derive(workflow_name) if workflow_name else workflow.run
+    )
+    with workflow_tracer(args=event_args):
+        if config is None:
+            return app.invoke(state)
+        return app.invoke(state, config=config)
+
+
+__all__ = [
+    "DEFAULT_FRAMEWORK_NAME",
+    "DEFAULT_WORKFLOW_EVENT_NAME",
+    "DEFAULT_NODE_OPERATION_KIND",
+    "DEFAULT_WORKFLOW_OPERATION_KIND",
+    "get_or_init_trace_context",
+    "traced_node",
+    "trace_invoke",
+]

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/python/dftracer/python/ai_common.py

@@ -1,6 +1,5 @@
 import functools
-import sys
-from typing import Any, Callable, Dict, Iterator, List, Optional, overload
+from typing import Any, Callable, Dict, Iterator, Optional, overload
 
 from dftracer.python.common import *
 from dftracer.python.env import *
@@ -8,34 +7,6 @@ from dftracer.python.env import *
 dft_fn = None  # type: ignore
 dftracer = None  # type: ignore
 
-if sys.version_info >= (3, 11):
-    from enum import StrEnum as StringEnum
-    from enum import auto
-    # StrEnum in 3.11 already lower-case
-else:
-    from enum import Enum, auto
-
-    # MIT License: https://github.com/irgeek/StrEnum/blob/master/strenum/__init__.py
-    class StringEnum(str, Enum):
-        # this is a specific version that will lowercase the values
-
-        def __new__(cls, value: Any, *args: Any, **kwargs: Any) -> "StringEnum":
-            if not isinstance(value, (str, auto)):
-                raise TypeError(  # pragma: no cover
-                    f"Values of StrEnums must be strings: {value!r} is a {type(value)}"
-                )
-            return super().__new__(cls, value, *args, **kwargs)
-
-        def __str__(self) -> str:
-            return str(self.value)
-
-        @staticmethod
-        def _generate_next_value_(
-            name: str, start: int, count: int, last_values: List[Any]
-        ) -> str:
-            return name.lower()
-
-
 ITER_COUNT_NAME = "count"
 INIT_NAME = "init"
 BLOCK_NAME = "block"
@@ -123,44 +94,16 @@ class _DFTracerAI:
         image_size: Optional[Any] = None,
         args: Optional[Dict[str, Any]] = None,
     ) -> Union[Callable[P, R], "DFTracerAI"]:
-        if epoch is not None:
-            self.profiler._arguments_int["epoch"] = TagValue(
-                epoch, TagDType.INT, TagType.KEY
-            ).value()
-        if step is not None:
-            self.profiler._arguments_int["step"] = TagValue(
-                step, TagDType.INT, TagType.KEY
-            ).value()
-        if image_idx is not None:
-            self.profiler._arguments_int["image_idx"] = TagValue(
-                image_idx, TagDType.INT, TagType.KEY
-            ).value()
-        if image_size is not None:
-            self.profiler._arguments_float["image_size"] = TagValue(
-                image_size, TagDType.FLOAT, TagType.KEY
-            ).value()
-        args = args or {}
-        for key, value in args.items():
-            if isinstance(value, TagValue):
-                new_value = value._value
-            else:
-                new_value = value
-            if isinstance(new_value, int):
-                self.profiler._arguments_int[key] = TagValue(
-                    new_value, TagDType.INT, TagType.KEY
-                ).value()
-            elif isinstance(new_value, float):
-                self.profiler._arguments_float[key] = TagValue(
-                    new_value, TagDType.FLOAT, TagType.KEY
-                ).value()
-            else:
-                self.profiler._arguments_string[key] = TagValue(
-                    str(new_value), TagDType.STRING, TagType.KEY
-                ).value()
-
         is_enabled = self.profiler._enable if enable is None else enable
 
         if fn:
+            self.update(
+                epoch=epoch,
+                step=step,
+                image_idx=image_idx,
+                image_size=image_size,
+                args=args,
+            )
 
             def _decorator(f: Callable[P, R]) -> Callable[P, R]:
                 @functools.wraps(f)
@@ -174,15 +117,25 @@ class _DFTracerAI:
 
             return _decorator(fn)
         else:
-            return DFTracerAI(
+            tracer = DFTracerAI(
                 cat=self.profiler._cat,
                 name=self.profiler._name,
+                enable=is_enabled,
+            )
+            if DFTRACER_ENABLE and tracer.profiler._enable:
+                tracer.profiler._arguments_int = dict(self.profiler._arguments_int)
+                tracer.profiler._arguments_float = dict(self.profiler._arguments_float)
+                tracer.profiler._arguments_string = dict(
+                    self.profiler._arguments_string
+                )
+            tracer.update(
                 epoch=epoch,
                 step=step,
                 image_idx=image_idx,
                 image_size=image_size,
-                enable=is_enabled,
+                args=args,
             )
+            return tracer
 
     def __enter__(self) -> "_DFTracerAI":
         self.profiler.__enter__()
@@ -285,6 +238,46 @@ class _DFTracerAI:
                     ).value()
         return self
 
+    def update_args(self, args: Dict[str, Any]) -> "_DFTracerAI":
+        """Update event arguments during span execution (e.g., token usage from LLM response).
+
+        This method allows adding or updating args after __enter__ but before __exit__,
+        which is useful for capturing metrics from API responses (like token counts).
+
+        Args:
+            args: Dictionary of argument names to values. Int, float, and str types are supported.
+
+        Returns:
+            self for chaining.
+
+        Example:
+            with llm.call(args={"model": "gpt-4"}) as span:
+                response = api.call()
+                span.update_args({
+                    "prompt_tokens": response["usage"]["prompt_tokens"],
+                    "completion_tokens": response["usage"]["completion_tokens"],
+                })
+        """
+        if DFTRACER_ENABLE and self.profiler._enable:
+            for key, value in args.items():
+                if isinstance(value, TagValue):
+                    new_value = value._value
+                else:
+                    new_value = value
+                if isinstance(new_value, int):
+                    self.profiler._arguments_int[key] = TagValue(
+                        new_value, TagDType.INT, TagType.KEY
+                    ).value()
+                elif isinstance(new_value, float):
+                    self.profiler._arguments_float[key] = TagValue(
+                        new_value, TagDType.FLOAT, TagType.KEY
+                    ).value()
+                else:
+                    self.profiler._arguments_string[key] = TagValue(
+                        str(new_value), TagDType.STRING, TagType.KEY
+                    ).value()
+        return self
+
     def iter(
         self,
         iterator: Iterator,

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/python/dftracer/python/common.py

@@ -4,7 +4,7 @@ import os
 import signal
 import sys
 from collections.abc import Iterable, Sequence
-from enum import Enum, IntEnum
+from enum import Enum, IntEnum, auto
 from functools import wraps
 from pathlib import Path
 from typing import (
@@ -29,6 +29,36 @@ P = ParamSpec("P")  # For capturing function parameters
 R = TypeVar("R")  # For capturing function return types
 T = TypeVar("T")  # For generic iterator types
 
+
+# StrEnum for Python < 3.11, shared by the ai and agent namespaces.
+if sys.version_info >= (3, 11):
+    from enum import StrEnum as StringEnum
+    from enum import auto
+    # StrEnum in 3.11 already lower-case
+else:
+    from enum import Enum, auto
+
+    # MIT License: https://github.com/irgeek/StrEnum/blob/master/strenum/__init__.py
+    class StringEnum(str, Enum):
+        # this is a specific version that will lowercase the values
+
+        def __new__(cls, value: Any, *args: Any, **kwargs: Any) -> "StringEnum":
+            if not isinstance(value, (str, auto)):
+                raise TypeError(  # pragma: no cover
+                    f"Values of StrEnums must be strings: {value!r} is a {type(value)}"
+                )
+            return super().__new__(cls, value, *args, **kwargs)
+
+        def __str__(self) -> str:
+            return str(self.value)
+
+        @staticmethod
+        def _generate_next_value_(
+            name: str, start: int, count: int, last_values: List[Any]
+        ) -> str:
+            return name.lower()
+
+
 # Type alias for TagValue return type
 TagValueTuple = Tuple[int, Union[int, float, str]]
 

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/python/dftracer/python/dbg/__init__.py

@@ -1,3 +1,4 @@
+from dftracer.python.dbg.agent import *
 from dftracer.python.dbg.ai import *
 from dftracer.python.dbg.logger import *
 from dftracer.python.dynamo import *

pydftracer-2.0.6.post2.dev0/python/dftracer/python/dbg/agent.py

@@ -0,0 +1,9 @@
+import dftracer.python.ai_common as ai_common
+from dftracer.python.dbg.logger import dft_fn as dft_fn_orig
+from dftracer.python.dbg.logger import dftracer as dftracer_orig
+
+ai_common.dft_fn = dft_fn_orig  # type: ignore[misc]
+ai_common.dftracer = dftracer_orig  # type: ignore[misc]
+
+from dftracer.python.agent_common import *
+from dftracer.python.agent_init import *

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/python/pydftracer.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydftracer
-Version: 2.0.5.post2.dev0
+Version: 2.0.6.post2.dev0
 Summary: Python libs for DFTracer
 Author-email: "Hariharan Devarajan (Hari)" <hariharandev1@llnl.gov>, Ray Andrew Sinurat <raydreww@gmail.com>
 Maintainer-email: "Hariharan Devarajan (Hari)" <hariharandev1@llnl.gov>, Ray Andrew Sinurat <raydreww@gmail.com>

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/python/pydftracer.egg-info/SOURCES.txt

@@ -19,6 +19,7 @@ docs/source/installation.rst
 docs/source/pytorch.rst
 docs/source/quickstart.rst
 docs/source/type_safety.rst
+docs/source/api/agent.rst
 docs/source/api/ai.rst
 docs/source/api/core.rst
 docs/source/api/dynamo.rst
@@ -26,6 +27,10 @@ docs/source/api/env.rst
 docs/source/api/index.rst
 python/dftracer/python/__init__.py
 python/dftracer/python/_version.py
+python/dftracer/python/agent.py
+python/dftracer/python/agent_common.py
+python/dftracer/python/agent_init.py
+python/dftracer/python/agent_langgraph.py
 python/dftracer/python/ai.py
 python/dftracer/python/ai_common.py
 python/dftracer/python/ai_init.py
@@ -35,6 +40,7 @@ python/dftracer/python/env.py
 python/dftracer/python/logger.py
 python/dftracer/python/py.typed
 python/dftracer/python/dbg/__init__.py
+python/dftracer/python/dbg/agent.py
 python/dftracer/python/dbg/ai.py
 python/dftracer/python/dbg/logger.py
 python/dftracer/python/dbg/torch/__init__.py

{pydftracer-2.0.5.post2.dev0 → pydftracer-2.0.6.post2.dev0}/python/pydftracer.egg-info/scm_file_list.json

@@ -14,6 +14,7 @@
     "docs/make.bat",
     "docs/requirements.txt",
     "docs/source/ai_ml_guide.rst",
+    "docs/source/api/agent.rst",
     "docs/source/api/ai.rst",
     "docs/source/api/core.rst",
     "docs/source/api/dynamo.rst",
@@ -30,11 +31,16 @@
     "docs/source/type_safety.rst",
     "pyproject.toml",
     "python/dftracer/python/__init__.py",
+    "python/dftracer/python/agent.py",
+    "python/dftracer/python/agent_common.py",
+    "python/dftracer/python/agent_init.py",
+    "python/dftracer/python/agent_langgraph.py",
     "python/dftracer/python/ai.py",
     "python/dftracer/python/ai_common.py",
     "python/dftracer/python/ai_init.py",
     "python/dftracer/python/common.py",
     "python/dftracer/python/dbg/__init__.py",
+    "python/dftracer/python/dbg/agent.py",
     "python/dftracer/python/dbg/ai.py",
     "python/dftracer/python/dbg/logger.py",
     "python/dftracer/python/dbg/torch/__init__.py",
@@ -51,6 +57,8 @@
     "tests/__init__.py",
     "tests/poc/pytorch_profiler/requirements.txt",
     "tests/poc/pytorch_profiler/test.py",
+    "tests/test_agent.py",
+    "tests/test_agent_langgraph.py",
     "tests/test_ai_logging.py",
     "tests/test_coverage_improvements.py",
     "tests/test_dftracer.py",

pydftracer-2.0.6.post2.dev0/python/pydftracer.egg-info/scm_version.json

@@ -0,0 +1,8 @@
+{
+  "tag": "2.0.6",
+  "distance": 2,
+  "node": "g68307eab60eb6cfecdcda3ad507585fb8241a2f6",
+  "dirty": false,
+  "branch": "develop",
+  "node_date": "2026-10-27"
+}

pydftracer-2.0.5.post2.dev0/python/pydftracer.egg-info/scm_version.json

@@ -1,8 +0,0 @@
-{
-  "tag": "2.0.5",
-  "distance": 2,
-  "node": "gb88b10fd4221cd6ed3653a3c34c1eab883bce041",
-  "dirty": false,
-  "branch": "develop",
-  "node_date": "2026-10-14"
-}