feat(client): cache-aware token accounting and a reusable span lifecycle

packages/client/src/launchdarkly_ai_server/__init__.py

@@ -34,6 +34,7 @@
     HandlerResult,
     HandlerStreamEvent,
     InitClientOptions,
+    InputTokenDetails,
     JudgeResult,
     JudgeRunResult,
     JudgeTask,
@@ -56,16 +57,26 @@
 )
 from .types_validation import parse_ai_config
 from .utils import (
+    RunUsage,
+    SpanUsage,
+    add_cached_tokens_to_input,
     create_handler,
+    create_run_usage,
+    end_span_once,
+    lang_chain_span_usage,
     make_track_data,
     normalize_mode,
+    number_or_zero,
     parse_json_with_possible_fences,
     parse_template,
     parse_usage,
     set_ld_span_attributes,
+    set_model_identity_attributes,
     set_openllmetry_completion,
     set_openllmetry_prompt,
+    set_usage_span_attributes,
     to_ld_context,
+    to_usage_dict,
 )
 
 __all__ = [  # noqa: RUF022
@@ -100,7 +111,18 @@
     "StreamDoneEvent",
     "StreamEvent",
     "TrackData",
+    "InputTokenDetails",
+    "RunUsage",
+    "SpanUsage",
     "UsageDict",
+    "add_cached_tokens_to_input",
+    "create_run_usage",
+    "end_span_once",
+    "lang_chain_span_usage",
+    "number_or_zero",
+    "set_model_identity_attributes",
+    "set_usage_span_attributes",
+    "to_usage_dict",
     "VariationMeta",
     # utils
     "create_handler",

packages/client/src/launchdarkly_ai_server/client.py

@@ -15,12 +15,12 @@
     ProviderHandler,
     ProviderResponse,
     StreamEvent,
-    UsageDict,
     VariationMeta,
 )
 from .utils import (
     parse_json_with_possible_fences,
     select_handler,
+    to_usage_dict,
 )
 
 
@@ -112,11 +112,7 @@ async def invoke(
             else json.dumps(parsed_response)
         )
 
-        usage_obj = UsageDict(
-            input=usage.get("input", 0),
-            output=usage.get("output", 0),
-            total=usage.get("total", 0),
-        )
+        usage_obj = to_usage_dict(usage)
 
         if self._skip_judges:
             judge_tasks = await build_judge_tasks(

packages/client/src/launchdarkly_ai_server/graph.py

@@ -668,7 +668,15 @@ async def invoke(
 
             return ProviderGraphResponse(
                 response=final_response,
-                usage=UsageDict(**total_usage),
+                # Named rather than splatted, so a new UsageDict member cannot silently arrive
+                # here from a dict that has no business filling it. Graph totals carry no cache
+                # breakdown: they are a sum across nodes, and the per-node detail is on the node's
+                # own spans.
+                usage=UsageDict(
+                    input=total_usage["input"],
+                    output=total_usage["output"],
+                    total=total_usage["total"],
+                ),
                 judge_results=judge_results,
             )
 

packages/client/src/launchdarkly_ai_server/judges.py

@@ -13,7 +13,6 @@
     NativeTool,
     ProviderHandler,
     TrackData,
-    UsageDict,
 )
 from .utils import (
     collapse_messages_to_instructions as _collapse_messages_to_instructions,
@@ -22,6 +21,7 @@
     normalize_mode,
     parse_json_with_possible_fences,
     to_ld_context,
+    to_usage_dict,
 )
 
 
@@ -409,11 +409,7 @@ def _matches(h: ProviderHandler) -> bool:
     reasoning = parsed.get("reasoning", "")
     raw_usage = result["usage"]
 
-    usage = UsageDict(
-        input=raw_usage.get("input", 0),
-        output=raw_usage.get("output", 0),
-        total=raw_usage.get("total", 0),
-    )
+    usage = to_usage_dict(raw_usage)
 
     merged_track_data: TrackData = {
         **task.parent_track_data,

packages/client/src/launchdarkly_ai_server/types.py

@@ -192,11 +192,26 @@ def has_stream(self) -> bool:
 # ---------------------------------------------------------------------------
 
 
+@dataclass
+class InputTokenDetails:
+    """The cache breakdown behind an inclusive ``input`` figure.
+
+    Present only when the provider reported at least one cache field. ``uncached + cache_read +
+    cache_creation`` equals :attr:`UsageDict.input`.
+    """
+
+    uncached: int = 0
+    cache_read: int = 0
+    cache_creation: int = 0
+
+
 @dataclass
 class UsageDict:
     input: int = 0
     output: int = 0
     total: int = 0
+    #: Cache breakdown, when the provider reported one. See :class:`InputTokenDetails`.
+    input_details: InputTokenDetails | None = None
 
 
 @dataclass