feat:mem0 add hook extension

Change-Id: I9696dfbe580630ed6a317e851882bc652385ac1d

Author: sunym1991

CHANGELOG-loongsuite.md

@@ -13,9 +13,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
-- `loongsuite-instrumentation-mem0`: add hook extension
-  ([#95](https://github.com/alibaba/loongsuite-python-agent/pull/95))
-
 - `loongsuite-instrumentation-mem0`: use memory handler
   ([#89](https://github.com/alibaba/loongsuite-python-agent/pull/89))
 
@@ -24,5 +21,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Added
 
+- `loongsuite-instrumentation-mem0`: add hook extension
+  ([#95](https://github.com/alibaba/loongsuite-python-agent/pull/95))
+  
 - `loongsuite-instrumentation-mem0`: add support for mem0
   ([#67](https://github.com/alibaba/loongsuite-python-agent/pull/67))

instrumentation-loongsuite/loongsuite-instrumentation-mem0/src/opentelemetry/instrumentation/mem0/__init__.py

@@ -23,6 +23,7 @@
     VectorStoreWrapper,
 )
 from opentelemetry.instrumentation.mem0.package import _instruments
+from opentelemetry.instrumentation.mem0.types import set_memory_hooks
 from opentelemetry.instrumentation.mem0.version import __version__
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.semconv.schemas import Schemas
@@ -439,7 +440,8 @@ def _instrument_memory_operations(
                 return
 
             wrapper = MemoryOperationWrapper(telemetry_handler)
-            wrapper.set_hooks(
+            set_memory_hooks(
+                wrapper,
                 memory_before_hook=memory_before_hook,
                 memory_after_hook=memory_after_hook,
             )
@@ -501,7 +503,8 @@ def _instrument_memory_client_operations(
                 return
 
             wrapper = MemoryOperationWrapper(telemetry_handler)
-            wrapper.set_hooks(
+            set_memory_hooks(
+                wrapper,
                 memory_before_hook=memory_before_hook,
                 memory_after_hook=memory_after_hook,
             )

instrumentation-loongsuite/loongsuite-instrumentation-mem0/src/opentelemetry/instrumentation/mem0/internal/_wrapper.py

@@ -6,7 +6,7 @@
 import inspect
 import logging
 from functools import wraps
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Optional
 
 from opentelemetry.instrumentation.mem0.config import (
     is_internal_phases_enabled,
@@ -26,6 +26,14 @@
     SemanticAttributes,
     SpanName,
 )
+from opentelemetry.instrumentation.mem0.types import (
+    HookContext,
+    InnerAfterHook,
+    InnerBeforeHook,
+    MemoryAfterHook,
+    MemoryBeforeHook,
+    safe_call_hook,
+)
 from opentelemetry.trace import (
     SpanKind,
     Status,
@@ -39,28 +47,6 @@
 
 logger = logging.getLogger(__name__)
 
-# Per-call hook context. Instrumentation only creates and passes it through.
-HookContext = Dict[str, Any]
-
-# Hook types are intentionally kept loose here to avoid coupling this package's runtime
-# to optional extension modules / type-checking configuration. Hooks are pure pass-through.
-MemoryBeforeHook = Optional[Callable[..., Any]]
-MemoryAfterHook = Optional[Callable[..., Any]]
-InnerBeforeHook = Optional[Callable[..., Any]]
-InnerAfterHook = Optional[Callable[..., Any]]
-
-
-def _safe_call_hook(hook: Optional[Callable], *args: Any) -> None:
-    """
-    Call a hook defensively: swallow hook exceptions to avoid breaking user code.
-    """
-    if not callable(hook):
-        return
-    try:
-        hook(*args)
-    except Exception as e:  # pragma: no cover - defensive
-        logger.debug("mem0 hook raised and was swallowed: %s", e)
-
 
 def _get_field(payload: dict, field_name: str) -> Any:
     """
@@ -180,20 +166,6 @@ def __init__(self, telemetry_handler: ExtendedTelemetryHandler):
         self._memory_before_hook: MemoryBeforeHook = None
         self._memory_after_hook: MemoryAfterHook = None
 
-    def set_hooks(
-        self,
-        *,
-        memory_before_hook: MemoryBeforeHook = None,
-        memory_after_hook: MemoryAfterHook = None,
-    ) -> None:
-        """
-        Set optional hooks for top-level memory operations.
-
-        Hooks are stored on the wrapper instance to avoid changing wrapt wrapper signatures.
-        """
-        self._memory_before_hook = memory_before_hook
-        self._memory_after_hook = memory_after_hook
-
     def wrap_operation(
         self,
         operation_name: str,
@@ -491,7 +463,7 @@ def _execute_with_handler(
         hook_context: HookContext = {}
         # Read current span after util handler starts memory (span should exist in most cases)
         span = get_current_span()
-        _safe_call_hook(
+        safe_call_hook(
             self._memory_before_hook,
             span,
             operation_name,
@@ -512,7 +484,7 @@ def _execute_with_handler(
                 extract_attributes_func=extract_attributes_func,
                 is_memory_client=is_memory_client,
             )
-            _safe_call_hook(
+            safe_call_hook(
                 self._memory_after_hook,
                 span,
                 operation_name,
@@ -526,7 +498,7 @@ def _execute_with_handler(
             self.telemetry_handler.stop_memory(invocation)
             return result
         except Exception as e:
-            _safe_call_hook(
+            safe_call_hook(
                 self._memory_after_hook,
                 span,
                 operation_name,
@@ -582,7 +554,7 @@ async def _execute_with_handler_async(
         self.telemetry_handler.start_memory(invocation)
         hook_context: HookContext = {}
         span = get_current_span()
-        _safe_call_hook(
+        safe_call_hook(
             self._memory_before_hook,
             span,
             operation_name,
@@ -603,7 +575,7 @@ async def _execute_with_handler_async(
                 extract_attributes_func=extract_attributes_func,
                 is_memory_client=is_memory_client,
             )
-            _safe_call_hook(
+            safe_call_hook(
                 self._memory_after_hook,
                 span,
                 operation_name,
@@ -617,7 +589,7 @@ async def _execute_with_handler_async(
             self.telemetry_handler.stop_memory(invocation)
             return result
         except Exception as e:
-            _safe_call_hook(
+            safe_call_hook(
                 self._memory_after_hook,
                 span,
                 operation_name,
@@ -686,7 +658,7 @@ def wrapper(
             ) as span:
                 result = None
                 hook_context: HookContext = {}
-                _safe_call_hook(
+                safe_call_hook(
                     self._inner_before_hook,
                     span,
                     "vector",
@@ -712,7 +684,7 @@ def wrapper(
                         span.set_attribute(key, value)
 
                     span.set_status(Status(StatusCode.OK))
-                    _safe_call_hook(
+                    safe_call_hook(
                         self._inner_after_hook,
                         span,
                         "vector",
@@ -732,7 +704,7 @@ def wrapper(
                     span.set_attribute(
                         SemanticAttributes.ERROR_TYPE, get_exception_type(e)
                     )
-                    _safe_call_hook(
+                    safe_call_hook(
                         self._inner_after_hook,
                         span,
                         "vector",
@@ -801,7 +773,7 @@ def wrapper(
             ) as span:
                 result = None
                 hook_context: HookContext = {}
-                _safe_call_hook(
+                safe_call_hook(
                     self._inner_before_hook,
                     span,
                     "graph",
@@ -827,7 +799,7 @@ def wrapper(
                         span.set_attribute(key, value)
 
                     span.set_status(Status(StatusCode.OK))
-                    _safe_call_hook(
+                    safe_call_hook(
                         self._inner_after_hook,
                         span,
                         "graph",
@@ -847,7 +819,7 @@ def wrapper(
                     span.set_attribute(
                         SemanticAttributes.ERROR_TYPE, get_exception_type(e)
                     )
-                    _safe_call_hook(
+                    safe_call_hook(
                         self._inner_after_hook,
                         span,
                         "graph",
@@ -919,7 +891,7 @@ def wrapper(
                 kind=SpanKind.CLIENT,
             ) as span:
                 hook_context: HookContext = {}
-                _safe_call_hook(
+                safe_call_hook(
                     self._inner_before_hook,
                     span,
                     "rerank",
@@ -944,7 +916,7 @@ def wrapper(
                     result = wrapped(*args, **kwargs)
 
                     span.set_status(Status(StatusCode.OK))
-                    _safe_call_hook(
+                    safe_call_hook(
                         self._inner_after_hook,
                         span,
                         "rerank",
@@ -964,7 +936,7 @@ def wrapper(
                     span.set_attribute(
                         SemanticAttributes.ERROR_TYPE, get_exception_type(e)
                     )
-                    _safe_call_hook(
+                    safe_call_hook(
                         self._inner_after_hook,
                         span,
                         "rerank",

instrumentation-loongsuite/loongsuite-instrumentation-mem0/src/opentelemetry/instrumentation/mem0/types.py

@@ -0,0 +1,45 @@
+"""
+Mem0 instrumentation public hook types and helpers.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+# Per-call hook context. Instrumentation only creates and passes it through.
+HookContext = Dict[str, Any]
+
+# Hook callables are kept intentionally loose: the open-source package only passes through
+# parameters, and commercial extensions are responsible for extracting/recording data.
+MemoryBeforeHook = Optional[Callable[..., Any]]
+MemoryAfterHook = Optional[Callable[..., Any]]
+InnerBeforeHook = Optional[Callable[..., Any]]
+InnerAfterHook = Optional[Callable[..., Any]]
+
+
+def safe_call_hook(hook: Optional[Callable[..., Any]], *args: Any) -> None:
+    """
+    Call a hook defensively: swallow hook exceptions to avoid breaking user code.
+    """
+    if not callable(hook):
+        return
+    try:
+        hook(*args)
+    except Exception as e:
+        logger.debug("mem0 hook raised and was swallowed: %s", e)
+
+
+def set_memory_hooks(
+    wrapper: Any,
+    *,
+    memory_before_hook: MemoryBeforeHook = None,
+    memory_after_hook: MemoryAfterHook = None,
+) -> None:
+    """
+    Configure top-level memory hooks on a MemoryOperationWrapper instance.
+    """
+    wrapper._memory_before_hook = memory_before_hook
+    wrapper._memory_after_hook = memory_after_hook

instrumentation-loongsuite/loongsuite-instrumentation-mem0/tests/test_hooks.py

@@ -22,6 +22,7 @@
     RerankerWrapper,
     VectorStoreWrapper,
 )
+from opentelemetry.instrumentation.mem0.types import set_memory_hooks
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.util.genai._extended_memory import MemoryInvocation
 
@@ -56,7 +57,7 @@ def after(
         calls.append(("after", hook_context))
 
     w = MemoryOperationWrapper(_DummyTelemetryHandler())
-    w.set_hooks(memory_before_hook=before, memory_after_hook=after)
+    set_memory_hooks(w, memory_before_hook=before, memory_after_hook=after)
 
     # Cover helper: _normalize_call_parameters with positional args mapping
     def _sig(self, memory_id, data, *, user_id=None):  # noqa: ARG001
@@ -165,7 +166,9 @@ def after_boom(
     ):
         raise RuntimeError("boom2")
 
-    w.set_hooks(memory_before_hook=before_boom, memory_after_hook=after_boom)
+    set_memory_hooks(
+        w, memory_before_hook=before_boom, memory_after_hook=after_boom
+    )
     assert (
         w._execute_with_handler(
             _fn_sync,
@@ -180,7 +183,7 @@ def after_boom(
     )
 
     # exception path calls after_hook with exception
-    w.set_hooks(memory_before_hook=before, memory_after_hook=after)
+    set_memory_hooks(w, memory_before_hook=before, memory_after_hook=after)
 
     def _fn_raises(*a, **k):
         raise ValueError("nope")