Repository: stanfordnlp/dspy
Pull request #8014 (2025-03-26 01:49:08)

Add usage tracking
Per community request, we are adding usage tracking for DSPy program. As an overview:

- We track the LM usage in the openai (litellm) format for each program call
- Local cache hit will cause no usage
- If multiple LMs are used in the process, they will be tracked separately. 
- The returned usage is a dictionary from lm name to openai-format usage data. The usage data is the merge result across all LM calls in one single program pass.


To use usage tracking, follow the code example below:

```
import dspy

dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=False), track_usage=True)

predict = dspy.ChainOfThought("question -> answer")
output = predict(question="What is the capital of France?")
print(output.get_lm_usage())
```

This usage tracking also works for multi-module program, for example:

```
import dspy

dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=False), track_usage=True)

class MyProgram(dspy.Module):
    def __init__(self):
        self.predict1 = dspy.ChainOfThought("question -> answer")
        self.predict2 = dspy.ChainOfThought("question, answer -> score")

    def __call__(self, question: str) -> str:
        answer = self.predict1(question=question)
        score = self.predict2(question=question, answer=answer)
        return score

program = MyProgram()
output = program(question="What is the capital of France?")
print(output.get_lm_usage())
```

The output is the combined output for the end to end program call. 

This usage tracking also works with multithreading:

```
class MyProgram(dspy.Module):
    def __init__(self, lm):
        self.lm = lm
        self.predict1 = dspy.ChainOfThought("question -> answer")
        self.predict2 = dspy.ChainOfThought("question, answer -> score")

    def __call__(self, question: str) -> str:
        with dspy.settings.context(lm=self.lm):
            answer = self.predict1(question=question)
            score = self.predict2(question=question, answer=answer)
            return score

dspy.settings.configure(track_usage=True)
program1 = MyProgram(lm=dspy.LM("openai/gpt-4o-mini", cache=False))
program2 = MyProgram(lm=dspy.LM("openai/gpt-3.5-turbo", cache=False))

parallelizer = dspy.Parallel()

results = parallelizer(
    [
        (program1, {"question": "What is the meaning of life?"}),
        (program2, {"question": "why did a chicken cross the kitchen?"}),
    ]
)
```
it is thread-safe, each program's output will only see the usage of itself. 


Patch:
diff --git a/.prettierignore b/.prettierignore
deleted file mode 100644
index 6d2eaaa4..00000000
--- a/.prettierignore
+++ /dev/null
@@ -1,1 +0,0 @@
-.github/
\ No newline at end of file
diff --git a/docs/docs/deep-dive/optimizers/miprov2.md b/docs/docs/deep-dive/optimizers/miprov2.md
index e1960122..5a030744 100644
--- a/docs/docs/deep-dive/optimizers/miprov2.md
+++ b/docs/docs/deep-dive/optimizers/miprov2.md
@@ -31,7 +31,7 @@ from dspy.datasets.gsm8k import GSM8K, gsm8k_metric
 
 gsm8k = GSM8K()
 
-trainset, devset = gsm8k.train, gms8k.dev
+trainset, devset = gsm8k.train, gsm8k.dev
 ```
 
 We'll now define a basic QA inline signature i.e. `question->answer` and pass it to `ChainOfThought` module, that applies necessary addition for CoT style prompting to the Signature.
diff --git a/docs/docs/learn/optimization/optimizers.md b/docs/docs/learn/optimization/optimizers.md
index 287025d4..9328f642 100644
--- a/docs/docs/learn/optimization/optimizers.md
+++ b/docs/docs/learn/optimization/optimizers.md
@@ -47,7 +47,7 @@ These optimizers extend the signature by automatically generating and including
 
 3. [**`BootstrapFewShotWithRandomSearch`**](/deep-dive/optimizers/bootstrap-fewshot): Applies `BootstrapFewShot` several times with random search over generated demonstrations, and selects the best program over the optimization. Parameters mirror those of `BootstrapFewShot`, with the addition of `num_candidate_programs`, which specifies the number of random programs evaluated over the optimization, including candidates of the uncompiled program, `LabeledFewShot` optimized program, `BootstrapFewShot` compiled program with unshuffled examples and `num_candidate_programs` of `BootstrapFewShot` compiled programs with randomized example sets.
 
-4. **`KNNFewShot`**. Uses k-Nearest Neighbors algorithm to find the nearest training example demonstrations for a given input example. These nearest neighbor demonstrations are then used as the trainset for the BootstrapFewShot optimization process. See [this notebook](https://github.com/stanfordnlp/dspy/blob/main/examples/knn.ipynb) for an example.
+4. **`KNNFewShot`**. Uses k-Nearest Neighbors algorithm to find the nearest training example demonstrations for a given input example. These nearest neighbor demonstrations are then used as the trainset for the BootstrapFewShot optimization process. See [this notebook](https://github.com/stanfordnlp/dspy/blob/main/examples/outdated_v2.4_examples/knn.ipynb) for an example.
 
 
 ### Automatic Instruction Optimization
@@ -147,14 +147,14 @@ optimized_program = teleprompter.compile(YOUR_PROGRAM_HERE, trainset=YOUR_TRAINS
         optimized_rag = tp.compile(RAG(), trainset=trainset, max_bootstrapped_demos=2, max_labeled_demos=2)
         ```
 
