Repository: bespokelabsai/curator

Issue:
Batch progress bar is sometimes confusing 
The reason why it looks like it is hanging is because the progress bar shows the # of completed requests are reported by OpenAI in each batch. However, the batch often spends a lot of time in the 'finalizing' status before it is ready to actually download. So the program shows 100% responses but is waiting for the 'completed' batch status.

Discussion:

Patch:
diff --git a/src/bespokelabs/curator/request_processor/base_request_processor.py b/src/bespokelabs/curator/request_processor/base_request_processor.py
index ad7bf91..dbfc8fe 100644
--- a/src/bespokelabs/curator/request_processor/base_request_processor.py
+++ b/src/bespokelabs/curator/request_processor/base_request_processor.py
@@ -96,6 +96,7 @@ class BaseRequestProcessor(ABC):
         """
         self.prompt_formatter = prompt_formatter
         self.working_dir = working_dir
+        self.total_requests = len(dataset) if dataset is not None else 1
 
         # load from already completed dataset
         output_dataset = self.attempt_loading_cached_dataset(parse_func_hash)
diff --git a/src/bespokelabs/curator/request_processor/batch/base_batch_request_processor.py b/src/bespokelabs/curator/request_processor/batch/base_batch_request_processor.py
index 4731739..2c8001c 100644
--- a/src/bespokelabs/curator/request_processor/batch/base_batch_request_processor.py
+++ b/src/bespokelabs/curator/request_processor/batch/base_batch_request_processor.py
@@ -5,7 +5,7 @@ import os
 from abc import abstractmethod
 from typing import Optional
 
-from tqdm import tqdm
+from litellm import model_cost
 
 from bespokelabs.curator.request_processor.base_request_processor import BaseRequestProcessor
 from bespokelabs.curator.request_processor.config import BatchRequestProcessorConfig
@@ -14,6 +14,7 @@ from bespokelabs.curator.status_tracker.batch_status_tracker import BatchStatusT
 from bespokelabs.curator.types.generic_batch import GenericBatch, GenericBatchRequestCounts, GenericBatchStatus
 from bespokelabs.curator.types.generic_request import GenericRequest
 from bespokelabs.curator.types.generic_response import GenericResponse
+from bespokelabs.curator.types.token_usage import TokenUsage
 
 logger = logging.getLogger(__name__)
 
@@ -37,7 +38,6 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
         prompt_formatter (PromptFormatter): Formatter for structuring prompts.
         tracker (BatchStatusTracker): Tracks status of submitted batches.
         batch_objects_file (str): Path to file storing batch objects.
-        batch_submit_pbar (tqdm | None): Progress bar for batch submission.
         semaphore (asyncio.Semaphore): Controls concurrent batch operations.
 
     Note:
@@ -52,6 +52,7 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
             config: Configuration object containing batch processing parameters.
         """
         super().__init__(config)
+        self._tracker_console = None  # Add this line to store console for testing
 
     @property
     def backend(self) -> str:
@@ -89,12 +90,12 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
         self.semaphore = asyncio.Semaphore(self.max_concurrent_batch_operations)
         self._batch_objects_file_lock = asyncio.Lock()
         self.batch_objects_file = os.path.join(self.working_dir, "batch_objects.jsonl")
-        self.batch_submit_pbar: tqdm | None = None
-        self.request_pbar: tqdm | None = None
+
+        self._attempt_loading_batch_status_tracker(generic_request_files)
 
         run_in_event_loop(self.submit_batches_from_request_files(generic_request_files))
-        logger.info(f"Submitted batches. These can be viewed in the web dashboard: {self.web_dashboard}")
         run_in_event_loop(self.poll_and_process_batches())
+        self.tracker.stop_tracker()
 
     @property
     @abstractmethod
@@ -281,10 +282,21 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
         if os.path.exists(self.batch_objects_file):
             with open(self.batch_objects_file, "r") as f:
                 self.tracker = BatchStatusTracker.model_validate_json(f.read())
-            logger.info(f"Loaded existing tracker from {self.batch_objects_file}:\n{self.tracker}")
+            logger.info(f"Loaded existing tracker from {self.batch_objects_file}")
         else:
             self.tracker = BatchStatusTracker(unsubmitted_request_files=set(request_files))
 
