agenta 0.20.0a11__tar.gz → 0.20.0a13__tar.gz

{agenta-0.20.0a11 → agenta-0.20.0a13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.20.0a11
+Version: 0.20.0a13
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Home-page: https://agenta.ai
 Keywords: LLMOps,LLM,evaluation,prompt engineering

{agenta-0.20.0a11 → agenta-0.20.0a13}/agenta/sdk/decorators/tracing.py

@@ -2,7 +2,7 @@
 import inspect
 import traceback
 from functools import wraps
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, List, Union
 
 # Own Imports
 import agenta as ag
@@ -40,9 +40,13 @@ class instrument(BaseDecorator):
         self,
         config: Optional[dict] = None,
         spankind: str = "workflow",
+        ignore_inputs: Union[List[str], bool] = False,
+        ignore_outputs: Union[List[str], bool] = False,
     ) -> None:
         self.config = config
         self.spankind = spankind
+        self.ignore_inputs = ignore_inputs
+        self.ignore_outputs = ignore_outputs
 
     def __call__(self, func: Callable[..., Any]):
         is_coroutine_function = inspect.iscoroutinefunction(func)
@@ -54,38 +58,54 @@ class instrument(BaseDecorator):
 
             return input_dict
 
+        def redact(io, blacklist):
+            return {
+                key: io[key]
+                for key in io.keys()
+                if key
+                not in (
+                    blacklist
+                    if isinstance(blacklist, list)
+                    else []
+                    if blacklist is False
+                    else io.keys()
+                )
+            }
+
+        def patch(result):
+            TRACE_DEFAULT_KEY = "__default__"
+
+            outputs = result
+
+            # PATCH : if result is not a dict, make it a dict
+            if not isinstance(result, dict):
+                outputs = {TRACE_DEFAULT_KEY: result}
+            else:
+                # PATCH : if result is a legacy dict, clean it up
+                if (
+                    "message" in result.keys()
+                    and "cost" in result.keys()
+                    and "usage" in result.keys()
+                ):
+                    outputs = {TRACE_DEFAULT_KEY: result["message"]}
+
+                    ag.tracing.store_cost(result["cost"])
+                    ag.tracing.store_usage(result["usage"])
+
+            return outputs
+
         @wraps(func)
         async def async_wrapper(*args, **kwargs):
             async def wrapped_func(*args, **kwargs):
                 with ag.tracing.Context(
                     name=func.__name__,
-                    input=get_inputs(*args, **kwargs),
+                    input=redact(get_inputs(*args, **kwargs), self.ignore_inputs),
                     spankind=self.spankind,
                     config=self.config,
                 ):
                     result = await func(*args, **kwargs)
 
-                    TRACE_DEFAULT_KEY = "__default__"
-
-                    outputs = result
-
-                    # PATCH : if result is not a dict, make it a dict
-                    if not isinstance(result, dict):
-                        outputs = {TRACE_DEFAULT_KEY: result}
-                    else:
-                        # PATCH : if result is a legacy dict, clean it up
-                        if (
-                            "message" in result.keys()
-                            and "cost" in result.keys()
-                            and "usage" in result.keys()
-                        ):
-                            outputs = {TRACE_DEFAULT_KEY: result["message"]}
-
-                            ag.tracing.store_cost(result["cost"])
-                            ag.tracing.store_usage(result["usage"])
-                    # END OF PATH
-
-                    ag.tracing.store_outputs(outputs)
+                    ag.tracing.store_outputs(redact(patch(result), self.ignore_outputs))
 
                     return result
 
@@ -96,33 +116,13 @@ class instrument(BaseDecorator):
             def wrapped_func(*args, **kwargs):
                 with ag.tracing.Context(
                     name=func.__name__,
-                    input=get_inputs(*args, **kwargs),
+                    input=redact(get_inputs(*args, **kwargs), self.ignore_inputs),
                     spankind=self.spankind,
                     config=self.config,
                 ):
                     result = func(*args, **kwargs)
 
