feat: Add unified token usage and API call tracking

README.md

@@ -26,6 +26,7 @@
   - [*Romeo and Juliet* Full Text Extraction](#romeo-and-juliet-full-text-extraction)
   - [Medication Extraction](#medication-extraction)
   - [Radiology Report Structuring: RadExtract](#radiology-report-structuring-radextract)
+  - [Tracking Token Usage & API Calls](#tracking-token-usage--api-calls)
 - [Community Providers](#community-providers)
 - [Contributing](#contributing)
 - [Testing](#testing)
@@ -175,6 +176,52 @@ result = lx.extract(
 
 This approach can extract hundreds of entities from full novels while maintaining high accuracy. The interactive visualization seamlessly handles large result sets, making it easy to explore hundreds of entities from the output JSONL file. **[See the full *Romeo and Juliet* extraction example →](https://github.com/google/langextract/blob/main/docs/examples/longer_text_example.md)** for detailed results and performance insights.
 
+### Tracking Token Usage & API Calls
+
+By default, the returned `AnnotatedDocument` contains aggregated token usage and API call metrics inside its `metadata` dictionary:
+
+```python
+result = lx.extract(
+    text_or_documents=input_text,
+    prompt_description=prompt,
+    examples=examples,
+    model_id="gemini-3.5-flash",
+)
+
+print(result.metadata)
+# Output:
+# {
+#   'token_usage': {'prompt_tokens': 197, 'completion_tokens': 82, 'total_tokens': 279},
+#   'api_calls': 1
+# }
+```
+
+> **Note:** The total number of `api_calls` is equal to `(number of chunks) * (extraction_passes)`.
+>
+> **Warning:** Enabling `track_api_call_details=True` on extremely large runs (with thousands of chunks or passes) can consume significant memory because details for every single API call are stored in memory. Only use it when detailed observability is required.
+
+For detailed observability (e.g., tracking the exact token cost of each chunk or extraction pass), you can opt-in to detailed call-level tracking by setting `track_api_call_details=True`:
+
+```python
+result = lx.extract(
+    text_or_documents=input_text,
+    prompt_description=prompt,
+    examples=examples,
+    model_id="gemini-3.5-flash",
+    track_api_call_details=True,
+)
+
+print(result.metadata["api_call_details"])
+# Output:
+# [
+#   {
+#     'pass_index': 0,
+#     'chunk_index': 0,
+#     'token_usage': {'prompt_tokens': 197, 'completion_tokens': 82, 'total_tokens': 279}
+#   }
+# ]
+```
+
 ### Vertex AI Batch Processing
 
 Save costs on large-scale tasks by enabling Vertex AI Batch API with

langextract/annotation.py

@@ -217,6 +217,7 @@ def annotate_documents(
       context_window_chars: int | None = None,
       show_progress: bool = True,
       tokenizer: tokenizer_lib.Tokenizer | None = None,
+      track_api_call_details: bool = False,
       **kwargs,
   ) -> Iterator[data.AnnotatedDocument]:
     """Annotates a sequence of documents with NLP extractions.
@@ -244,6 +245,9 @@ def annotate_documents(
         resolution across chunk boundaries. Defaults to None (disabled).
       show_progress: Whether to show progress bar. Defaults to True.
       tokenizer: Optional tokenizer to use. If None, uses default tokenizer.
+      track_api_call_details: Whether to track detailed tokens of individual API calls.
+        Warning: Enabling this on extremely large runs with thousands of chunks/passes can consume
+        significant memory.
       **kwargs: Additional arguments passed to LanguageModel.infer and
         Resolver.
 
@@ -266,6 +270,7 @@ def annotate_documents(
           show_progress,
           context_window_chars=context_window_chars,
           tokenizer=tokenizer,
+          track_api_call_details=track_api_call_details,
           **kwargs,
       )
     else:
@@ -279,6 +284,7 @@ def annotate_documents(
           show_progress,
           context_window_chars=context_window_chars,
           tokenizer=tokenizer,
+          track_api_call_details=track_api_call_details,
           **kwargs,
       )
 
@@ -293,6 +299,8 @@ def _annotate_documents_single_pass(
       context_window_chars: int | None = None,
       tokenizer: tokenizer_lib.Tokenizer | None = None,
       suppress_parse_errors: bool = False,
+      track_api_call_details: bool = False,
+      pass_num: int = 0,
       **kwargs,
   ) -> Iterator[data.AnnotatedDocument]:
     """Single-pass annotation with stable ordering and streaming emission.
@@ -309,6 +317,14 @@ def _annotate_documents_single_pass(
     per_doc: DefaultDict[str, list[data.Extraction]] = collections.defaultdict(
         list
     )
+    doc_usage = collections.defaultdict(
+        lambda: {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+    )
+    # doc_successful_api_calls tracks the number of successful inference API requests
+    # that completed and returned a scored output.
+    doc_successful_api_calls = collections.defaultdict(int)
+    doc_api_call_details = collections.defaultdict(list)
+    chunk_counters = collections.defaultdict(int)
     next_emit_idx = 0
 
     def _capture_docs(src: Iterable[data.Document]) -> Iterator[data.Document]:
@@ -336,13 +352,26 @@ def _emit_docs_iter(
       limit = max(0, len(doc_order) - 1) if keep_last_doc else len(doc_order)
       while next_emit_idx < limit:
         document_id = doc_order[next_emit_idx]
+        metadata = {
+            "token_usage": doc_usage.get(document_id),
+            "api_calls": doc_successful_api_calls.get(document_id, 0),
+        }
+        if track_api_call_details:
+          metadata["api_call_details"] = doc_api_call_details.get(
+              document_id, []
+          )
         yield data.AnnotatedDocument(
             document_id=document_id,
             extractions=per_doc.get(document_id, []),
             text=doc_text_by_id.get(document_id, ""),
+            metadata=metadata,
         )
         per_doc.pop(document_id, None)
         doc_text_by_id.pop(document_id, None)
+        doc_usage.pop(document_id, None)
+        doc_successful_api_calls.pop(document_id, None)
+        doc_api_call_details.pop(document_id, None)
+        chunk_counters.pop(document_id, None)
         next_emit_idx += 1
 
     chunk_iter = _document_chunk_iterator(
@@ -401,8 +430,37 @@ def _emit_docs_iter(
                 "No scored outputs from language model."
             )
 
+          doc_id = text_chunk.document_id
+          doc_successful_api_calls[doc_id] += 1
+          scored_output = scored_outputs[0]
+          usage = scored_output.token_usage
+          chunk_idx = chunk_counters[doc_id]
+          chunk_counters[doc_id] += 1
+
+          if usage is not None:
+            for key in ("prompt_tokens", "completion_tokens", "total_tokens"):
+              val = getattr(usage, key, None)
+              if isinstance(val, (int, float)):
+                doc_usage[doc_id][key] += val
+
+          if track_api_call_details:
+            detail = {
+                "pass_index": pass_num,
+                "chunk_index": chunk_idx,
+                "token_usage": {
+                    "prompt_tokens": usage.prompt_tokens if usage else None,
+                    "completion_tokens": (
+                        usage.completion_tokens if usage else None
+                    ),
+                    "total_tokens": usage.total_tokens if usage else None,
+                },
+            }
+            if scored_output.request_id is not None:
+              detail["request_id"] = scored_output.request_id
+            doc_api_call_details[doc_id].append(detail)
+
           resolved_extractions = resolver.resolve(
-              scored_outputs[0].output,
+              scored_output.output,
               debug=debug,
               suppress_parse_errors=suppress_parse_errors,
               **kwargs,
@@ -455,6 +513,7 @@ def _annotate_documents_sequential_passes(
       show_progress: bool = True,
       context_window_chars: int | None = None,
       tokenizer: tokenizer_lib.Tokenizer | None = None,
+      track_api_call_details: bool = False,
       **kwargs,
   ) -> Iterator[data.AnnotatedDocument]:
     """Sequential extraction passes logic for improved recall."""
@@ -469,6 +528,10 @@ def _annotate_documents_sequential_passes(
 
     document_extractions_by_pass: dict[str, list[list[data.Extraction]]] = {}
     document_texts: dict[str, str] = {}
+    document_usage = {}
+    document_api_calls = {}
+    document_api_call_details = {}
+
     # Preserve text up-front so we can emit documents even if later passes
     # produce no extractions.
     for _doc in document_list:
@@ -488,18 +551,43 @@ def _annotate_documents_sequential_passes(
           show_progress=show_progress if pass_num == 0 else False,
           context_window_chars=context_window_chars,
           tokenizer=tokenizer,
+          track_api_call_details=track_api_call_details,
+          pass_num=pass_num,
           **kwargs,
       ):
         doc_id = annotated_doc.document_id
 
         if doc_id not in document_extractions_by_pass:
           document_extractions_by_pass[doc_id] = []
-          # Keep first-seen text (already pre-filled above).
+          document_usage[doc_id] = {
+              "prompt_tokens": 0,
+              "completion_tokens": 0,
+              "total_tokens": 0,
+          }
+          document_api_calls[doc_id] = 0
+          document_api_call_details[doc_id] = []
 
         document_extractions_by_pass[doc_id].append(
             annotated_doc.extractions or []
         )
 
+        if annotated_doc.metadata:
+          meta = annotated_doc.metadata
+          usage = meta.get("token_usage")
+          if usage:
+            document_usage[doc_id]["prompt_tokens"] += (
+                usage.get("prompt_tokens") or 0
+            )
+            document_usage[doc_id]["completion_tokens"] += (
+                usage.get("completion_tokens") or 0
+            )
+            document_usage[doc_id]["total_tokens"] += (
+                usage.get("total_tokens") or 0
+            )
+          document_api_calls[doc_id] += meta.get("api_calls", 0)
+          if "api_call_details" in meta:
+            document_api_call_details[doc_id].extend(meta["api_call_details"])
+
     # Emit results strictly in original input order.
     for doc in document_list:
       doc_id = doc.document_id
@@ -521,10 +609,18 @@ def _annotate_documents_sequential_passes(
             len(merged_extractions),
         )
 
+      metadata = {
+          "token_usage": document_usage.get(doc_id),
+          "api_calls": document_api_calls.get(doc_id, 0),
+      }
+      if track_api_call_details:
+        metadata["api_call_details"] = document_api_call_details.get(doc_id, [])
+
       yield data.AnnotatedDocument(
           document_id=doc_id,
           extractions=merged_extractions,
           text=document_texts.get(doc_id, doc.text or ""),
+          metadata=metadata,
       )
 
     logging.info("Sequential extraction passes completed.")
@@ -541,6 +637,7 @@ def annotate_text(
       context_window_chars: int | None = None,
       show_progress: bool = True,
       tokenizer: tokenizer_lib.Tokenizer | None = None,
+      track_api_call_details: bool = False,
       **kwargs,
   ) -> data.AnnotatedDocument:
     """Annotates text with NLP extractions for text input.
@@ -562,6 +659,9 @@ def annotate_text(
         (disabled).
       show_progress: Whether to show progress bar. Defaults to True.
       tokenizer: Optional tokenizer instance.
+      track_api_call_details: Whether to track detailed tokens of individual API calls.
+        Warning: Enabling this on extremely large runs with thousands of chunks/passes can consume
+        significant memory.
       **kwargs: Additional arguments for inference and resolver_lib.
 
     Returns:
@@ -593,6 +693,7 @@ def annotate_text(
             context_window_chars=context_window_chars,
             show_progress=show_progress,
             tokenizer=tokenizer,
+            track_api_call_details=track_api_call_details,
             **kwargs,
         )
     )
@@ -622,4 +723,5 @@ def annotate_text(
         document_id=annotations[0].document_id,
         extractions=annotations[0].extractions,
         text=annotations[0].text,
+        metadata=annotations[0].metadata,
     )

langextract/core/data.py

@@ -17,6 +17,7 @@
 
 import dataclasses
 import enum
+from typing import Any
 import uuid
 
 from langextract.core import tokenizer
@@ -212,10 +213,14 @@ class AnnotatedDocument:
     extractions: List of extractions in the document.
     text: Raw text representation of the document.
     tokenized_text: Tokenized text of the document, computed from `text`.
+    metadata: Metadata dict (e.g. token_usage, api_calls) for the document.
   """
 
   extractions: list[Extraction] | None = None
   text: str | None = None
+  metadata: dict[str, Any] | None = dataclasses.field(
+      default=None, compare=False
+  )
   _document_id: str | None = dataclasses.field(
       default=None, init=False, repr=False, compare=False
   )
@@ -229,9 +234,11 @@ def __init__(
       document_id: str | None = None,
       extractions: list[Extraction] | None = None,
       text: str | None = None,
+      metadata: dict[str, Any] | None = None,
   ):
     self.extractions = extractions
     self.text = text
+    self.metadata = metadata
     self._document_id = document_id
 
   @property

langextract/core/types.py

@@ -66,12 +66,25 @@ class Constraint:
   constraint_type: ConstraintType = ConstraintType.NONE
 
 
+@dataclasses.dataclass(frozen=True)
+class TokenUsage:
+  """Token usage details for a model request/run."""
+
+  prompt_tokens: int | None = None
+  completion_tokens: int | None = None
+  total_tokens: int | None = None
+
+
 @dataclasses.dataclass(frozen=True)
 class ScoredOutput:
   """Scored output from language model inference."""
 
   score: float | None = None
   output: str | None = None
+  token_usage: TokenUsage | None = dataclasses.field(
+      default=None, compare=False
+  )
+  request_id: str | None = dataclasses.field(default=None, compare=False)
 
   def __str__(self) -> str:
     score_str = '-' if self.score is None else f'{self.score:.2f}'

langextract/data_lib.py

@@ -79,6 +79,9 @@ def annotated_document_to_dict(
 
   result["document_id"] = adoc.document_id
 
+  if adoc.metadata is None:
+    result.pop("metadata", None)
+
   return result
 
 
@@ -121,4 +124,5 @@ def dict_to_annotated_document(
       extractions=[
           data.Extraction(**ent) for ent in adoc_dic.get("extractions", [])
       ],
+      metadata=adoc_dic.get("metadata"),
   )