+        self.tracker.model = self.prompt_formatter.model_name
+        self.tracker.n_total_requests = self.total_requests
+
+        # Set cost information if available
+        if self.prompt_formatter.model_name in model_cost:
+            # Batch requests are 50% cheaper
+            self.tracker.input_cost_per_million = (model_cost[self.prompt_formatter.model_name]["input_cost_per_token"] * 1_000_000) * 0.5
+            self.tracker.output_cost_per_million = (model_cost[self.prompt_formatter.model_name]["output_cost_per_token"] * 1_000_000) * 0.5
+        # Start the tracker with the console from constructor
+        self.tracker.start_tracker(self._tracker_console)
+
     async def update_batch_objects_file(self):
         """Update batch objects file with current tracker state.
 
@@ -325,7 +337,7 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
         # Join requests with newlines and encode to bytes for upload
         file_content = "\n".join(json.dumps(r) for r in api_specific_requests).encode()
         file_content_size = len(file_content)
-        logger.debug(f"Batch file content size: {file_content_size / (1024*1024):.2f} MB ({file_content_size:,} bytes)")
+        logger.debug(f"Batch file content size: {file_content_size / (1024 * 1024):.2f} MB ({file_content_size:,} bytes)")
         if file_content_size > self.max_bytes_per_batch:
             raise ValueError(
                 f"Batch file content size {file_content_size:,} bytes "
@@ -361,7 +373,6 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
 
         self.tracker.mark_as_submitted(batch, len(requests))
         await self.update_batch_objects_file()
-        self.batch_submit_pbar.update(1)
 
     def requests_from_generic_request_file(self, request_file: str, completed_request_ids: set[int]) -> list[dict]:
         """Reads and converts generic requests from a file into API-specific request format.
@@ -406,18 +417,25 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
             - Writes GenericResponse objects in JSONL format
             - Calculates costs with batch discount
             - Handles failed requests with error details
+            - Updates tracker with token usage and cost stats
         """
         request_file = batch.request_file
         request_dir = os.path.dirname(request_file)
         request_filename = os.path.basename(request_file)
         response_filename = request_filename.replace("requests_", "responses_")
         response_file = os.path.join(request_dir, response_filename)
+
+        # Load generic requests from request file
         generic_request_map = {}
         with open(request_file, "r") as f:
             for line in f:
                 generic_request = GenericRequest.model_validate_json(line)
                 generic_request_map[generic_request.original_row_idx] = generic_request
 
+        # Track total token usage and cost for this batch
+        total_token_usage = TokenUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
+        total_cost = 0.0
+
         # appending allows for the resubmitted resumed batch
         with open(response_file, "a") as f:
             for raw_response in responses:
@@ -427,6 +445,18 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
                 json.dump(generic_response.model_dump(), f, default=str)
                 f.write("\n")
 
+                # Update token and cost totals
+                if generic_response.token_usage:
+                    total_token_usage.prompt_tokens += generic_response.token_usage.prompt_tokens
+                    total_token_usage.completion_tokens += generic_response.token_usage.completion_tokens
+                    total_token_usage.total_tokens += generic_response.token_usage.total_tokens
+                if generic_response.response_cost:
+                    total_cost += generic_response.response_cost
+
+        # Update tracker with token usage and cost stats
+        self.tracker.update_token_and_cost(total_token_usage, total_cost)
+        return response_file
+
     async def submit_batches_from_request_files(
         self,
         request_files: set[str],
@@ -446,8 +476,6 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
             - Submits batches concurrently using asyncio
             - Updates batch objects file
         """
-        self._attempt_loading_batch_status_tracker(request_files)
-
         tasks = []
 
         # check existing response files for resuming
@@ -466,16 +494,7 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
         if len(tasks) == 0:
             return
 
-        # submit batches
-        self.batch_submit_pbar = tqdm(
-            total=len(tasks) + self.tracker.n_submitted_finished_or_downloaded_batches,
-            desc="Submitting batches",
-            unit="batch",
-            initial=self.tracker.n_submitted_finished_or_downloaded_batches,
-        )
-
         await asyncio.gather(*tasks)
-        self.batch_submit_pbar.close()
         assert self.tracker.unsubmitted_request_files == set()
 
     async def check_batch_status(self, batch: GenericBatch) -> GenericBatch | None:
