feat: add token limit handling

Author: andrewklatzke

packages/optimization/src/ldai_optimizer/client.py

@@ -157,6 +157,7 @@ def __init__(self, ldClient: LDAIClient) -> None:
         self._last_succeeded_context: Optional[OptimizationContext] = None
         self._last_optimization_result_id: Optional[str] = None
         self._initial_tool_keys: List[str] = []
+        self._total_token_usage: int = 0
 
         if os.environ.get("LAUNCHDARKLY_API_KEY"):
             self._has_api_key = True
@@ -966,12 +967,14 @@ async def _run_ground_truth_optimization(
             on_passing_result=gt_options.on_passing_result,
             on_failing_result=gt_options.on_failing_result,
             on_status_update=gt_options.on_status_update,
+            token_limit=gt_options.token_limit,
         )
         self._options = bridge
         self._agent_config = agent_config
         self._last_run_succeeded = False
         self._last_succeeded_context = None
         self._last_optimization_result_id = None
+        self._total_token_usage = 0
         self._initialize_class_members_from_config(agent_config)
 
         # Seed from the first model choice on the first iteration
@@ -1036,6 +1039,26 @@ async def _run_ground_truth_optimization(
                     linear_iter,
                     expected_response=sample.expected_response,
                 )
+                self._accumulate_tokens(optimize_context)
+                if self._is_token_limit_exceeded():
+                    logger.error(
+                        "[GT Attempt %d] -> Token limit exceeded on sample %d (total=%d)",
+                        attempt,
+                        i + 1,
+                        self._total_token_usage,
+                    )
+                    attempt_results.append(optimize_context)
+                    self._last_run_succeeded = False
+                    self._last_succeeded_context = None
+                    self._safe_status_update("failure", optimize_context, linear_iter)
+                    if self._options.on_failing_result:
+                        try:
+                            self._options.on_failing_result(optimize_context)
+                        except Exception:
+                            logger.exception(
+                                "[GT Attempt %d] -> on_failing_result callback failed", attempt
+                            )
+                    return attempt_results
 
                 # Per-sample pass/fail check
                 if self._options.on_turn is not None:
@@ -1681,6 +1704,7 @@ def _persist_and_forward(
                 on_passing_result=options.on_passing_result,
                 on_failing_result=options.on_failing_result,
                 on_status_update=_persist_and_forward,
+                token_limit=config.get("tokenLimit"),
             )
 
         variable_choices: List[Dict[str, Any]] = config["variableChoices"] or [{}]
@@ -1700,6 +1724,7 @@ def _persist_and_forward(
             on_passing_result=options.on_passing_result,
             on_failing_result=options.on_failing_result,
             on_status_update=_persist_and_forward,
+            token_limit=config.get("tokenLimit"),
         )
 
     async def _execute_agent_turn(
@@ -1780,6 +1805,31 @@ async def _execute_agent_turn(
 
         return result_ctx
 
+    def _accumulate_tokens(self, optimize_context: OptimizationContext) -> None:
+        """Add token usage from a completed turn to the running total.
+
+        Sums the agent's token usage and each judge's token usage from the given
+        context and adds them to ``_total_token_usage``.
+
+        :param optimize_context: The completed turn context containing usage data.
+        """
+        if optimize_context.usage is not None:
+            self._total_token_usage += optimize_context.usage.total or 0
+        for judge_result in optimize_context.scores.values():
+            if judge_result.usage is not None:
+                self._total_token_usage += judge_result.usage.total or 0
+
+    def _is_token_limit_exceeded(self) -> bool:
+        """Return True if the accumulated token usage has met or exceeded the configured limit.
+
+        Returns False when no token limit is set so callers can use this as a
+        simple guard without needing to check for ``None`` themselves.
+
+        :return: True if token limit is set and ``_total_token_usage >= token_limit``.
+        """
+        limit: Optional[int] = getattr(self._options, "token_limit", None)
+        return limit is not None and self._total_token_usage >= limit
+
     def _evaluate_response(self, optimize_context: OptimizationContext) -> bool:
         """
         Determine whether the current iteration's scores meet all judge thresholds.
@@ -2091,6 +2141,15 @@ async def _run_validation_phase(
             )
             self._safe_status_update("generating", val_ctx, val_iter)
             val_ctx = await self._execute_agent_turn(val_ctx, val_iter)
+            self._accumulate_tokens(val_ctx)
+            if self._is_token_limit_exceeded():
+                logger.error(
+                    "[Validation %d/%d] -> Token limit exceeded (total=%d)",
+                    i + 1,
+                    validation_count,
+                    self._total_token_usage,
+                )
+                return False, val_ctx
 
             if options.on_turn is not None:
                 try:
@@ -2147,6 +2206,7 @@ async def _run_optimization(
         self._last_run_succeeded = False
         self._last_succeeded_context = None
         self._last_optimization_result_id = None
+        self._total_token_usage = 0
         self._initialize_class_members_from_config(agent_config)
 
         # If the LD flag doesn't carry a model name, seed from the first model choice
@@ -2192,6 +2252,14 @@ async def _run_optimization(
             optimize_context = await self._execute_agent_turn(
                 optimize_context, iteration
             )
+            self._accumulate_tokens(optimize_context)
+            if self._is_token_limit_exceeded():
+                logger.error(
+                    "[Iteration %d] -> Token limit exceeded (total=%d)",
+                    iteration,
+                    self._total_token_usage,
+                )
+                return self._handle_failure(optimize_context, iteration)
 
             # Manual path: on_turn callback gives caller full control over pass/fail
             if self._options.on_turn is not None:
@@ -2229,6 +2297,8 @@ async def _run_optimization(
                 )
                 if all_valid:
                     return self._handle_success(optimize_context, iteration)
+                if self._is_token_limit_exceeded():
+                    return self._handle_failure(last_ctx, iteration)
                 # Validation failed — treat as a normal failed attempt.
                 # Use optimize_context (the main iteration) for terminal API events so
                 # the persisted record's completionResponse and userInput stay aligned.

packages/optimization/src/ldai_optimizer/dataclasses.py

@@ -346,6 +346,7 @@ class OptimizationOptions:
     on_failing_result: Optional[Callable[[OptimizationContext], None]] = None
     # called to provide status updates during the optimization flow
     on_status_update: Optional[Callable[[_StatusLiteral, OptimizationContext], None]] = None
+    token_limit: Optional[int] = None  # stop the run when total token usage reaches this value
 
     def __post_init__(self):
         """Validate required options."""
@@ -433,6 +434,7 @@ class GroundTruthOptimizationOptions:
     project_key: Optional[str] = None  # required when auto_commit=True
     output_key: Optional[str] = None   # variation key/name; auto-generated if omitted
     base_url: Optional[str] = None  # override to target a non-default LD instance
+    token_limit: Optional[int] = None  # stop the run when total token usage reaches this value
 
     def __post_init__(self):
         """Validate required options."""

packages/optimization/src/ldai_optimizer/ld_api_client.py

@@ -89,6 +89,7 @@ class AgentOptimizationConfig(_AgentOptimizationConfigRequired, total=False):
 
     groundTruthResponses: List[str]
     metricKey: str
+    tokenLimit: int
 
 
 # ---------------------------------------------------------------------------