-        For a complete RAG example that you can run, start this [tutorial](http://127.0.0.1:8000/quick-start/getting-started-01/). It improves the quality of a RAG system over a subset of StackExchange communities from 53% to 61%.
+        For a complete RAG example that you can run, start this [tutorial](/tutorials/rag/). It improves the quality of a RAG system over a subset of StackExchange communities from 53% to 61%.
 
     === "Optimizing weights for Classification"
         This is a minimal but fully runnable example of setting up a `dspy.ChainOfThought` module that classifies
         short texts into one of 77 banking labels and then using `dspy.BootstrapFinetune` with 2000 text-label pairs
         from the `Banking77` to finetune the weights of GPT-4o-mini for this task. We use the variant
         `dspy.ChainOfThoughtWithHint`, which takes an optional `hint` at bootstrapping time, to maximize the utility of
-        the training data. Naturally, hints are not available at test time.
+        the training data. Naturally, hints are not available at test time. More can be found in this [tutorial](/tutorials/classification_finetuning/).
 
         <details><summary>Click to show dataset setup code.</summary>
 
diff --git a/docs/docs/learn/programming/modules.md b/docs/docs/learn/programming/modules.md
index 3275ee0a..792a89af 100644
--- a/docs/docs/learn/programming/modules.md
+++ b/docs/docs/learn/programming/modules.md
@@ -244,3 +244,93 @@ class Hop(dspy.Module):
         
         return dspy.Prediction(notes=notes, titles=list(set(titles)))
 ```
+
+## How do I track LM usage?
+
+!!! note "Version Requirement"
+    LM usage tracking is available in DSPy version 2.6.16 and later.
+
+DSPy provides built-in tracking of language model usage across all module calls. To enable tracking:
+
+```python
+dspy.settings.configure(track_usage=True)
+```
+
+Once enabled, you can access usage statistics from any `dspy.Prediction` object:
+
+```python
+usage = prediction_instance.get_lm_usage()
+```
+
+The usage data is returned as a dictionary that maps each language model name to its usage statistics. Here's a complete example:
+
+```python
+import dspy
+
+# Configure DSPy with tracking enabled
+dspy.settings.configure(
+    lm=dspy.LM("openai/gpt-4o-mini", cache=False),
+    track_usage=True
+)
+
+# Define a simple program that makes multiple LM calls
+class MyProgram(dspy.Module):
+    def __init__(self):
+        self.predict1 = dspy.ChainOfThought("question -> answer")
+        self.predict2 = dspy.ChainOfThought("question, answer -> score")
+
+    def __call__(self, question: str) -> str:
+        answer = self.predict1(question=question)
+        score = self.predict2(question=question, answer=answer)
+        return score
+
+# Run the program and check usage
+program = MyProgram()
+output = program(question="What is the capital of France?")
+print(output.get_lm_usage())
+```
+
+This will output usage statistics like:
+
+```python
+{
+    'openai/gpt-4o-mini': {
+        'completion_tokens': 61,
+        'prompt_tokens': 260,
+        'total_tokens': 321,
+        'completion_tokens_details': {
+            'accepted_prediction_tokens': 0,
+            'audio_tokens': 0,
+            'reasoning_tokens': 0,
+            'rejected_prediction_tokens': 0,
+            'text_tokens': None
+        },
+        'prompt_tokens_details': {
+            'audio_tokens': 0,
+            'cached_tokens': 0,
+            'text_tokens': None,
+            'image_tokens': None
+        }
+    }
+}
+```
+
+When using DSPy's caching features (either in-memory or on-disk via litellm), cached responses won't count toward usage statistics. For example:
+
+```python
+# Enable caching
+dspy.settings.configure(
+    lm=dspy.LM("openai/gpt-4o-mini", cache=True),
+    track_usage=True
+)
+
+program = MyProgram()
+
+# First call - will show usage statistics
+output = program(question="What is the capital of Zambia?")
+print(output.get_lm_usage())  # Shows token usage
+
+# Second call - same question, will use cache
+output = program(question="What is the capital of Zambia?")
+print(output.get_lm_usage())  # Shows empty dict: {}
+```
diff --git a/docs/requirements.txt b/docs/requirements.txt
new file mode 100644
index 00000000..8083e88f
--- /dev/null
+++ b/docs/requirements.txt
@@ -0,0 +1,9 @@
+git+https://github.com/stanfordnlp/dspy.git
+mkdocs-material
+mkdocs-jupyter
+mkdocs-material[imaging]
+mkdocs-redirects
+mkdocstrings
+mkdocstrings-python
+urllib3==1.26.6
+mistune==3.0.2
diff --git a/dspy/__init__.py b/dspy/__init__.py
index 74cad5a3..70e7de70 100644
--- a/dspy/__init__.py
+++ b/dspy/__init__.py
@@ -13,6 +13,7 @@ from dspy.utils.logging_utils import configure_dspy_loggers, disable_logging, en
 from dspy.utils.asyncify import asyncify
 from dspy.utils.saving import load
 from dspy.utils.streaming import streamify
+from dspy.utils.usage_tracker import track_usage
 
 from dspy.dsp.utils.settings import settings
 
diff --git a/dspy/clients/__init__.py b/dspy/clients/__init__.py
index 19217fb7..48fa730a 100644
--- a/dspy/clients/__init__.py
+++ b/dspy/clients/__init__.py
@@ -13,6 +13,14 @@ logger = logging.getLogger(__name__)
 DISK_CACHE_DIR = os.environ.get("DSPY_CACHEDIR") or os.path.join(Path.home(), ".dspy_cache")
 DISK_CACHE_LIMIT = int(os.environ.get("DSPY_CACHE_LIMIT", 3e10))  # 30 GB default
 
+
+def _litellm_track_cache_hit_callback(kwargs, completion_response, start_time, end_time):
+    # Access the cache_hit information
+    completion_response.cache_hit = kwargs.get("cache_hit", False)
+
+
+litellm.success_callback = [_litellm_track_cache_hit_callback]
+
 try:
     # TODO: There's probably value in getting litellm to support FanoutCache and to separate the limit for
     # the LM cache from the embeddings cache. Then we can lower the default 30GB limit.
diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py
index adaba55e..39dcbc0b 100644
--- a/dspy/clients/lm.py
+++ b/dspy/clients/lm.py
@@ -18,6 +18,7 @@ import dspy
 from dspy.clients.openai import OpenAIProvider
 from dspy.clients.provider import Provider, TrainingJob
 from dspy.clients.utils_finetune import TrainDataFormat
+from dspy.dsp.utils.settings import settings
 from dspy.utils.callback import BaseCallback, with_callbacks
 
 from .base_lm import BaseLM
@@ -108,20 +109,24 @@ class LM(BaseLM):
         if cache_in_memory:
             completion = cached_litellm_completion if self.model_type == "chat" else cached_litellm_text_completion
 
-            return completion(
+            results = completion(
                 request=dict(model=self.model, messages=messages, **kwargs),
                 num_retries=self.num_retries,
             )
         else:
             completion = litellm_completion if self.model_type == "chat" else litellm_text_completion
 
-            return completion(
+            results = completion(
                 request=dict(model=self.model, messages=messages, **kwargs),
                 num_retries=self.num_retries,
                 # only leverage LiteLLM cache in this case
                 cache={"no-cache": not cache, "no-store": not cache},
             )
 
+        if not getattr(results, "cache_hit", False) and dspy.settings.usage_tracker and hasattr(results, "usage"):
+            settings.usage_tracker.add_usage(self.model, dict(results.usage))
+        return results
+
     def launch(self, launch_kwargs: Optional[Dict[str, Any]] = None):
         self.provider.launch(self, launch_kwargs)
 
@@ -259,6 +264,12 @@ def request_cache(maxsize: Optional[int] = None):
                 # If the cache key cannot be computed (e.g. because it contains a value that cannot
                 # be converted to JSON), bypass the cache and call the target function directly
                 return func(request, *args, **kwargs)
+            cache_hit = key in func_cached.cache
+            output = func_cached(key, request, *args, **kwargs)
+            if cache_hit and hasattr(output, "usage"):
+                # Clear the usage data when cache is hit, because no LM call is made
+                output.usage = {}
+
             return func_cached(key, request, *args, **kwargs)
 
         return wrapper
diff --git a/dspy/dsp/utils/settings.py b/dspy/dsp/utils/settings.py
index 98d96d6a..fad3cc67 100644
--- a/dspy/dsp/utils/settings.py
+++ b/dspy/dsp/utils/settings.py
@@ -20,6 +20,8 @@ DEFAULT_CONFIG = dotdict(
     async_max_workers=8,
     send_stream=None,
     disable_history=False,
+    track_usage=False,
+    usage_tracker=None,
 )
 
 # Global base configuration and owner tracking
diff --git a/dspy/predict/parallel.py b/dspy/predict/parallel.py
index 9fa74ebe..0c03ad96 100644
--- a/dspy/predict/parallel.py
+++ b/dspy/predict/parallel.py
@@ -1,6 +1,5 @@
 import threading
-
-from typing import Tuple, List, Any
+from typing import Any, List, Tuple
 
 from dspy.primitives.example import Example
 from dspy.utils.parallelizer import ParallelExecutor
diff --git a/dspy/primitives/prediction.py b/dspy/primitives/prediction.py
index 63b74873..b05b3857 100644
--- a/dspy/primitives/prediction.py
+++ b/dspy/primitives/prediction.py
@@ -9,6 +9,13 @@ class Prediction(Example):
         del self._input_keys
 
         self._completions = None
+        self._lm_usage = None
+
+    def get_lm_usage(self):
+        return self._lm_usage
+
+    def set_lm_usage(self, value):
+        self._lm_usage = value
 
     @classmethod
     def from_completions(cls, list_or_dict, signature=None):
@@ -29,7 +36,7 @@ class Prediction(Example):
 
     def __str__(self):
         return self.__repr__()
-    
+
     def __float__(self):
         if "score" not in self._store:
             raise ValueError("Prediction object does not have a 'score' field to convert to float.")
@@ -41,7 +48,7 @@ class Prediction(Example):
         elif isinstance(other, Prediction):
             return self.__float__() + float(other)
         raise TypeError(f"Unsupported type for addition: {type(other)}")
-    
+
     def __radd__(self, other):
         if isinstance(other, (float, int)):
             return other + self.__float__()
@@ -55,7 +62,7 @@ class Prediction(Example):
         elif isinstance(other, Prediction):
             return self.__float__() / float(other)
         raise TypeError(f"Unsupported type for division: {type(other)}")
-    
+
     def __rtruediv__(self, other):
         if isinstance(other, (float, int)):
             return other / self.__float__()
@@ -69,21 +76,21 @@ class Prediction(Example):
         elif isinstance(other, Prediction):
             return self.__float__() < float(other)
         raise TypeError(f"Unsupported type for comparison: {type(other)}")
-    
+
     def __le__(self, other):
         if isinstance(other, (float, int)):
             return self.__float__() <= other
         elif isinstance(other, Prediction):
             return self.__float__() <= float(other)
         raise TypeError(f"Unsupported type for comparison: {type(other)}")
-    
+
     def __gt__(self, other):
         if isinstance(other, (float, int)):
             return self.__float__() > other
         elif isinstance(other, Prediction):
             return self.__float__() > float(other)
         raise TypeError(f"Unsupported type for comparison: {type(other)}")
-    
+
     def __ge__(self, other):
         if isinstance(other, (float, int)):
             return self.__float__() >= other
diff --git a/dspy/primitives/program.py b/dspy/primitives/program.py
index 07b5cd4b..a0592df3 100644
--- a/dspy/primitives/program.py
+++ b/dspy/primitives/program.py
@@ -1,8 +1,10 @@
 import magicattr
 
+from dspy.dsp.utils.settings import settings
 from dspy.predict.parallel import Parallel
 from dspy.primitives.module import BaseModule
 from dspy.utils.callback import with_callbacks
+from dspy.utils.usage_tracker import track_usage
 
 
 class ProgramMeta(type):
@@ -19,6 +21,12 @@ class Module(BaseModule, metaclass=ProgramMeta):
 
     @with_callbacks
     def __call__(self, *args, **kwargs):
+        if settings.track_usage and settings.usage_tracker is None:
+            with track_usage() as usage_tracker:
+                output = self.forward(*args, **kwargs)
+                output.set_lm_usage(usage_tracker.get_total_tokens())
+                return output
+
         return self.forward(*args, **kwargs)
 
     def named_predictors(self):
diff --git a/dspy/teleprompt/simba.py b/dspy/teleprompt/simba.py
index 9d41c563..8554d0b1 100644
--- a/dspy/teleprompt/simba.py
+++ b/dspy/teleprompt/simba.py
@@ -47,7 +47,10 @@ class SIMBA(Teleprompter):
         self.temperature_for_sampling = temperature_for_sampling
         self.temperature_for_candidates = temperature_for_candidates
 
-        self.strategies = [append_a_demo(demo_input_field_maxlen), append_a_rule]
+        if self.max_demos > 0:
+            self.strategies = [append_a_demo(demo_input_field_maxlen), append_a_rule]
+        else:
+            self.strategies = [append_a_rule]
 
     def compile(self, student: dspy.Module, *, trainset: list[dspy.Example], seed: int = 0):
         # Basic checks
@@ -117,7 +120,10 @@ class SIMBA(Teleprompter):
         # Parallel runner
         run_parallel = dspy.Parallel(access_examples=False, num_threads=self.num_threads)
 
+        trial_logs = {}
         for batch_idx in range(self.max_steps):
+            trial_logs[batch_idx] = {}
+
             logger.info(f"Starting batch {batch_idx+1} of {self.max_steps}.")
 
             # STEP 1: Get next batch
@@ -205,12 +211,14 @@ class SIMBA(Teleprompter):
                 name2predictor = {}
                 num_demos_list = []
 
+                max_demos_tmp = self.max_demos if self.max_demos > 0 else 3
+
                 for name, predictor in system_candidate.named_predictors():
                     name2predictor[name] = predictor
                     num_demos_list.append(len(predictor.demos))
 
                 num_demos = max(num_demos_list) if num_demos_list else 0
-                num_demos_to_drop = max(rng_np.poisson(num_demos / self.max_demos), int(num_demos >= self.max_demos))
+                num_demos_to_drop = max(rng_np.poisson(num_demos / max_demos_tmp), int(num_demos >= max_demos_tmp))
                 num_demos_to_drop = min(num_demos_to_drop, num_demos)
                 demos_to_drop = [rng.randrange(num_demos) for _ in range(num_demos_to_drop)]
 
@@ -268,7 +276,7 @@ class SIMBA(Teleprompter):
             if candidate_scores:
                 best_idx_among_candidates = candidate_scores.index(max(candidate_scores))
                 best_program = system_candidates[best_idx_among_candidates]
-                winning_programs.append(best_program)
+                winning_programs.append(best_program.deepcopy())
 
             # STEP 8: Register all new candidate systems in our global pool
             for idx_cand, cand_sys in enumerate(system_candidates):
@@ -276,7 +284,7 @@ class SIMBA(Teleprompter):
                 end = (idx_cand + 1) * self.bsize
                 sys_scores = [outputs[i]["score"] for i in range(start, end)]
                 register_new_program(cand_sys, sys_scores)
-
+            
         M = len(winning_programs) - 1
         N = self.num_candidates + 1
         if M < 1:
@@ -286,7 +294,7 @@ class SIMBA(Teleprompter):
             program_idxs = [round(i * M / (N - 1)) for i in range(N)]
         program_idxs = list(dict.fromkeys(program_idxs))
 
-        candidate_programs = [winning_programs[i] for i in program_idxs]
+        candidate_programs = [winning_programs[i].deepcopy() for i in program_idxs]
         logger.info(f"VALIDATION: Evaluating {len(candidate_programs)} programs on the full trainset.")
         exec_pairs = [(wrap_program(sys, self.metric), ex) for sys in candidate_programs for ex in trainset]
         outputs = run_parallel(exec_pairs)
@@ -298,7 +306,9 @@ class SIMBA(Teleprompter):
             sys_scores = [outputs[i]["score"] for i in range(start, end)]
             avg_score = sum(sys_scores) / len(sys_scores) if sys_scores else 0.0
             scores.append(avg_score)
-
+            if idx_prog != 0:
+                trial_logs[idx_prog-1]["train_score"] = avg_score
+        
         best_idx = scores.index(max(scores)) if scores else 0
         best_program = candidate_programs[best_idx]
         logger.info(
@@ -309,5 +319,6 @@ class SIMBA(Teleprompter):
         # FIXME: Attach all program candidates in decreasing average score to the best program.
         best_program.candidate_programs = candidate_programs
         best_program.winning_programs = winning_programs
+        best_program.trial_logs = trial_logs
 
         return best_program
diff --git a/dspy/teleprompt/utils.py b/dspy/teleprompt/utils.py
index bbf9ac50..954ce560 100644
--- a/dspy/teleprompt/utils.py
+++ b/dspy/teleprompt/utils.py
@@ -5,7 +5,6 @@ import os
 import random
 import shutil
 import sys
-
 import numpy as np
 
 try:
@@ -250,6 +249,46 @@ def setup_logging(log_dir):
     console_handler.setFormatter(console_formatter)
     logger.addHandler(console_handler)
 
+def get_token_usage(model) -> tuple[int, int]:
+    """
+    Extract total input tokens and output tokens from a model's interaction history.
+    Returns (total_input_tokens, total_output_tokens).
+    """
+    if not hasattr(model, "history"):
+        return 0, 0
+
+    input_tokens = []
+    output_tokens = []
+    for interaction in model.history:
+        usage = interaction.get("usage", {})
+        _input_tokens = usage.get("prompt_tokens", 0)
+        _output_tokens = usage.get("completion_tokens", 0)
+        input_tokens.append(_input_tokens)
+        output_tokens.append(_output_tokens)
+
+    total_input_tokens = int(np.sum(input_tokens))
+    total_output_tokens = int(np.sum(output_tokens))
+
+    return total_input_tokens, total_output_tokens
+
+
+def log_token_usage(trial_logs, trial_num, model_dict):
+    """
+    Extract total input and output tokens used by each model and log to trial_logs[trial_num]["token_usage"].
+    """
+
+    token_usage_dict = {}
+
+    for model_name, model in model_dict.items():
+        in_tokens, out_tokens = get_token_usage(model)
+        token_usage_dict[model_name] = {
+            "total_input_tokens": in_tokens,
+            "total_output_tokens": out_tokens
+        }
+
+    # Store token usage info in trial logs
+    trial_logs[trial_num]["token_usage"] = token_usage_dict
+
 
 ### OTHER UTILS ###
 
diff --git a/dspy/utils/parallelizer.py b/dspy/utils/parallelizer.py
index 5541e380..404aab1d 100644
--- a/dspy/utils/parallelizer.py
+++ b/dspy/utils/parallelizer.py
@@ -1,13 +1,14 @@
-import sys
-import tqdm
-import signal
+import contextlib
+import copy
 import logging
+import signal
+import sys
 import threading
-import traceback
 import time
-import contextlib
+import traceback
+from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
 
-from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED
+import tqdm
 
 logger = logging.getLogger(__name__)
 
@@ -27,7 +28,7 @@ class ParallelExecutor:
         Offers isolation between the tasks (dspy.settings) irrespective of whether num_threads == 1 or > 1.
         Handles also straggler timeouts.
         """
-        
+
         self.num_threads = num_threads
         self.max_errors = max_errors
         self.disable_progress_bar = disable_progress_bar
@@ -59,10 +60,7 @@ class ParallelExecutor:
                 if self.provide_traceback:
                     logger.error(f"Error for {item}: {e}\n{traceback.format_exc()}")
                 else:
-                    logger.error(
-                        f"Error for {item}: {e}. "
-                        "Set `provide_traceback=True` for traceback."
-                    )
+                    logger.error(f"Error for {item}: {e}. " "Set `provide_traceback=True` for traceback.")
                 return None
 
         return safe_func
@@ -89,6 +87,9 @@ class ParallelExecutor:
 
             original = thread_local_overrides.overrides
             thread_local_overrides.overrides = parent_overrides.copy()
+            if parent_overrides.get("usage_tracker"):
+                # Usage tracker needs to be deep copied across threads so that each thread tracks its own usage
+                thread_local_overrides.overrides["usage_tracker"] = copy.deepcopy(parent_overrides["usage_tracker"])
 
             try:
                 return index, function(item)
@@ -126,9 +127,7 @@ class ParallelExecutor:
                 submission_counter = 0
 
                 for idx, item in enumerate(data):
-                    f = executor.submit(
-                        worker, parent_overrides, submission_counter, idx, item
-                    )
+                    f = executor.submit(worker, parent_overrides, submission_counter, idx, item)
                     futures_map[f] = (submission_counter, idx, item)
                     futures_set.add(f)
                     submission_counter += 1
@@ -146,9 +145,7 @@ class ParallelExecutor:
                 while futures_set and not self.cancel_jobs.is_set():
                     if all_done():
                         break
-                    done, not_done = wait(
-                        futures_set, timeout=1, return_when=FIRST_COMPLETED
-                    )
+                    done, not_done = wait(futures_set, timeout=1, return_when=FIRST_COMPLETED)
                     for f in done:
                         futures_set.remove(f)
                         try:
diff --git a/dspy/utils/usage_tracker.py b/dspy/utils/usage_tracker.py
new file mode 100644
index 00000000..857b1d02
--- /dev/null
+++ b/dspy/utils/usage_tracker.py
@@ -0,0 +1,69 @@
+"""Usage tracking utilities for DSPy."""
+
+from collections import defaultdict
+from contextlib import contextmanager
+from typing import Any
+
+from dspy.dsp.utils.settings import settings
+
+
+class UsageTracker:
+    """Tracks LM usage data within a context."""
+
+    def __init__(self):
+        # Map of LM name to list of usage entries. For example:
+        # {
+        #     "openai/gpt-4o-mini": [
+        #         {"prompt_tokens": 100, "completion_tokens": 200},
+        #         {"prompt_tokens": 300, "completion_tokens": 400},
+        #     ],
+        # }
+        self.usage_data = defaultdict(list)
+
+    def _flatten_usage_entry(self, usage_entry) -> dict[str, dict[str, Any]]:
+        result = dict(usage_entry)
+        if "completion_tokens_details" in result:
+            result["completion_tokens_details"] = dict(result["completion_tokens_details"])
+        if "prompt_tokens_details" in result:
+            result["prompt_tokens_details"] = dict(result["prompt_tokens_details"])
+        return result
+
+    def _merge_usage_entries(self, usage_entry1, usage_entry2) -> dict[str, dict[str, Any]]:
+        if usage_entry1 is None or len(usage_entry1) == 0:
+            return dict(usage_entry2)
+        if usage_entry2 is None or len(usage_entry2) == 0:
+            return dict(usage_entry1)
+
+        result = dict(usage_entry2)
+        for k, v in usage_entry1.items():
+            if k in result:
+                if isinstance(v, dict):
+                    result[k] = self._merge_usage_entries(result[k], v)
+                else:
+                    result[k] = result[k] or 0
+                    result[k] += v if v else 0
+        return result
+
+    def add_usage(self, lm: str, usage_entry: dict):
+        """Add a usage entry to the tracker."""
+        if len(usage_entry) > 0:
+            self.usage_data[lm].append(self._flatten_usage_entry(usage_entry))
+
+    def get_total_tokens(self) -> dict[str, dict[str, Any]]:
+        """Calculate total tokens from all tracked usage."""
+        total_usage_by_lm = {}
+        for lm, usage_entries in self.usage_data.items():
+            total_usage = {}
+            for usage_entry in usage_entries:
+                total_usage = self._merge_usage_entries(total_usage, usage_entry)
+            total_usage_by_lm[lm] = total_usage
+        return total_usage_by_lm
+
+
+@contextmanager
+def track_usage():
+    """Context manager for tracking LM usage."""
+    tracker = UsageTracker()
+
+    with settings.context(usage_tracker=tracker):
+        yield tracker
diff --git a/poetry.lock b/poetry.lock
index babcf669..24851d21 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -384,7 +384,7 @@ files = [
     {file = "async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c"},
     {file = "async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3"},
 ]
-markers = {main = "python_version <= \"3.10\"", dev = "sys_platform != \"win32\" and python_full_version < \"3.11.3\" or python_version <= \"3.10\""}
+markers = {main = "python_version == \"3.9\" or python_version == \"3.10\"", dev = "sys_platform != \"win32\" and python_full_version < \"3.11.3\" and python_version == \"3.11\" or python_version == \"3.10\" or python_version == \"3.9\""}
 
 [[package]]
 name = "asyncer"
@@ -719,8 +719,8 @@ files = [
 jmespath = ">=0.7.1,<2.0.0"
 python-dateutil = ">=2.1,<3.0.0"
 urllib3 = [
-    {version = ">=1.25.4,<2.2.0 || >2.2.0,<3", markers = "python_version >= \"3.10\""},
     {version = ">=1.25.4,<1.27", markers = "python_version < \"3.10\""},
+    {version = ">=1.25.4,<2.2.0 || >2.2.0,<3", markers = "python_version >= \"3.10\""},
 ]
 
 [package.extras]
@@ -1217,9 +1217,9 @@ isort = ">=4.3.21,<6.0"
 jinja2 = ">=2.10.1,<4.0"
 packaging = "*"
 pydantic = [
+    {version = ">=1.5.1,<2.4.0 || >2.4.0,<3.0", extras = ["email"], markers = "python_version < \"3.10\""},
     {version = ">=1.10.0,<2.0.0 || >2.0.0,<2.0.1 || >2.0.1,<2.4.0 || >2.4.0,<3.0", extras = ["email"], markers = "python_version >= \"3.12\" and python_version < \"4.0\""},
     {version = ">=1.10.0,<2.4.0 || >2.4.0,<3.0", extras = ["email"], markers = "python_version == \"3.11\""},
-    {version = ">=1.5.1,<2.4.0 || >2.4.0,<3.0", extras = ["email"], markers = "python_version < \"3.10\""},
     {version = ">=1.9.0,<2.4.0 || >2.4.0,<3.0", extras = ["email"], markers = "python_version == \"3.10\""},
 ]
 pyyaml = ">=6.0.1"
@@ -1499,7 +1499,7 @@ description = "Backport of PEP 654 (exception groups)"
 optional = false
 python-versions = ">=3.7"
 groups = ["main", "dev"]
-markers = "python_version <= \"3.10\""
+markers = "python_version == \"3.9\" or python_version == \"3.10\""
 files = [
     {file = "exceptiongroup-1.2.2-py3-none-any.whl", hash = "sha256:3111b9d131c238bec2f8f516e123e14ba243563fb135d3fe885990585aa7795b"},
     {file = "exceptiongroup-1.2.2.tar.gz", hash = "sha256:47c2edf7c6738fafb49fd34290706d1a1a2f4d1c6df275526b62cbb4aa5393cc"},
@@ -1583,13 +1583,13 @@ huggingface-hub = ">=0.20,<1.0"
 loguru = ">=0.7.2,<0.8.0"
 mmh3 = ">=4.1.0,<6.0.0"
 numpy = [
+    {version = ">=1.21,<2.1.0", markers = "python_version < \"3.10\""},
     {version = ">=1.26", markers = "python_version == \"3.12\""},
     {version = ">=1.21", markers = "python_version >= \"3.10\" and python_version < \"3.12\""},
-    {version = ">=1.21,<2.1.0", markers = "python_version < \"3.10\""},
 ]
 onnxruntime = [
-    {version = ">=1.17.0,<1.20.0 || >1.20.0", markers = "python_version >= \"3.10\" and python_version < \"3.13\""},
     {version = ">=1.17.0,<1.20.0", markers = "python_version < \"3.10\""},
+    {version = ">=1.17.0,<1.20.0 || >1.20.0", markers = "python_version >= \"3.10\" and python_version < \"3.13\""},
 ]
 pillow = ">=10.3.0,<12.0.0"
 py-rust-stemmers = ">=0.1.0,<0.2.0"
@@ -2515,7 +2515,7 @@ description = "IPython: Productive Interactive Computing"
 optional = false
 python-versions = ">=3.9"
 groups = ["main", "dev"]
-markers = "python_version < \"3.10\""
+markers = "python_version == \"3.9\""
 files = [
     {file = "ipython-8.18.1-py3-none-any.whl", hash = "sha256:e8267419d72d81955ec1177f8a29aaa90ac80ad647499201119e2f05e99aa397"},
     {file = "ipython-8.18.1.tar.gz", hash = "sha256:ca6f079bb33457c66e233e4580ebfc4128855b4cf6370dddd73842a9563e8a27"},
@@ -3934,7 +3934,7 @@ description = "Python package for creating and manipulating graphs and networks"
 optional = false
 python-versions = ">=3.9"
 groups = ["dev"]
-markers = "python_version < \"3.10\""
+markers = "python_version == \"3.9\""
 files = [
     {file = "networkx-3.2.1-py3-none-any.whl", hash = "sha256:f18c69adc97877c42332c170849c96cefa91881c99a7cb3e95b7c659ebdc1ec2"},
     {file = "networkx-3.2.1.tar.gz", hash = "sha256:9f1bb5cf3409bf324e0a722c20bdb4c20ee39bf1c30ce8ae499c8502b0b5e0c6"},
@@ -4253,7 +4253,7 @@ description = "ONNX Runtime is a runtime accelerator for Machine Learning models
 optional = true
 python-versions = "*"
 groups = ["main"]
-markers = "python_version < \"3.10\""
+markers = "python_version == \"3.9\""
 files = [
     {file = "onnxruntime-1.19.2-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:84fa57369c06cadd3c2a538ae2a26d76d583e7c34bdecd5769d71ca5c0fc750e"},
     {file = "onnxruntime-1.19.2-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bdc471a66df0c1cdef774accef69e9f2ca168c851ab5e4f2f3341512c7ef4666"},
@@ -4297,7 +4297,7 @@ description = "ONNX Runtime is a runtime accelerator for Machine Learning models
 optional = true
 python-versions = "*"
 groups = ["main"]
-markers = "python_version < \"3.10\""
+markers = "python_version == \"3.9\""
 files = [
     {file = "onnxruntime-1.20.1-cp310-cp310-macosx_13_0_universal2.whl", hash = "sha256:e50ba5ff7fed4f7d9253a6baf801ca2883cc08491f9d32d78a80da57256a5439"},
     {file = "onnxruntime-1.20.1-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7b2908b50101a19e99c4d4e97ebb9905561daf61829403061c1adc1b588bc0de"},
@@ -4778,9 +4778,9 @@ files = [
 
 [package.dependencies]
 numpy = [
+    {version = ">=1.22.4", markers = "python_version < \"3.11\""},
     {version = ">=1.26.0", markers = "python_version >= \"3.12\""},
     {version = ">=1.23.2", markers = "python_version == \"3.11\""},
-    {version = ">=1.22.4", markers = "python_version < \"3.11\""},
 ]
 python-dateutil = ">=2.8.2"
 pytz = ">=2020.1"
@@ -4846,7 +4846,7 @@ description = "Pexpect allows easy control of interactive console applications."
 optional = false
 python-versions = "*"
 groups = ["main", "dev"]
-markers = "sys_platform != \"win32\" and sys_platform != \"emscripten\" or python_version < \"3.10\" and sys_platform != \"win32\""
+markers = "python_version == \"3.9\" and sys_platform != \"win32\" or sys_platform != \"win32\" and sys_platform != \"emscripten\""
 files = [
     {file = "pexpect-4.9.0-py2.py3-none-any.whl", hash = "sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523"},
     {file = "pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f"},
@@ -5252,7 +5252,7 @@ description = "Run a subprocess in a pseudo terminal"
 optional = false
 python-versions = "*"
 groups = ["main", "dev"]
-markers = "sys_platform != \"win32\" and sys_platform != \"emscripten\" or python_version < \"3.10\" and sys_platform != \"win32\""
+markers = "python_version == \"3.9\" and sys_platform != \"win32\" or sys_platform != \"win32\" and sys_platform != \"emscripten\""
 files = [
     {file = "ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35"},
     {file = "ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220"},
@@ -6029,7 +6029,7 @@ files = [
     {file = "pywin32-310-cp39-cp39-win32.whl", hash = "sha256:851c8d927af0d879221e616ae1f66145253537bbdd321a77e8ef701b443a9a1a"},
     {file = "pywin32-310-cp39-cp39-win_amd64.whl", hash = "sha256:96867217335559ac619f00ad70e513c0fcf84b8a3af9fc2bba3b59b97da70475"},
 ]
-markers = {main = "sys_platform == \"win32\" and platform_python_implementation != \"PyPy\" or platform_system == \"Windows\"", dev = "sys_platform == \"win32\" and platform_python_implementation != \"PyPy\""}
+markers = {main = "(sys_platform == \"win32\" or platform_system == \"Windows\") and (platform_python_implementation != \"PyPy\" or platform_system == \"Windows\")", dev = "sys_platform == \"win32\" and platform_python_implementation != \"PyPy\""}
 
 [[package]]
 name = "pyyaml"
@@ -6232,9 +6232,9 @@ grpcio = ">=1.41.0"
 grpcio-tools = ">=1.41.0"
 httpx = {version = ">=0.20.0", extras = ["http2"]}
 numpy = [
+    {version = ">=1.21,<2.1.0", markers = "python_version < \"3.10\""},
     {version = ">=1.26", markers = "python_version == \"3.12\""},
     {version = ">=1.21", markers = "python_version >= \"3.10\" and python_version < \"3.12\""},
-    {version = ">=1.21,<2.1.0", markers = "python_version < \"3.10\""},
 ]
 portalocker = ">=2.7.0,<3.0.0"
 pydantic = ">=1.10.8"
@@ -6746,7 +6746,7 @@ unleash = ["UnleashClient (>=6.0.1)"]
 name = "setuptools"
 version = "77.0.3"
 description = "Easily download, build, install, upgrade, and uninstall Python packages"
-optional = false
+optional = true
 python-versions = ">=3.9"
 groups = ["main", "dev"]
 files = [
@@ -7343,7 +7343,7 @@ description = "Python Library for Tom's Obvious, Minimal Language"
 optional = false
 python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*"
 groups = ["dev"]
-markers = "python_version <= \"3.10\""
+markers = "python_version == \"3.9\" or python_version == \"3.10\""
 files = [
     {file = "toml-0.10.2-py2.py3-none-any.whl", hash = "sha256:806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b"},
     {file = "toml-0.10.2.tar.gz", hash = "sha256:b3bda1d108d5dd99f4a20d24d9c348e91c4db7ab1b749200bded2f839ccbe68f"},
@@ -7356,7 +7356,7 @@ description = "A lil' TOML parser"
 optional = false
 python-versions = ">=3.8"
 groups = ["main", "dev"]
-markers = "python_version <= \"3.10\""
+markers = "python_version == \"3.9\" or python_version == \"3.10\""
 files = [
     {file = "tomli-2.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:678e4fa69e4575eb77d103de3df8a895e1591b48e740211bd1067378c69e8249"},
     {file = "tomli-2.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:023aa114dd824ade0100497eb2318602af309e5a55595f76b626d6d9f3b7b0a6"},
@@ -7630,7 +7630,7 @@ files = [
     {file = "typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d"},
     {file = "typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8"},
 ]
-markers = {doc = "python_version <= \"3.10\""}
+markers = {doc = "python_version == \"3.9\" or python_version == \"3.10\""}
 
 [[package]]
 name = "tzdata"
@@ -7759,7 +7759,7 @@ description = "HTTP library with thread-safe connection pooling, file post, and
 optional = false
 python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7"
 groups = ["main", "dev", "doc"]
-markers = "python_version < \"3.10\""
+markers = "python_version == \"3.9\""
 files = [
     {file = "urllib3-1.26.20-py2.py3-none-any.whl", hash = "sha256:0ed14ccfbf1c30a9072c7ca157e4319b70d65f623e91e7b32fadb2853431016e"},
     {file = "urllib3-1.26.20.tar.gz", hash = "sha256:40c2dc0c681e47eb8f90e7e27bf6ff7df2e677421fd46756da1161c39ca70d32"},
@@ -8561,4 +8561,4 @@ weaviate = ["weaviate-client"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<3.13"
-content-hash = "eb10091cda3277c5fde78efd70b100009ddb5798a06bb52faf161025eea4d347"
+content-hash = "997d49fb24a6bf3a51f96124232f7ea34a51460ff1e5f7fe1c7213f67ccaaadd"
diff --git a/pyproject.toml b/pyproject.toml
index d50f491f..c7ea5aa7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,8 @@ dependencies = [
     "cachetools>=5.5.0",
     "cloudpickle>=3.0.0",
     "rich>=13.7.1",
-    "numpy>=1.26.0",
+    "numpy>=1.26.0,<2.2; python_version == '3.9'",
+    "numpy>=1.26.0; python_version >= '3.10'",
 ]
 
 [tool.setuptools.packages.find]
@@ -122,7 +123,10 @@ asyncer = "0.0.8"
 cachetools = "^5.5.0"
 cloudpickle = "^3.0.0"
 rich = "^13.7.1"
-numpy = ">=1.26.0"
+numpy = [
+    { version = ">=1.26.0,<2.2", markers = "python_version == '3.9'" },
+    { version = ">=1.26.0", markers = "python_version >= '3.10'" }
+]
 
 # Optional dependencies (now declared inline)
 anthropic = { version = ">=0.18.0,<1.0.0", optional = true }


Test patch:
diff --git a/tests/primitives/test_module.py b/tests/primitives/test_module.py
index 1d458edd..88621384 100644
--- a/tests/primitives/test_module.py
+++ b/tests/primitives/test_module.py
@@ -3,6 +3,8 @@ import threading
 from dspy.utils.dummies import DummyLM
 import logging
 from unittest.mock import patch
+import pytest
+import os
 
 
 def test_deepcopy_basic():
@@ -161,3 +163,83 @@ def test_load_with_version_mismatch(tmp_path):
         # Clean up: restore original level and remove handler
         logger.setLevel(original_level)
         logger.removeHandler(handler)
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skip the test if OPENAI_API_KEY is not set.")
+def test_single_module_call_with_usage_tracker():
+    dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=False), track_usage=True)
+
+    predict = dspy.ChainOfThought("question -> answer")
+    output = predict(question="What is the capital of France?")
+
+    lm_usage = output.get_lm_usage()
+    assert len(lm_usage) == 1
+    assert lm_usage["openai/gpt-4o-mini"]["prompt_tokens"] > 0
+    assert lm_usage["openai/gpt-4o-mini"]["completion_tokens"] > 0
+    assert lm_usage["openai/gpt-4o-mini"]["total_tokens"] > 0
+
+    # Test no usage being tracked when cache is enabled
+    dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=True), track_usage=True)
+    for _ in range(2):
+        output = predict(question="What is the capital of France?")
+
+    assert len(output.get_lm_usage()) == 0
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skip the test if OPENAI_API_KEY is not set.")
+def test_multi_module_call_with_usage_tracker():
+    dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=False), track_usage=True)
+
+    class MyProgram(dspy.Module):
+        def __init__(self):
+            self.predict1 = dspy.ChainOfThought("question -> answer")
+            self.predict2 = dspy.ChainOfThought("question, answer -> score")
+
+        def __call__(self, question: str) -> str:
+            answer = self.predict1(question=question)
+            score = self.predict2(question=question, answer=answer)
+            return score
+
+    program = MyProgram()
+    output = program(question="What is the capital of France?")
+
+    lm_usage = output.get_lm_usage()
+    assert len(lm_usage) == 1
+    assert lm_usage["openai/gpt-4o-mini"]["prompt_tokens"] > 0
+    assert lm_usage["openai/gpt-4o-mini"]["prompt_tokens"] > 0
+    assert lm_usage["openai/gpt-4o-mini"]["completion_tokens"] > 0
+    assert lm_usage["openai/gpt-4o-mini"]["total_tokens"] > 0
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skip the test if OPENAI_API_KEY is not set.")
+def test_usage_tracker_in_parallel():
+    class MyProgram(dspy.Module):
+        def __init__(self, lm):
+            self.lm = lm
+            self.predict1 = dspy.ChainOfThought("question -> answer")
+            self.predict2 = dspy.ChainOfThought("question, answer -> score")
+
+        def __call__(self, question: str) -> str:
+            with dspy.settings.context(lm=self.lm):
+                answer = self.predict1(question=question)
+                score = self.predict2(question=question, answer=answer)
+                return score
+
+    dspy.settings.configure(track_usage=True)
+    program1 = MyProgram(lm=dspy.LM("openai/gpt-4o-mini", cache=False))
+    program2 = MyProgram(lm=dspy.LM("openai/gpt-3.5-turbo", cache=False))
+
+    parallelizer = dspy.Parallel()
+
+    results = parallelizer(
+        [
+            (program1, {"question": "What is the meaning of life?"}),
+            (program2, {"question": "why did a chicken cross the kitchen?"}),
+        ]
+    )
+
+    assert results[0].get_lm_usage() is not None
+    assert results[1].get_lm_usage() is not None
+
+    assert results[0].get_lm_usage().keys() == set(["openai/gpt-4o-mini"])
+    assert results[1].get_lm_usage().keys() == set(["openai/gpt-3.5-turbo"])
diff --git a/tests/utils/test_usage_tracker.py b/tests/utils/test_usage_tracker.py
new file mode 100644
index 00000000..a9accef6
--- /dev/null
+++ b/tests/utils/test_usage_tracker.py
@@ -0,0 +1,159 @@
+import pytest
+import dspy
+from unittest.mock import Mock, patch, MagicMock
+from dspy.utils.usage_tracker import UsageTracker, track_usage
+from dspy.utils.dummies import DummyLM
+import os
+
+
+def test_add_usage_entry():
+    """Test adding usage entries to the tracker."""
+    tracker = UsageTracker()
+
+    # Test with a single usage entry
+    usage_entry = {
+        "prompt_tokens": 1117,
+        "completion_tokens": 46,
+        "total_tokens": 1163,
+        "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0},
+        "completion_tokens_details": {
+            "reasoning_tokens": 0,
+            "audio_tokens": 0,
+            "accepted_prediction_tokens": 0,
+            "rejected_prediction_tokens": 0,
+        },
+    }
+
+    tracker.add_usage("gpt-4o-mini", usage_entry)
+    assert len(tracker.usage_data["gpt-4o-mini"]) == 1
+    assert tracker.usage_data["gpt-4o-mini"][0] == usage_entry
+
+
+def test_get_total_tokens():
+    """Test calculating total tokens from usage entries."""
+    tracker = UsageTracker()
+
+    # Add multiple usage entries for the same model
+    usage_entries = [
+        {
+            "prompt_tokens": 1117,
+            "completion_tokens": 46,
+            "total_tokens": 1163,
+            "prompt_tokens_details": {"cached_tokens": 200, "audio_tokens": 50},
+            "completion_tokens_details": {
+                "reasoning_tokens": 20,
+                "audio_tokens": 10,
+                "accepted_prediction_tokens": 16,
+                "rejected_prediction_tokens": 0,
+            },
+        },
+        {
+            "prompt_tokens": 800,
+            "completion_tokens": 100,
+            "total_tokens": 900,
+            "prompt_tokens_details": {"cached_tokens": 300, "audio_tokens": 0},
+            "completion_tokens_details": {
+                "reasoning_tokens": 50,
+                "audio_tokens": 0,
+                "accepted_prediction_tokens": 40,
+                "rejected_prediction_tokens": 10,
+            },
+        },
+        {
+            "prompt_tokens": 500,
+            "completion_tokens": 80,
+            "total_tokens": 580,
+            "prompt_tokens_details": {"cached_tokens": 100, "audio_tokens": 25},
+            "completion_tokens_details": {
+                "reasoning_tokens": 30,
+                "audio_tokens": 15,
+                "accepted_prediction_tokens": 25,
+                "rejected_prediction_tokens": 10,
+            },
+        },
+    ]
+
+    for entry in usage_entries:
+        tracker.add_usage("gpt-4o-mini", entry)
+
+    total_usage = tracker.get_total_tokens()
+    assert "gpt-4o-mini" in total_usage
+    assert total_usage["gpt-4o-mini"]["prompt_tokens"] == 2417  # 1117 + 800 + 500
+    assert total_usage["gpt-4o-mini"]["completion_tokens"] == 226  # 46 + 100 + 80
+    assert total_usage["gpt-4o-mini"]["total_tokens"] == 2643  # 1163 + 900 + 580
+    assert total_usage["gpt-4o-mini"]["prompt_tokens_details"]["cached_tokens"] == 600  # 200 + 300 + 100
+    assert total_usage["gpt-4o-mini"]["prompt_tokens_details"]["audio_tokens"] == 75  # 50 + 0 + 25
+    assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["reasoning_tokens"] == 100  # 20 + 50 + 30
+    assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["audio_tokens"] == 25  # 10 + 0 + 15
+    assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["accepted_prediction_tokens"] == 81  # 16 + 40 + 25
+    assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["rejected_prediction_tokens"] == 20  # 0 + 10 + 10
+
+
+def test_track_usage_with_multiple_models():
+    """Test tracking usage across multiple models."""
+    tracker = UsageTracker()
+
+    # Add usage entries for different models
+    usage_entries = [
+        {
+            "model": "gpt-4o-mini",
+            "usage": {
+                "prompt_tokens": 1117,
+                "completion_tokens": 46,
+                "total_tokens": 1163,
+                "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0},
+                "completion_tokens_details": {
+                    "reasoning_tokens": 0,
+                    "audio_tokens": 0,
+                    "accepted_prediction_tokens": 0,
+                    "rejected_prediction_tokens": 0,
+                },
+            },
+        },
+        {
+            "model": "gpt-3.5-turbo",
+            "usage": {
+                "prompt_tokens": 800,
+                "completion_tokens": 100,
+                "total_tokens": 900,
+                "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0},
+                "completion_tokens_details": {
+                    "reasoning_tokens": 0,
+                    "audio_tokens": 0,
+                    "accepted_prediction_tokens": 0,
+                    "rejected_prediction_tokens": 0,
+                },
+            },
+        },
+    ]
+
+    for entry in usage_entries:
+        tracker.add_usage(entry["model"], entry["usage"])
+
+    total_usage = tracker.get_total_tokens()
+    assert "gpt-4o-mini" in total_usage
+    assert "gpt-3.5-turbo" in total_usage
+    assert total_usage["gpt-4o-mini"]["total_tokens"] == 1163
+    assert total_usage["gpt-3.5-turbo"]["total_tokens"] == 900
+
+
+@pytest.mark.skipif(
+    not os.getenv("OPENAI_API_KEY"),
+    reason="Skip the test if OPENAI_API_KEY is not set.",
+)
+def test_track_usage_context_manager():
+    lm = dspy.LM("openai/gpt-4o-mini", cache=False)
+    dspy.settings.configure(lm=lm)
+
+    predict = dspy.ChainOfThought("question -> answer")
+    with track_usage() as tracker:
+        predict(question="What is the capital of France?")
+        predict(question="What is the capital of Italy?")
+
+    assert len(tracker.usage_data) > 0
+    assert len(tracker.usage_data["openai/gpt-4o-mini"]) == 2
+
+    total_usage = tracker.get_total_tokens()
+    assert "openai/gpt-4o-mini" in total_usage
+    assert len(total_usage.keys()) == 1
+    assert isinstance(total_usage["openai/gpt-4o-mini"], dict)