@@ -514,6 +533,7 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
                 if batch.status == GenericBatchStatus.FINISHED:
                     logger.debug(f"Batch {batch.id} finished with status: {batch.raw_status}")
                     self.tracker.mark_as_finished(batch)
+                    await self.update_batch_objects_file()
 
     async def poll_and_process_batches(self) -> None:
         """Monitor and process batches until completion.
@@ -535,14 +555,6 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
             - Logs progress and status information
             - Sleeps between status checks
         """
-        # progress bar for finished requests
-        self.request_pbar = tqdm(
-            total=self.tracker.n_total_requests,
-            desc="Finished requests in batches",
-            unit="request",
-            initial=self.tracker.n_finished_or_downloaded_succeeded_requests,
-        )
-
         # loop until all batches have been returned
         all_response_files = []
         while self.tracker.n_submitted_batches + self.tracker.n_finished_batches > 0:
@@ -551,10 +563,6 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
             await asyncio.gather(*status_tasks)
             await self.update_batch_objects_file()
 
-            # update progress bar
-            self.request_pbar.n = self.tracker.n_finished_or_downloaded_succeeded_requests
-            self.request_pbar.refresh()
-
             download_tasks = [self.download_batch_to_response_file(batch) for batch in self.tracker.finished_batches.values()]
             # Failed downloads return None and print any errors that occurred
             all_response_files.extend(await asyncio.gather(*download_tasks))
@@ -564,46 +572,24 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
                 f"Requests completed: {self.tracker.n_finished_or_downloaded_succeeded_requests:,}/{self.tracker.n_total_requests:,}"
             )
 
+            self.tracker.update_display()
             if self.tracker.n_submitted_batches + self.tracker.n_finished_batches > 0:
                 logger.debug(f"Sleeping for {self.config.batch_check_interval} seconds...")
                 await asyncio.sleep(self.config.batch_check_interval)
 
-        self.request_pbar.close()
         response_files = filter(None, all_response_files)
         if self.tracker.n_downloaded_batches == 0 or not response_files:
-            raise RuntimeError("None of the submitted batches completed successfully. " f"Please check the logs above and {self.web_dashboard} for errors.")
+            raise RuntimeError(f"None of the submitted batches completed successfully. Please check the logs above and {self.web_dashboard} for errors.")
 
     async def download_batch_to_response_file(self, batch: GenericBatch) -> str | None:
-        """Download and process completed batch results.
-
-        Downloads batch results, converts responses to generic format, and handles
-        cleanup of completed batches including file deletion if configured.
-
-        Args:
-            batch: The completed batch object to process.
-
-        Returns:
-            str | None: Path to response file if successful, None if failed.
-
-        Side Effects:
-            - Downloads batch results from API
-            - Creates response file with processed results
-            - Updates batch tracking state
-            - Updates batch objects file
-            - Optionally deletes API provider's batch files
-            - Logs download progress and completion
-        """
+        """Download and process completed batch results."""
         file_content = await self.download_batch(batch)
 
         if file_content is None:
             return None
 
-        request_file = batch.request_file
-        request_dir = os.path.dirname(request_file)
-        request_filename = os.path.basename(request_file)
-        response_filename = request_filename.replace("requests_", "responses_")
-        response_file = os.path.join(request_dir, response_filename)
-        self.generic_response_file_from_responses(file_content, batch)
+        # Write responses to file and update stats
+        response_file = self.generic_response_file_from_responses(file_content, batch)
 
         logger.debug(f"Batch {batch.id} written to {response_file}")
 
@@ -611,8 +597,10 @@ class BaseBatchRequestProcessor(BaseRequestProcessor):
             await self.delete_file(batch.input_file_id, self.semaphore)
             await self.delete_file(batch.output_file_id, self.semaphore)
 
+        # Update tracker with downloaded batch
         self.tracker.mark_as_downloaded(batch)
         await self.update_batch_objects_file()
+
         return response_file
 
     async def cancel_batches(self):
diff --git a/src/bespokelabs/curator/request_processor/online/base_online_request_processor.py b/src/bespokelabs/curator/request_processor/online/base_online_request_processor.py
index 0b8d1d6..374f4b1 100644
--- a/src/bespokelabs/curator/request_processor/online/base_online_request_processor.py
+++ b/src/bespokelabs/curator/request_processor/online/base_online_request_processor.py
@@ -222,11 +222,8 @@ class BaseOnlineRequestProcessor(BaseRequestProcessor, ABC):
         completed_request_ids = self.validate_existing_response_file(response_file)
 
         # Count total requests
