update callbacks

Author: jgbradley1

backend/src/logger/application_insights_workflow_callbacks.py

@@ -13,7 +13,7 @@
 )
 
 from azure.monitor.opentelemetry.exporter import AzureMonitorLogExporter
-from datashaper.workflow.workflow_callbacks import NoopWorkflowCallbacks
+from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 from opentelemetry._logs import (
     get_logger_provider,
     set_logger_provider,
@@ -121,7 +121,7 @@ def _format_details(self, details: Dict[str, Any] | None = None) -> Dict[str, An
             return {}
         return {"custom_dimensions": {**self._properties, **unwrap_dict(details)}}
 
-    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._workflow_name = name
         self._processed_workflow_steps.append(name)
@@ -142,7 +142,7 @@ def on_workflow_start(self, name: str, instance: object) -> None:
             message, stack_info=False, extra=self._format_details(details=details)
         )
 
-    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."""
         message = f"Index: {self._index_name} -- " if self._index_name else ""
         workflow_progress = (
@@ -161,7 +161,7 @@ def on_workflow_end(self, name: str, instance: object) -> None:
             message, stack_info=False, extra=self._format_details(details=details)
         )
 
-    def on_error(
+    def error(
         self,
         message: str,
         cause: Optional[BaseException] = None,
@@ -178,24 +178,18 @@ def on_error(
             extra=self._format_details(details=details),
         )
 
-    def on_warning(self, message: str, details: Optional[dict] = None) -> None:
+    def warning(self, message: str, details: Optional[dict] = None) -> None:
         """A call back handler for when a warning occurs."""
         self._logger.warning(
             message, stack_info=False, extra=self._format_details(details=details)
         )
 
-    def on_log(self, message: str, details: Optional[dict] = None) -> None:
+    def log(self, message: str, details: Optional[dict] = None) -> None:
         """A call back handler for when a log message occurs."""
         self._logger.info(
             message, stack_info=False, extra=self._format_details(details=details)
         )
 
-    def on_measure(
-        self, name: str, value: float, details: Optional[dict] = None
-    ) -> None:
-        """A call back handler for when a measurement occurs."""
-        raise NotImplementedError("on_measure() not supported by this logger.")
-
 
 def unwrap_dict(input_dict, parent_key="", sep="_"):
     """

backend/src/logger/blob_workflow_callbacks.py

@@ -4,12 +4,11 @@
 from datetime import datetime
 from typing import (
     Any,
-    Optional,
 )
 
 from azure.storage.blob import BlobServiceClient
-from datashaper import NoopWorkflowCallbacks
 from devtools import pformat
+from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 
 
 class BlobWorkflowCallbacks(NoopWorkflowCallbacks):
@@ -72,23 +71,7 @@ def _write_log(self, log: dict[str, Any]):
         blob_client.append_block(pformat(log, indent=2) + "\n")
         self._num_blocks += 1
 
-    def on_error(
-        self,
-        message: str,
-        cause: BaseException | None = None,
-        stack: str | None = None,
-        details: dict | None = None,
-    ):
-        """Report an error."""
-        self._write_log({
-            "type": "error",
-            "data": message,
-            "cause": str(cause),
-            "stack": stack,
-            "details": details,
-        })
-
-    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._workflow_name = name
         self._processed_workflow_steps.append(name)
@@ -111,7 +94,7 @@ def on_workflow_start(self, name: str, instance: object) -> None:
             "details": details,
         })
 
-    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."""
         message = f"Index: {self._index_name} -- " if self._index_name else ""
         workflow_progress = (
@@ -132,16 +115,26 @@ def on_workflow_end(self, name: str, instance: object) -> None:
             "details": details,
         })
 
-    def on_warning(self, message: str, details: dict | None = None):
+    def error(
+        self,
+        message: str,
+        cause: BaseException | None = None,
+        stack: str | None = None,
+        details: dict | None = None,
+    ):
+        """Report an error."""
+        self._write_log({
+            "type": "error",
+            "data": message,
+            "cause": str(cause),
+            "stack": stack,
+            "details": details,
+        })
+
+    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})
-
-    def on_measure(
-        self, name: str, value: float, details: Optional[dict] = None
-    ) -> None:
-        """A call back handler for when a measurement occurs."""
-        pass

backend/src/logger/console_workflow_callbacks.py

@@ -11,7 +11,7 @@
     Optional,
 )
 
-from datashaper.workflow.workflow_callbacks import NoopWorkflowCallbacks
+from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 
 
 class ConsoleWorkflowCallbacks(NoopWorkflowCallbacks):
@@ -107,7 +107,7 @@ def _format_details(self, details: Dict[str, Any] | None = None) -> Dict[str, An
             details = {}
         return {**self._properties, **details}
 
-    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._workflow_name = name
         self._processed_workflow_steps.append(name)
@@ -128,7 +128,7 @@ def on_workflow_start(self, name: str, instance: object) -> None:
             message, stack_info=False, extra=self._format_details(details=details)
         )
 
-    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."""
         message = f"Index: {self._index_name} -- " if self._index_name else ""
         workflow_progress = (
@@ -147,7 +147,7 @@ def on_workflow_end(self, name: str, instance: object) -> None:
             message, stack_info=False, extra=self._format_details(details=details)
         )
 
