Export observe cost helpers (#1) (#12)

Author: Vam-si-krish

grapharc/observe/__init__.py

@@ -13,6 +13,8 @@
     ThreadCost,
     attribute,
     attribute_thread,
+    by_node,
+    tokens_by_model,
 )
 from grapharc.observe.metrics import RunMetrics, summarize, to_mermaid
 from grapharc.observe.otel import (
@@ -62,6 +64,7 @@
     "TraceRecorder",
     "attribute",
     "attribute_thread",
+    "by_node",
     "diff_runs",
     "diff_trace",
     "export_run",
@@ -73,4 +76,5 @@
     "summarize",
     "to_mermaid",
     "to_spans",
+    "tokens_by_model",
 ]

tests/test_replay.py

@@ -931,6 +931,14 @@ def test_attribute_thread_sums_across_resumes(trace):
     )
 
 
+def test_cost_helpers_are_exported_from_observe():
+    from grapharc.observe import by_node as exported_by_node
+    from grapharc.observe import tokens_by_model as exported_tokens_by_model
+
+    assert exported_by_node is by_node
+    assert exported_tokens_by_model is tokens_by_model
+
+
 def test_attribute_by_node_ranks_across_every_run(trace):
     trace.event(run_id="r1", graph="demo", node="cheap", phase="start", step=1)
     trace.event(run_id="r1", graph="demo", node="cheap", phase="end", step=1,