-        with open(generic_request_filepath) as f:
-            total_requests = sum(1 for _ in f)
-
         status_tracker.num_tasks_already_completed = len(completed_request_ids)
-        status_tracker.total_requests = total_requests
+        status_tracker.total_requests = self.total_requests
         status_tracker.model = self.prompt_formatter.model_name
         status_tracker.start_tracker(self._tracker_console)
 
diff --git a/src/bespokelabs/curator/status_tracker/batch_status_tracker.py b/src/bespokelabs/curator/status_tracker/batch_status_tracker.py
index fbe8ae4..e9899f6 100644
--- a/src/bespokelabs/curator/status_tracker/batch_status_tracker.py
+++ b/src/bespokelabs/curator/status_tracker/batch_status_tracker.py
@@ -1,10 +1,17 @@
 """Module for tracking the status of batches during curation."""
 
 import logging
+import time
+from typing import Optional
 
 from pydantic import BaseModel, Field
+from rich import box
+from rich.console import Console
+from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn
+from rich.table import Table
 
 from bespokelabs.curator.types.generic_batch import GenericBatch, GenericBatchStatus
+from bespokelabs.curator.types.generic_response import TokenUsage
 
 logger = logging.getLogger(__name__)
 
@@ -17,12 +24,203 @@ class BatchStatusTracker(BaseModel):
     update the status of batches and requests throughout the curation process.
     """
 
+    model_config = {
+        "arbitrary_types_allowed": True,  # Allow non-serializable types
+        "exclude": {"console", "progress", "task_id"},  # Exclude from serialization
+    }
+
     n_total_requests: int = Field(default=0)
     unsubmitted_request_files: set[str] = Field(default_factory=set)
     submitted_batches: dict[str, GenericBatch] = Field(default_factory=dict)
     finished_batches: dict[str, GenericBatch] = Field(default_factory=dict)
     downloaded_batches: dict[str, GenericBatch] = Field(default_factory=dict)
 
+    # Add fields for tracking costs and tokens
+    total_prompt_tokens: int = Field(default=0)
+    total_completion_tokens: int = Field(default=0)
+    total_tokens: int = Field(default=0)
+    total_cost: float = Field(default=0.0)
+
+    # Model information
+    model: str = Field(default="")
+    input_cost_per_million: Optional[float] = Field(default=None)
+    output_cost_per_million: Optional[float] = Field(default=None)
+
+    # Track start time
+    start_time: float = Field(default_factory=time.time)
+
+    def start_tracker(self, console: Optional[Console] = None):
+        """Start the progress tracker with rich console output."""
+        self._console = Console() if console is None else console
+        self._progress = Progress(
+            TextColumn(
+                "[cyan]{task.description}[/cyan]\n"
+                "{task.fields[batches_text]}\n"
+                "{task.fields[requests_text]}\n"
+                "{task.fields[tokens_text]}\n"
+                "{task.fields[cost_text]}\n"
+                "{task.fields[price_text]}",
+                justify="left",
+            ),
+            TextColumn("\n\n\n\n\n\n"),  # Spacer
+            BarColumn(),
+            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+            TextColumn("[bold white]•[/bold white]"),
+            TimeElapsedColumn(),
+            console=self._console,
+        )
+
+        self._task_id = self._progress.add_task(
+            description=f"[cyan]Processing batches using {self.model}",
+            total=self.n_total_requests,
+            # Since we don't automatically retry failed requests within a run, we can count
+            # failed downloaded requests as "completed". Users who require
+            # 100% success rate can set require_all_responses to True and
+            # manually retry.
+            completed=self.n_downloaded_succeeded_requests + self.n_downloaded_failed_requests,
+            batches_text="[bold white]Batches:[/bold white] [dim]--[/dim]",
+            requests_text="[bold white]Requests:[/bold white] [dim]--[/dim]",
+            tokens_text="[bold white]Tokens:[/bold white] [dim]--[/dim]",
+            cost_text="[bold white]Cost:[/bold white] [dim]--[/dim]",
+            price_text="[bold white]Model Pricing:[/bold white] [dim]--[/dim]",
+        )
+        self._progress.start()
+
+    def stop_tracker(self):
+        """Stop the tracker and display final statistics."""
+        if self._progress:
+            self._progress.stop()
+            self.display_final_stats()
+
+    def update_display(self):
+        """Update statistics with token usage and cost information."""
+        # Format display texts
+        batches_text = (
+            "[bold white]Batches:[/bold white] "
+            f"[white]Total:[/white] [blue]{self.n_total_batches}[/blue] "
+            f"[white]•[/white] "
+            f"[white]Submitted:[/white] [yellow]{self.n_submitted_batches}⋯[/yellow] "
+            f"[white]•[/white] "
+            f"[white]Downloaded:[/white] [green]{self.n_downloaded_batches}✓[/green]"
+        )
+
+        n_submitted_requests = self.n_total_requests - (self.n_downloaded_succeeded_requests + self.n_downloaded_failed_requests)
+        requests_text = (
+            "[bold white]Requests:[/bold white] "
+            f"[white]Total:[/white] [blue]{self.n_total_requests}[/blue] "
+            f"[white]•[/white] "
+            f"[white]Submitted:[/white] [yellow]{n_submitted_requests}⋯[/yellow] "
+            f"[white]•[/white] "
+            f"[white]Succeeded:[/white] [green]{self.n_downloaded_succeeded_requests}✓[/green] "
+            f"[white]•[/white] "
+            f"[white]Failed:[/white] [red]{self.n_downloaded_failed_requests}✗[/red] "
+        )
+
+        avg_prompt = self.total_prompt_tokens / max(1, self.n_finished_or_downloaded_succeeded_requests)
+        avg_completion = self.total_completion_tokens / max(1, self.n_finished_or_downloaded_succeeded_requests)
+
+        tokens_text = (
+            "[bold white]Tokens:[/bold white] "
+            f"[white]Avg Input:[/white] [blue]{avg_prompt:.0f}[/blue] "
+            f"[white]•[/white] "
+            f"[white]Avg Output:[/white] [blue]{avg_completion:.0f}[/blue] "
+        )
+
+        avg_cost = self.total_cost / max(1, self.n_downloaded_succeeded_requests)
+        projected_cost = avg_cost * self.n_total_requests
+
+        cost_text = (
+            "[bold white]Cost:[/bold white] "
+            f"[white]Current:[/white] [magenta]${self.total_cost:.3f}[/magenta] "
+            f"[white]•[/white] "
+            f"[white]Projected:[/white] [magenta]${projected_cost:.3f}[/magenta] "
+            f"[white]•[/white] "
+            f"[white]Rate:[/white] [magenta]${avg_cost:.3f}/request[/magenta]"
+        )
+
+        input_cost_str = f"${self.input_cost_per_million:.3f}" if self.input_cost_per_million else "N/A"
+        output_cost_str = f"${self.output_cost_per_million:.3f}" if self.output_cost_per_million else "N/A"
+
+        price_text = (
+            "[bold white]Model Pricing:[/bold white] "
+            f"[white]Per 1M tokens:[/white] "
+            f"[white]Input:[/white] [red]{input_cost_str}[/red] "
+            f"[white]•[/white] "
+            f"[white]Output:[/white] [red]{output_cost_str}[/red]"
+        )
+
+        # Update progress display
+        self._progress.update(
+            self._task_id,
+            completed=self.n_downloaded_succeeded_requests + self.n_downloaded_failed_requests,
+            batches_text=batches_text,
+            requests_text=requests_text,
+            tokens_text=tokens_text,
+            cost_text=cost_text,
+            price_text=price_text,
+        )
+
+    def display_final_stats(self):
+        """Display final statistics."""
+        table = Table(title="Final Curator Statistics", box=box.ROUNDED)
+        table.add_column("Section/Metric", style="cyan")
+        table.add_column("Value", style="yellow")
+
+        # Model Information
+        table.add_row("Model", "", style="bold magenta")
+        table.add_row("Model", f"[blue]{self.model}[/blue]")
+
+        # Batch Statistics
+        table.add_row("Batches", "", style="bold magenta")
+        table.add_row("Total Batches", str(self.n_total_batches))
+        table.add_row("Submitted", f"[yellow]{self.n_submitted_batches}[/yellow]")
+        table.add_row("Downloaded", f"[green]{self.n_downloaded_batches}[/green]")
+
+        # Request Statistics
+        table.add_row("Requests", "", style="bold magenta")
+        table.add_row("Total Requests", str(self.n_total_requests))
+        table.add_row("Successful", f"[green]{self.n_finished_or_downloaded_succeeded_requests}[/green]")
+        table.add_row("Failed", f"[red]{self.n_finished_failed_requests + self.n_downloaded_failed_requests}[/red]")
+
+        # Token Statistics
+        table.add_row("Tokens", "", style="bold magenta")
+        table.add_row("Total Tokens Used", f"{self.total_tokens:,}")
+        table.add_row("Total Input Tokens", f"{self.total_prompt_tokens:,}")
+        table.add_row("Total Output Tokens", f"{self.total_completion_tokens:,}")
+        if self.n_finished_or_downloaded_succeeded_requests > 0:
+            table.add_row("Average Tokens per Request", f"{int(self.total_tokens / self.n_finished_or_downloaded_succeeded_requests)}")
+            table.add_row("Average Input Tokens", f"{int(self.total_prompt_tokens / self.n_finished_or_downloaded_succeeded_requests)}")
+            table.add_row("Average Output Tokens", f"{int(self.total_completion_tokens / self.n_finished_or_downloaded_succeeded_requests)}")
+
+        # Cost Statistics
+        table.add_row("Costs", "", style="bold magenta")
+        table.add_row("Total Cost", f"[red]${self.total_cost:.4f}[/red]")
+        if self.n_finished_or_downloaded_succeeded_requests > 0:
+            table.add_row("Average Cost per Request", f"[red]${self.total_cost / self.n_finished_or_downloaded_succeeded_requests:.4f}[/red]")
+
+        # Handle None values for cost per million tokens
+        input_cost_str = f"[red]${self.input_cost_per_million:.4f}[/red]" if self.input_cost_per_million else "[dim]N/A[/dim]"
+        output_cost_str = f"[red]${self.output_cost_per_million:.4f}[/red]" if self.output_cost_per_million else "[dim]N/A[/dim]"
+
+        table.add_row("Input Cost per 1M Tokens", input_cost_str)
+        table.add_row("Output Cost per 1M Tokens", output_cost_str)
+
+        # Performance Statistics
+        table.add_row("Performance", "", style="bold magenta")
+        elapsed_time = time.time() - self.start_time
+        elapsed_minutes = elapsed_time / 60
+        rpm = (self.n_downloaded_succeeded_requests + self.n_downloaded_failed_requests) / max(0.001, elapsed_minutes)
+        input_tpm = self.total_prompt_tokens / max(0.001, elapsed_minutes)
+        output_tpm = self.total_completion_tokens / max(0.001, elapsed_minutes)
+
+        table.add_row("Total Time", f"{elapsed_time:.2f}s")
+        table.add_row("Average Time per Request", f"{elapsed_time / max(1, self.n_finished_or_downloaded_succeeded_requests):.2f}s")
+        table.add_row("Requests per Minute", f"{rpm:.1f}")
+        table.add_row("Input Tokens per Minute", f"{input_tpm:.1f}")
+        table.add_row("Output Tokens per Minute", f"{output_tpm:.1f}")
+
+        self._console.print(table)
+
     @property
     def n_total_batches(self) -> int:
         """Get the total number of batches across all states."""
@@ -118,11 +316,11 @@ class BatchStatusTracker(BaseModel):
         batch.status = GenericBatchStatus.SUBMITTED
         if batch.request_file in self.unsubmitted_request_files:
             self.unsubmitted_request_files.remove(batch.request_file)
-            self.n_total_requests += n_requests
         else:
             logger.warning(f"Request file {batch.request_file} is being re-submitted.")
         self.submitted_batches[batch.id] = batch
         logger.debug(f"Marked {batch.request_file} as submitted with batch {batch.id}")
+        self.update_display()
 
     def mark_as_finished(self, batch: GenericBatch):
         """Mark a batch as finished and move it to finished batches.
