Refactor callbacks (#1583)

* Unify Workflow and Verb callbacks interfaces

* Semver

* Fix storage class instantiation (#1582)

---------

Co-authored-by: Josh Bradley <[email protected]>

.semversioner/next-release/patch-20250103231659816022.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Simplify callbacks model."
+}

docs/examples_notebooks/index_migration.ipynb

@@ -207,7 +207,7 @@
    "outputs": [],
    "source": [
     "from graphrag.cache.factory import create_cache\n",
-    "from graphrag.callbacks.noop_verb_callbacks import NoopVerbCallbacks\n",
+    "from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks\n",
     "from graphrag.index.flows.generate_text_embeddings import generate_text_embeddings\n",
     "\n",
     "# We only need to re-run the embeddings workflow, to ensure that embeddings for all required search fields are in place\n",
@@ -219,7 +219,7 @@
     "config = workflow.config\n",
     "text_embed = config.get(\"text_embed\", {})\n",
     "embedded_fields = config.get(\"embedded_fields\", {})\n",
-    "callbacks = NoopVerbCallbacks()\n",
+    "callbacks = NoopWorkflowCallbacks()\n",
     "cache = create_cache(pipeline_config.cache, PROJECT_DIRECTORY)\n",
     "\n",
     "await generate_text_embeddings(\n",

graphrag/api/prompt_tune.py

@@ -13,7 +13,7 @@
 
 from pydantic import PositiveInt, validate_call
 
-from graphrag.callbacks.noop_verb_callbacks import NoopVerbCallbacks
+from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 from graphrag.config.models.graph_rag_config import GraphRagConfig
 from graphrag.index.llm.load_llm import load_llm
 from graphrag.logger.print_progress import PrintProgressLogger
@@ -99,7 +99,7 @@ async def generate_indexing_prompts(
         "prompt_tuning",
         config.llm,
         cache=None,
-        callbacks=NoopVerbCallbacks(),
+        callbacks=NoopWorkflowCallbacks(),
     )
 
     if not domain:

graphrag/callbacks/blob_workflow_callbacks.py

@@ -84,7 +84,7 @@ def _write_log(self, log: dict[str, Any]):
         # update the blob's block count
         self._num_blocks += 1
 
-    def on_error(
+    def error(
         self,
         message: str,
         cause: BaseException | None = None,
@@ -100,10 +100,10 @@ def on_error(
             "details": details,
         })
 
-    def on_warning(self, message: str, details: dict | None = None):
+    def warning(self, message: str, details: dict | None = None):
         """Report a warning."""
         self._write_log({"type": "warning", "data": message, "details": details})
 
-    def on_log(self, message: str, details: dict | None = None):
+    def log(self, message: str, details: dict | None = None):
         """Report a generic log message."""
         self._write_log({"type": "log", "data": message, "details": details})

graphrag/callbacks/console_workflow_callbacks.py

@@ -9,7 +9,7 @@
 class ConsoleWorkflowCallbacks(NoopWorkflowCallbacks):
     """A logger that writes to a console."""
 
-    def on_error(
+    def error(
         self,
         message: str,
         cause: BaseException | None = None,
@@ -19,11 +19,11 @@ def on_error(
         """Handle when an error occurs."""
         print(message, str(cause), stack, details)  # noqa T201
 
-    def on_warning(self, message: str, details: dict | None = None):
+    def warning(self, message: str, details: dict | None = None):
         """Handle when a warning occurs."""
         _print_warning(message)
 
-    def on_log(self, message: str, details: dict | None = None):
+    def log(self, message: str, details: dict | None = None):
         """Handle when a log message is produced."""
         print(message, details)  # noqa T201
 

graphrag/callbacks/file_workflow_callbacks.py

@@ -25,7 +25,7 @@ def __init__(self, directory: str):
             Path(directory) / "logs.json", "a", encoding="utf-8", errors="strict"
         )
 
-    def on_error(
+    def error(
         self,
         message: str,
         cause: BaseException | None = None,
@@ -50,7 +50,7 @@ def on_error(
         message = f"{message} details={details}"
         log.info(message)
 
-    def on_warning(self, message: str, details: dict | None = None):
+    def warning(self, message: str, details: dict | None = None):
         """Handle when a warning occurs."""
         self._out_stream.write(
             json.dumps(
@@ -61,7 +61,7 @@ def on_warning(self, message: str, details: dict | None = None):
         )
         _print_warning(message)
 
-    def on_log(self, message: str, details: dict | None = None):
+    def log(self, message: str, details: dict | None = None):
         """Handle when a log message is produced."""
         self._out_stream.write(
             json.dumps(

graphrag/callbacks/noop_workflow_callbacks.py

@@ -3,31 +3,23 @@
 
 """A no-op implementation of WorkflowCallbacks."""
 
-from typing import Any
-
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
 from graphrag.logger.progress import Progress
 
 
 class NoopWorkflowCallbacks(WorkflowCallbacks):
     """A no-op implementation of WorkflowCallbacks."""
 
-    def on_workflow_start(self, name: str, instance: object) -> None:
+    def workflow_start(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow starts."""
 
-    def on_workflow_end(self, name: str, instance: object) -> None:
+    def workflow_end(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow ends."""
 
-    def on_step_start(self, step_name: str) -> None:
-        """Execute this callback every time a step starts."""
-
-    def on_step_end(self, step_name: str, result: Any) -> None:
-        """Execute this callback every time a step ends."""
-
-    def on_step_progress(self, step_name: str, progress: Progress) -> None:
+    def progress(self, progress: Progress) -> None:
         """Handle when progress occurs."""
 
-    def on_error(
+    def error(
         self,
         message: str,
         cause: BaseException | None = None,
@@ -36,11 +28,8 @@ def on_error(
     ) -> None:
         """Handle when an error occurs."""
 
-    def on_warning(self, message: str, details: dict | None = None) -> None:
+    def warning(self, message: str, details: dict | None = None) -> None:
         """Handle when a warning occurs."""
 
-    def on_log(self, message: str, details: dict | None = None) -> None:
+    def log(self, message: str, details: dict | None = None) -> None:
         """Handle when a log message occurs."""
-
-    def on_measure(self, name: str, value: float, details: dict | None = None) -> None:
-        """Handle when a measurement occurs."""

graphrag/callbacks/progress_workflow_callbacks.py

@@ -3,8 +3,6 @@
 
 """A workflow callback manager that emits updates."""
 
-from typing import Any
-
 from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 from graphrag.logger.base import ProgressLogger
 from graphrag.logger.progress import Progress
@@ -31,23 +29,14 @@ def _push(self, name: str) -> None:
     def _latest(self) -> ProgressLogger:
         return self._progress_stack[-1]
 
-    def on_workflow_start(self, name: str, instance: object) -> None:
+    def workflow_start(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow starts."""
         self._push(name)
 
-    def on_workflow_end(self, name: str, instance: object) -> None:
+    def workflow_end(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow ends."""
         self._pop()
 
-    def on_step_start(self, step_name: str) -> None:
-        """Execute this callback every time a step starts."""
-        self._push(f"Step {step_name}")
-        self._latest(Progress(percent=0))
-
-    def on_step_end(self, step_name: str, result: Any) -> None:
-        """Execute this callback every time a step ends."""
-        self._pop()
-
-    def on_step_progress(self, step_name: str, progress: Progress) -> None:
+    def progress(self, progress: Progress) -> None:
         """Handle when progress occurs."""
         self._latest(progress)

graphrag/callbacks/workflow_callbacks.py

@@ -3,7 +3,7 @@
 
 """Collection of callbacks that can be used to monitor the workflow execution."""
 
-from typing import Any, Protocol
+from typing import Protocol
 
 from graphrag.logger.progress import Progress
 
@@ -15,27 +15,19 @@ class WorkflowCallbacks(Protocol):
     This base class is a "noop" implementation so that clients may implement just the callbacks they need.
     """
 
-    def on_workflow_start(self, name: str, instance: object) -> None:
+    def workflow_start(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow starts."""
         ...
 
-    def on_workflow_end(self, name: str, instance: object) -> None:
+    def workflow_end(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow ends."""
         ...
 
-    def on_step_start(self, step_name: str) -> None:
-        """Execute this callback every time a step starts."""
-        ...
-
-    def on_step_end(self, step_name: str, result: Any) -> None:
-        """Execute this callback every time a step ends."""
-        ...
-
-    def on_step_progress(self, step_name: str, progress: Progress) -> None:
+    def progress(self, progress: Progress) -> None:
         """Handle when progress occurs."""
         ...
 
-    def on_error(
+    def error(
         self,
         message: str,
         cause: BaseException | None = None,
@@ -45,14 +37,10 @@ def on_error(
         """Handle when an error occurs."""
         ...
 
-    def on_warning(self, message: str, details: dict | None = None) -> None:
+    def warning(self, message: str, details: dict | None = None) -> None:
         """Handle when a warning occurs."""
         ...
 
-    def on_log(self, message: str, details: dict | None = None) -> None:
+    def log(self, message: str, details: dict | None = None) -> None:
         """Handle when a log message occurs."""
         ...
-
-    def on_measure(self, name: str, value: float, details: dict | None = None) -> None:
-        """Handle when a measurement occurs."""
-        ...