Merge pull request #2269 from Agenta-AI/feature/allow-custom-redact

[Feautre] Add custom redact functionality to SDK

agenta-cli/agenta/__init__.py

@@ -1,3 +1,5 @@
+from typing import Any, Callable, Optional
+
 from .sdk.utils.preinit import PreInitObject
 
 import agenta.client.backend.types as client_types  # pylint: disable=wrong-import-order
@@ -18,7 +20,7 @@
 )
 
 from .sdk.utils.logging import log as logging
-from .sdk.tracing import Tracing
+from .sdk.tracing import Tracing, get_tracer
 from .sdk.decorators.tracing import instrument
 from .sdk.tracing.conventions import Reference
 from .sdk.decorators.routing import entrypoint, app, route
@@ -36,15 +38,36 @@
 DEFAULT_AGENTA_SINGLETON_INSTANCE = AgentaSingleton()
 
 types = client_types
-tracing = None
+
 api = None
 async_api = None
 
+tracing = DEFAULT_AGENTA_SINGLETON_INSTANCE.tracing  # type: ignore
+tracer = get_tracer(tracing)
 
-def init(*args, **kwargs):
-    global api, async_api, tracing, config
-    _init(*args, **kwargs)
 
-    tracing = DEFAULT_AGENTA_SINGLETON_INSTANCE.tracing  # type: ignore
+def init(
+    host: Optional[str] = None,
+    api_key: Optional[str] = None,
+    config_fname: Optional[str] = None,
+    redact: Optional[Callable[..., Any]] = None,
+    redact_on_error: Optional[bool] = True,
+    # DEPRECATING
+    app_id: Optional[str] = None,
+):
+    global api, async_api, tracing, tracer  # pylint: disable=global-statement
+
+    _init(
+        host=host,
+        api_key=api_key,
+        config_fname=config_fname,
+        redact=redact,
+        redact_on_error=redact_on_error,
+        app_id=app_id,
+    )
+
     api = DEFAULT_AGENTA_SINGLETON_INSTANCE.api  # type: ignore
     async_api = DEFAULT_AGENTA_SINGLETON_INSTANCE.async_api  # type: ignore
+
+    tracing = DEFAULT_AGENTA_SINGLETON_INSTANCE.tracing  # type: ignore
+    tracer = get_tracer(tracing)

agenta-cli/agenta/sdk/__init__.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Callable, Any
 
 from .utils.preinit import PreInitObject  # always the first import!
 
@@ -43,17 +43,21 @@
 
 def init(
     host: Optional[str] = None,
-    app_id: Optional[str] = None,
     api_key: Optional[str] = None,
     config_fname: Optional[str] = None,
+    redact: Optional[Callable[..., Any]] = None,
+    redact_on_error: Optional[bool] = True,
+    # DEPRECATING
+    app_id: Optional[str] = None,
 ):
-    global api, async_api, tracing, tracer
+    global api, async_api, tracing, tracer  # pylint: disable=global-statement
 
     _init(
         host=host,
         api_key=api_key,
         config_fname=config_fname,
-        # DEPRECATING
+        redact=redact,
+        redact_on_error=redact_on_error,
         app_id=app_id,
     )
 

agenta-cli/agenta/sdk/agenta_init.py

@@ -1,7 +1,7 @@
 import logging
 import toml
 from os import getenv
-from typing import Optional
+from typing import Optional, Callable, Any
 from importlib.metadata import version
 
 from agenta.sdk.utils.logging import log
@@ -36,6 +36,8 @@ def init(
         host: Optional[str] = None,
         api_key: Optional[str] = None,
         config_fname: Optional[str] = None,
+        redact: Optional[Callable[..., Any]] = None,
+        redact_on_error: Optional[bool] = True,
         # DEPRECATING
         app_id: Optional[str] = None,
     ) -> None:
@@ -91,6 +93,8 @@ def init(
 
         self.tracing = Tracing(
             url=f"{self.host}/api/observability/v1/otlp/traces",  # type: ignore
+            redact=redact,
+            redact_on_error=redact_on_error,
         )
 
         self.tracing.configure(
@@ -258,7 +262,9 @@ def init(
     host: Optional[str] = None,
     api_key: Optional[str] = None,
     config_fname: Optional[str] = None,
-    # DEPRECATED
+    redact: Optional[Callable[..., Any]] = None,
+    redact_on_error: Optional[bool] = True,
+    # DEPRECATING
     app_id: Optional[str] = None,
 ):
     """Main function to initialize the agenta sdk.
@@ -289,7 +295,8 @@ def init(
         host=host,
         api_key=api_key,
         config_fname=config_fname,
-        # DEPRECATED
+        redact=redact,
+        redact_on_error=redact_on_error,
         app_id=app_id,
     )
 

agenta-cli/agenta/sdk/decorators/tracing.py

@@ -19,6 +19,8 @@ def __init__(
         config: Optional[Dict[str, Any]] = None,
         ignore_inputs: Optional[bool] = None,
         ignore_outputs: Optional[bool] = None,
+        redact: Optional[Callable[..., Any]] = None,
+        redact_on_error: Optional[bool] = True,
         max_depth: Optional[int] = 2,
         # DEPRECATING
         kind: str = "task",
@@ -29,6 +31,8 @@ def __init__(
         self.config = config
         self.ignore_inputs = ignore_inputs
         self.ignore_outputs = ignore_outputs
+        self.redact = redact
+        self.redact_on_error = redact_on_error
         self.max_depth = max_depth
 
     def __call__(self, func: Callable[..., Any]):
@@ -109,12 +113,10 @@ def _pre_instrument(
                 )
 
             _inputs = self._redact(
-                self._parse(
-                    func,
-                    *args,
-                    **kwargs,
-                ),
-                self.ignore_inputs,
+                name=span.name,
+                field="inputs",
+                io=self._parse(func, *args, **kwargs),
+                ignore=self.ignore_inputs,
             )
             span.set_attributes(
                 attributes={"inputs": _inputs},
@@ -153,7 +155,12 @@ def _post_instrument(
                 namespace="metrics.unit.tokens",
             )
 
-            _outputs = self._redact(self._patch(result), self.ignore_outputs)
+            _outputs = self._redact(
+                name=span.name,
+                field="outputs",
+                io=self._patch(result),
+                ignore=self.ignore_outputs,
+            )
             span.set_attributes(
                 attributes={"outputs": _outputs},
                 namespace="data",
@@ -192,6 +199,9 @@ def _parse(
 
     def _redact(
         self,
+        *,
+        name: str,
+        field: str,
         io: Dict[str, Any],
         ignore: Union[List[str], bool] = False,
     ) -> Dict[str, Any]:
@@ -220,6 +230,20 @@ def _redact(
             )
         }
 
+        if self.redact is not None:
+            try:
+                io = self.redact(name, field, io)
+            except:  # pylint: disable=bare-except
+                if self.redact_on_error:
+                    io = {}
+
+        if ag.tracing.redact is not None:
+            try:
+                io = ag.tracing.redact(name, field, io)
+            except:  # pylint: disable=bare-except
+                if ag.tracing.redact_on_error:
+                    io = {}
+
         return io
 
     def _patch(

agenta-cli/agenta/sdk/tracing/tracing.py

@@ -1,4 +1,4 @@
-from typing import Optional, Any, Dict
+from typing import Optional, Any, Dict, Callable
 from enum import Enum
 
 from httpx import get as check
@@ -32,6 +32,8 @@ class Tracing(metaclass=Singleton):
     def __init__(
         self,
         url: str,
+        redact: Optional[Callable[..., Any]] = None,
+        redact_on_error: Optional[bool] = True,
     ) -> None:
         # ENDPOINT (OTLP)
         self.otlp_url = url
@@ -49,6 +51,10 @@ def __init__(
         # INLINE SPANS for INLINE TRACES (INLINE PROCESSOR)
         self.inline_spans: Dict[str, Any] = dict()
 
+        # REDACT
+        self.redact = redact
+        self.redact_on_error = redact_on_error
+
     # PUBLIC
 
     def configure(

agenta-cli/tests/redact/01_ignore_all.py

@@ -0,0 +1,18 @@
+import agenta as ag
+
+ag.init()
+
+
+@ag.entrypoint
+@ag.instrument(
+    spankind="WORKFLOW",
+    ignore_inputs=True,
+    ignore_outputs=True,
+)
+def embed(description: str):
+    return {
+        "embedding": "somedata",
+        "ignored": "ignored",
+        "cost": 15,
+        "usage": 20,
+    }

agenta-cli/tests/redact/02_ignore_all_inputs.py

@@ -0,0 +1,17 @@
+import agenta as ag
+
+ag.init()
+
+
+@ag.entrypoint
+@ag.instrument(
+    spankind="WORKFLOW",
+    ignore_inputs=True,
+)
+def embed(description: str):
+    return {
+        "embedding": "somedata",
+        "ignored": "ignored",
+        "cost": 15,
+        "usage": 20,
+    }

agenta-cli/tests/redact/03_ignore_all_outputs.py

@@ -0,0 +1,17 @@
+import agenta as ag
+
+ag.init()
+
+
+@ag.entrypoint
+@ag.instrument(
+    spankind="WORKFLOW",
+    ignore_inputs=True,
+)
+def embed(description: str):
+    return {
+        "embedding": "somedata",
+        "ignored": "ignored",
+        "cost": 15,
+        "usage": 20,
+    }

agenta-cli/tests/redact/04_ignore_some.py

@@ -0,0 +1,18 @@
+import agenta as ag
+
+ag.init()
+
+
+@ag.entrypoint
+@ag.instrument(
+    spankind="WORKFLOW",
+    ignore_inputs=["description"],
+    ignore_outputs=["embedding", "ignored"],
+)
+def embed(description: str, theme: str):
+    return {
+        "embedding": "somedata",
+        "ignored": "ignored",
+        "cost": 15,
+        "usage": 20,
+    }

agenta-cli/tests/redact/05_redact_instrument.py

@@ -0,0 +1,30 @@
+import agenta as ag
+
+ag.init()
+
+
+def redact(name, field, io):
+    print(">", name, field, io)
+
+    if name == "embed" and field == "inputs":
+        io = {key: value for key, value in io.items() if key not in ("description",)}
+
+    if name == "embed" and field == "outputs":
+        io = {key: value for key, value in io.items() if key not in ("embedding",)}
+
+    print("<", io)
+    return io
+
+
+@ag.entrypoint
+@ag.instrument(
+    spankind="WORKFLOW",
+    redact=redact,
+)
+def embed(description: str, theme: str):
+    return {
+        "embedding": "somedata",
+        "ignored": "ignored",
+        "cost": 15,
+        "usage": 20,
+    }

agenta-cli/tests/redact/06_redact_init.py

@@ -0,0 +1,31 @@
+import agenta as ag
+
+
+def redact(name, field, io):
+    print(">", name, field, io)
+
+    if name == "embed" and field == "inputs":
+        io = {key: value for key, value in io.items() if key not in ("description",)}
+
+    if name == "embed" and field == "outputs":
+        io = {key: value for key, value in io.items() if key not in ("embedding",)}
+
+    print("<", io)
+
+    return io
+
+
+ag.init(redact=redact)
+
+
+@ag.entrypoint
+@ag.instrument(
+    spankind="WORKFLOW",
+)
+def embed(description: str, theme: str):
+    return {
+        "embedding": "somedata",
+        "ignored": "ignored",
+        "cost": 15,
+        "usage": 20,
+    }