@@ -135,6 +333,7 @@ class BatchStatusTracker(BaseModel):
         self.submitted_batches.pop(batch.id)
         self.finished_batches[batch.id] = batch
         logger.debug(f"Marked batch {batch.id} as finished")
+        self.update_display()
 
     def mark_as_downloaded(self, batch: GenericBatch):
         """Mark a batch as downloaded and move it to downloaded batches.
@@ -147,6 +346,7 @@ class BatchStatusTracker(BaseModel):
         self.finished_batches.pop(batch.id)
         self.downloaded_batches[batch.id] = batch
         logger.debug(f"Marked batch {batch.id} as downloaded")
+        self.update_display()
 
     def update_submitted(self, batch: GenericBatch):
         """Update the request counts for a submitted batch.
@@ -157,6 +357,7 @@ class BatchStatusTracker(BaseModel):
         assert batch.id in self.submitted_batches
         self.submitted_batches[batch.id] = batch
         logger.debug(f"Updated submitted batch {batch.id} with new request counts")
+        self.update_display()
 
     def __str__(self) -> str:
         """Return a human-readable string representation of the batch status."""
@@ -174,3 +375,13 @@ class BatchStatusTracker(BaseModel):
             f"Downloaded succeeded requests: {self.n_downloaded_succeeded_requests}",
         ]
         return "\n".join(status_lines)