-                    TRACE_DEFAULT_KEY = "__default__"
-
-                    outputs = result
-
-                    # PATCH : if result is not a dict, make it a dict
-                    if not isinstance(result, dict):
-                        outputs = {TRACE_DEFAULT_KEY: result}
-                    else:
-                        # PATCH : if result is a legacy dict, clean it up
-                        if (
-                            "message" in result.keys()
-                            and "cost" in result.keys()
-                            and "usage" in result.keys()
-                        ):
-                            outputs = {"message": result["message"]}
-
-                            ag.tracing.store_cost(result["cost"])
-                            ag.tracing.store_usage(result["usage"])
-                    # END OF PATH
-
-                    ag.tracing.store_outputs(outputs)
+                    ag.tracing.store_outputs(redact(patch(result), self.ignore_outputs))
 
                     return result
 

{agenta-0.20.0a11 → agenta-0.20.0a13}/agenta/sdk/tracing/callbacks.py

@@ -66,9 +66,10 @@ def litellm_handler():
         def log_stream_event(self, kwargs, response_obj, start_time, end_time):
             logging.info(f"log_stream_event({self.span.id})")
             ag.tracing.set_status(status="OK", span_id=self.span.id)
-            ag.tracing.store_cost(kwargs.get("response_cost"))
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
             ag.tracing.store_usage(
-                response_obj.usage.dict() if hasattr(response_obj, "usage") else None
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
             )
             ag.tracing.store_outputs(
                 # the complete streamed response (only set if `completion(..stream=True)`
@@ -81,9 +82,10 @@ def litellm_handler():
         def log_success_event(self, kwargs, response_obj, start_time, end_time):
             logging.info(f"log_success_event({self.span.id})")
             ag.tracing.set_status(status="OK", span_id=self.span.id)
-            ag.tracing.store_cost(kwargs.get("response_cost"))
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
             ag.tracing.store_usage(
-                response_obj.usage.dict() if hasattr(response_obj, "usage") else None
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
             )
             ag.tracing.store_outputs(
                 outputs={TRACE_DEFAULT_KEY: response_obj.choices[0].message.content},
@@ -106,9 +108,10 @@ def litellm_handler():
                 },
                 span_id=self.span.id,
             )
-            ag.tracing.store_cost(kwargs.get("response_cost"))
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
             ag.tracing.store_usage(
-                response_obj.usage.dict() if hasattr(response_obj, "usage") else None
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
             )
             ag.tracing.store_outputs(
                 # the Exception raised
@@ -123,9 +126,10 @@ def litellm_handler():
         ):
             logging.info(f"async_log_stream_event({self.span.id})")
             ag.tracing.set_status(status="OK", span_id=self.span.id)
-            ag.tracing.store_cost(kwargs.get("response_cost"))
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
             ag.tracing.store_usage(
-                response_obj.usage.dict() if hasattr(response_obj, "usage") else None
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
             )
             ag.tracing.store_outputs(
                 # the complete streamed response (only set if `completion(..stream=True)`)
@@ -140,9 +144,10 @@ def litellm_handler():
         ):
             logging.info(f"async_log_success_event({self.span.id})")
             ag.tracing.set_status(status="OK", span_id=self.span.id)
-            ag.tracing.store_cost(kwargs.get("response_cost"))
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
             ag.tracing.store_usage(
-                response_obj.usage.dict() if hasattr(response_obj, "usage") else None
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
             )
             ag.tracing.store_outputs(
                 outputs={TRACE_DEFAULT_KEY: response_obj.choices[0].message.content},
@@ -167,9 +172,10 @@ def litellm_handler():
                 },
                 span_id=self.span.id,
             )
-            ag.tracing.store_cost(kwargs.get("response_cost"))
+            ag.tracing.store_cost(kwargs.get("response_cost"), span_id=self.span.id)
             ag.tracing.store_usage(
-                response_obj.usage.dict() if hasattr(response_obj, "usage") else None
+                response_obj.usage.dict() if hasattr(response_obj, "usage") else None,
+                span_id=self.span.id,
             )
             ag.tracing.store_outputs(
                 # the Exception raised

{agenta-0.20.0a11 → agenta-0.20.0a13}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agenta"
-version = "0.20.0a11"
+version = "0.20.0a13"
 description = "The SDK for agenta is an open-source LLMOps platform."
 readme = "README.md"
 authors = ["Mahmoud Mabrouk <mahmoud@agenta.ai>"]

agenta-0.20.0a11/agenta/sdk/tracing/context_manager.py

@@ -1,13 +0,0 @@
-# Own Imports
-from agenta.sdk.tracing.llm_tracing import Tracing
-
-
-class TracingContextManager:
-    def __init__(self, tracing: Tracing):
-        ...
-
-    def __enter__(self):
-        ...
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        ...