| """A simple progress bar for the console.""" |
|
|
| import threading |
| from typing import Any, Dict, Optional, Sequence |
| from uuid import UUID |
|
|
| from langchain_core.callbacks import base as base_callbacks |
| from langchain_core.documents import Document |
| from langchain_core.outputs import LLMResult |
|
|
|
|
| class ProgressBarCallback(base_callbacks.BaseCallbackHandler): |
| """A simple progress bar for the console.""" |
|
|
| def __init__(self, total: int, ncols: int = 50, **kwargs: Any): |
| """Initialize the progress bar. |
| |
| Args: |
| total: int, the total number of items to be processed. |
| ncols: int, the character width of the progress bar. |
| """ |
| self.total = total |
| self.ncols = ncols |
| self.counter = 0 |
| self.lock = threading.Lock() |
| self._print_bar() |
|
|
| def increment(self) -> None: |
| """Increment the counter and update the progress bar.""" |
| with self.lock: |
| self.counter += 1 |
| self._print_bar() |
|
|
| def _print_bar(self) -> None: |
| """Print the progress bar to the console.""" |
| progress = self.counter / self.total |
| arrow = "-" * int(round(progress * self.ncols) - 1) + ">" |
| spaces = " " * (self.ncols - len(arrow)) |
| print(f"\r[{arrow + spaces}] {self.counter}/{self.total}", end="") |
|
|
| def on_chain_error( |
| self, |
| error: BaseException, |
| *, |
| run_id: UUID, |
| parent_run_id: Optional[UUID] = None, |
| **kwargs: Any, |
| ) -> Any: |
| if parent_run_id is None: |
| self.increment() |
|
|
| def on_chain_end( |
| self, |
| outputs: Dict[str, Any], |
| *, |
| run_id: UUID, |
| parent_run_id: Optional[UUID] = None, |
| **kwargs: Any, |
| ) -> Any: |
| if parent_run_id is None: |
| self.increment() |
|
|
| def on_retriever_error( |
| self, |
| error: BaseException, |
| *, |
| run_id: UUID, |
| parent_run_id: Optional[UUID] = None, |
| **kwargs: Any, |
| ) -> Any: |
| if parent_run_id is None: |
| self.increment() |
|
|
| def on_retriever_end( |
| self, |
| documents: Sequence[Document], |
| *, |
| run_id: UUID, |
| parent_run_id: Optional[UUID] = None, |
| **kwargs: Any, |
| ) -> Any: |
| if parent_run_id is None: |
| self.increment() |
|
|
| def on_llm_error( |
| self, |
| error: BaseException, |
| *, |
| run_id: UUID, |
| parent_run_id: Optional[UUID] = None, |
| **kwargs: Any, |
| ) -> Any: |
| if parent_run_id is None: |
| self.increment() |
|
|
| def on_llm_end( |
| self, |
| response: LLMResult, |
| *, |
| run_id: UUID, |
| parent_run_id: Optional[UUID] = None, |
| **kwargs: Any, |
| ) -> Any: |
| if parent_run_id is None: |
| self.increment() |
|
|
| def on_tool_error( |
| self, |
| error: BaseException, |
| *, |
| run_id: UUID, |
| parent_run_id: Optional[UUID] = None, |
| **kwargs: Any, |
| ) -> Any: |
| if parent_run_id is None: |
| self.increment() |
|
|
| def on_tool_end( |
| self, |
| output: str, |
| *, |
| run_id: UUID, |
| parent_run_id: Optional[UUID] = None, |
| **kwargs: Any, |
| ) -> Any: |
| if parent_run_id is None: |
| self.increment() |
|
|