+
+    def update_token_and_cost(self, token_usage: TokenUsage, cost: float):
+        """Update statistics with token usage and cost information."""
+        if token_usage:
+            self.total_prompt_tokens += token_usage.prompt_tokens
+            self.total_completion_tokens += token_usage.completion_tokens
+            self.total_tokens += token_usage.total_tokens
+        if cost:
+            self.total_cost += cost
+        self.update_display()


Test patch:
diff --git a/tests/integrations/test_all.py b/tests/integrations/test_all.py
index 1a0a21a..aa4239c 100644
--- a/tests/integrations/test_all.py
+++ b/tests/integrations/test_all.py
@@ -223,6 +223,19 @@ def test_basic_batch(temp_working_dir, mock_dataset):
         "anthropic": "f38e7406448e95160ebe4d9b6148920ef37b019f23a4e2c57094fdd4bafb09be",
     }
     with vcr_config.use_cassette("basic_batch_completion.yaml"):
-        dataset = helper.create_basic(temp_working_dir, mock_dataset, batch=True, backend=backend)
+        output = StringIO()
+        console = Console(file=output, width=300)
+
+        dataset = helper.create_basic(temp_working_dir, mock_dataset, batch=True, backend=backend, tracker_console=console)
         recipes = "".join([recipe[0] for recipe in dataset.to_pandas().values.tolist()])
         assert _hash_string(recipes) == hash_book[backend]