-    def on_error(
+    def error(
         self,
         message: str,
         cause: Optional[BaseException] = None,
@@ -164,20 +164,14 @@ def on_error(
             extra=self._format_details(details=details),
         )
 
-    def on_warning(self, message: str, details: Optional[dict] = None) -> None:
+    def warning(self, message: str, details: Optional[dict] = None) -> None:
         """A call back handler for when a warning occurs."""
         self._logger.warning(
             message, stack_info=False, extra=self._format_details(details=details)
         )
 
-    def on_log(self, message: str, details: Optional[dict] = None) -> None:
+    def log(self, message: str, details: Optional[dict] = None) -> None:
         """A call back handler for when a log message occurs."""
         self._logger.info(
             message, stack_info=False, extra=self._format_details(details=details)
         )
-
-    def on_measure(
-        self, name: str, value: float, details: Optional[dict] = None
-    ) -> None:
-        """A call back handler for when a measurement occurs."""
-        pass

backend/src/logger/load_logger.py

@@ -5,8 +5,8 @@
 from pathlib import Path
 from typing import List
 
-from datashaper import WorkflowCallbacks, WorkflowCallbacksManager
 from graphrag.callbacks.file_workflow_callbacks import FileWorkflowCallbacks
+from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
 
 from src.api.azure_clients import AzureClientManager
 from src.logger.application_insights_workflow_callbacks import (
@@ -23,7 +23,7 @@ def load_pipeline_logger(
     index_name: str = "",
     num_workflow_steps: int = 0,
 ) -> WorkflowCallbacks:
-    """Create a callback manager and register a list of loggers.
+    """Create and load a list of loggers.
 
     Loggers may be configured as generic loggers or associated with a specified indexing job.
     """
@@ -32,7 +32,7 @@ def load_pipeline_logger(
         reporters.append(Reporters.CONSOLE)
 
     azure_client_manager = AzureClientManager()
-    callback_manager = WorkflowCallbacksManager()
+    logger_callbacks = []
     for reporter in reporters:
         match reporter:
             case Reporters.BLOB:
@@ -48,7 +48,7 @@ def load_pipeline_logger(
                 ).exists():
                     blob_service_client.create_container(container_root)
                 # register the blob reporter
-                callback_manager.register(
+                logger_callbacks.append(
                     BlobWorkflowCallbacks(
                         blob_service_client=blob_service_client,
                         container_name=container_name,
@@ -57,10 +57,10 @@ def load_pipeline_logger(
                     )
                 )
             case Reporters.FILE:
-                callback_manager.register(FileWorkflowCallbacks(dir=reporting_dir))
+                logger_callbacks.append(FileWorkflowCallbacks(dir=reporting_dir))
             case Reporters.APP_INSIGHTS:
                 if os.getenv("APP_INSIGHTS_CONNECTION_STRING"):
-                    callback_manager.register(
+                    logger_callbacks.append(
                         ApplicationInsightsWorkflowCallbacks(
                             connection_string=os.environ[
                                 "APP_INSIGHTS_CONNECTION_STRING"
@@ -70,11 +70,11 @@ def load_pipeline_logger(
                         )
                     )
             case Reporters.CONSOLE:
-                callback_manager.register(
+                logger_callbacks.append(
                     ConsoleWorkflowCallbacks(
                         index_name=index_name, num_workflow_steps=num_workflow_steps
                     )
                 )
             case _:
                 print(f"WARNING: unknown reporter type: {reporter}. Skipping.")
-    return callback_manager
+    return logger_callbacks

backend/src/logger/logger_singleton.py

@@ -4,7 +4,7 @@
 import os
 from urllib.parse import urlparse
 
-from datashaper import WorkflowCallbacks
+from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
 
 from src.logger.load_logger import load_pipeline_logger
 from src.logger.typing import Reporters

backend/src/logger/pipeline_job_workflow_callbacks.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from datashaper.workflow.workflow_callbacks import NoopWorkflowCallbacks
+from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 
 from src.typing.pipeline import PipelineJobState
 from src.utils.pipeline import PipelineJob