+
+        # Verify status tracker output
+        captured = output.getvalue()
+        assert "Processing batches using" in captured, captured
+        assert "Batches: Total: 1 • Submitted: 0⋯ • Downloaded: 1✓" in captured, captured
+        assert "Requests: Total: 3 • Submitted: 0⋯ • Succeeded: 3✓ • Failed: 0✗" in captured, captured
+        assert "Final Curator Statistics" in captured, captured
+        assert "Total Requests             │ 3" in captured, captured
+        assert "Successful                 │ 3" in captured, captured
+        assert "Failed                     │ 0" in captured, captured
diff --git a/tests/unittests/test_status_tracker.py b/tests/unittests/test_status_tracker.py
index 5805ed5..809a0f0 100644
--- a/tests/unittests/test_status_tracker.py
+++ b/tests/unittests/test_status_tracker.py
@@ -1,11 +1,15 @@
+import datetime
 from io import StringIO
 
 from rich.console import Console
 
+from bespokelabs.curator.status_tracker.batch_status_tracker import BatchStatusTracker
 from bespokelabs.curator.status_tracker.online_status_tracker import OnlineStatusTracker
+from bespokelabs.curator.types.generic_batch import GenericBatch, GenericBatchRequestCounts, GenericBatchStatus
+from bespokelabs.curator.types.generic_response import TokenUsage
 
 
-def test_status_tracker_display():
+def test_online_status_tracker_display():
     """Test that the status tracker display works correctly."""
     # Create a string buffer to capture output
     output = StringIO()
@@ -79,3 +83,112 @@ def test_status_tracker_final_stats():
     assert "1,000" in captured  # Input tokens
     assert "2,000" in captured  # Output tokens
     assert "$0.123" in captured  # Cost
+
+
+def test_batch_status_tracker_display():
+    """Test that the batch status tracker display works correctly."""
+    # Create a string buffer to capture output
+    output = StringIO()
+    # Set width=200 and force_terminal=True to prevent truncation
+    console = Console(file=output, width=200, force_terminal=True)
+
+    # Create and configure the status tracker
+    tracker = BatchStatusTracker()
+    tracker.model = "test-model"
+    tracker.input_cost_per_million = 0.15
+    tracker.output_cost_per_million = 0.60
+
+    # Start display
+    tracker.start_tracker(console)
+
+    # Create a test batch
+    batch = GenericBatch(
+        id="test-batch-1",
+        request_file="test.jsonl",
+        status=GenericBatchStatus.SUBMITTED,
+        created_at=datetime.datetime.now(),
+        finished_at=None,
+        api_key_suffix="1234",
+        raw_status="submitted",
+        raw_batch={},
+        request_counts=GenericBatchRequestCounts(total=20, succeeded=15, failed=5, raw_request_counts_object={"total": 20, "succeeded": 15, "failed": 5}),
+    )
+
+    # Simulate batch progress
+    tracker.mark_as_submitted(batch, n_requests=20)
+
+    # Update token usage and cost
+    token_usage = TokenUsage(prompt_tokens=1000, completion_tokens=2000, total_tokens=3000)
+    tracker.update_token_and_cost(token_usage, cost=0.123)
+
+    # Mark batch as downloaded
+    tracker.mark_as_finished(batch)
+    tracker.mark_as_downloaded(batch)  # Add this to move it to downloaded state
+
+    # Stop display
+    tracker.stop_tracker()
+
+    # Get the captured output
+    captured = output.getvalue()
+
+    # Verify key elements are present
+    print(captured)
+    assert "test-model" in captured
+    assert "1✓" in captured.replace(" ", "")  # Downloaded batch count
+    assert "15✓" in captured.replace(" ", "")  # Success count
+    assert "5✗" in captured.replace(" ", "")  # Failed count
+    assert "$0.123" in captured  # Cost
+
+
+def test_batch_status_tracker_final_stats():
+    """Test that the final statistics table for batch tracker is formatted correctly."""
+    output = StringIO()
+    console = Console(file=output, width=200, force_terminal=True)
+
+    tracker = BatchStatusTracker()
+    tracker.model = "test-model"
+    tracker.input_cost_per_million = 0.15
+    tracker.output_cost_per_million = 0.60
+
+    tracker.start_tracker(console)
+
+    # Create and submit a test batch
+    batch = GenericBatch(
+        id="test-batch-1",
+        request_file="test.jsonl",
+        status=GenericBatchStatus.SUBMITTED,
+        created_at=datetime.datetime.now(),
+        finished_at=None,
+        api_key_suffix="1234",
+        raw_status="submitted",
+        raw_batch={},
+        request_counts=GenericBatchRequestCounts(total=50, succeeded=45, failed=5, raw_request_counts_object={"total": 50, "succeeded": 45, "failed": 5}),
+    )
+
+    tracker.mark_as_submitted(batch, n_requests=50)
+
+    # Update token usage and cost
+    token_usage = TokenUsage(prompt_tokens=1000, completion_tokens=2000, total_tokens=3000)
+    tracker.update_token_and_cost(token_usage, cost=0.123)
+
+    # Mark as finished and downloaded
+    tracker.mark_as_finished(batch)
+    tracker.mark_as_downloaded(batch)
+
+    # Generate final stats
+    tracker.stop_tracker()
+
+    # Get the captured output
+    captured = output.getvalue()
+
+    # Verify the final statistics table
+    assert "Final Curator Statistics" in captured, captured
+    assert "test-model" in captured, captured
+    assert "Total Tokens Used" in captured, captured
+    assert "1,000" in captured, captured  # Input tokens
+    assert "2,000" in captured, captured  # Output tokens
+    assert "$0.123" in captured, captured  # Cost
+    assert "45" in captured, captured  # Successful requests
+    assert "5" in captured, captured  # Failed requests
+    assert "$0.15" in captured, captured  # Input cost per million
+    assert "$0.60" in captured, captured  # Output cost per million