Adds documentation for tracing

burr/visibility/tracing.py

@@ -210,9 +210,22 @@ def __exit__(
         return None
 
     def log_attribute(self, key: str, value: Any):
+        """Logs a single attribute to the UI. Note that this must
+        be paired with a tracker or a tracking hook to be useful, otherwise it
+        will be a no-op.
+
+        :param key: Name of the attribute (must be unique per action/span)
+        :param value: Value of the attribute.
+        """
         self.log_attributes(**{key: value})
 
     def log_attributes(self, **attributes):
+        """Logs a set of attributes to the UI. Note that this must
+        be paired with a tracker or a tracking hook to be useful, otherwise it
+        will be a no-op.
+
+        :param attributes: Attributes to log
+        """
         self.lifecycle_adapters.call_all_lifecycle_hooks_sync(
             "do_log_attributes",
             attributes=attributes,

docs/concepts/additional-visibility.rst

@@ -17,7 +17,7 @@ the `OpenTelemetry <https://opentelemetry.io/>`_ sdk, although it is simplified
 To add the tracing capability, the action first has to declare a ``__tracer`` input. This is a
 an object that instantiates spans and tracks state.
 
-Then, using the `__tracer` as a callable, you can instantiate spans and track state.
+Then, using the ``__tracer`` as a callable, you can instantiate spans and track state.
 
 For the function-based API, this would look as follows:
 
@@ -26,22 +26,23 @@ For the function-based API, this would look as follows:
     from burr.visibility import TracingFactory
     from burr.core import action
 
-    @action(reads=['input_var'], writes=['output_var'])
+    @action(reads=['prompt'], writes=['response'])
     def my_action(state: State, __tracer: TracingFactory) -> State:
-        with __tracer('process_data'):
-            initial_data = _process_data(state['input_var'])
-            with __tracer('validate_data'):
-                _validate(initial_data)
-        with __tracer('transform_data', dependencies=['process_data']):
-            transformed_data = _transform(initial_data)
-        return state.update({'output_var': transformed_data})
+        with __tracer('create_prompt'):
+            modified_prompt = _modify_prompt(state["prompt"])
+            with __tracer('check_prompt_safety'):
+                if _is_unsafe(modified_prompt):
+                    modified_prompt = _fix(modified_prompt)
+        with __tracer('call_llm', dependencies=['create_prompt']):
+            response = _query(prompt=modified_prompt)
+        return state.update({'response': response})
 
 
 This would create the following traces:
 
-#. ``process_data``
-#. ``validate_data`` as a child of ``process_data``
-#. ``transform_data`` as a causal dependent of ``process_data``
+#. ``create_prompt``
+#. ``check_prompt_safety`` as a child of ``create_prompt``
+#. ``call_llm`` as a causal dependent of ``create_prompt``
 
 Dependencies are used to express [dag](-style structures of spans within actions. This is useful for gaining visibility into the internal structure
 of an action, but is likely best used with integrations with micro-orchestration systems for implementating actions, such as Hamilton or Lanchain.
@@ -62,33 +63,35 @@ These just function as callbacks (on enter/exit). The :py:class:`LocalTrackingCl
 Observations
 ------------
 
-(This is a work in progress, and is not complete)
+You can make observations on the state by calling out to the :py:meth:`log_attribute <burr.visibility.ActionSpanTracer.log_attribute>` or :py:meth:`log_attributes <burr.visibility.ActionSpanTracer.log_attributes>` method on the ``__tracer`` context manager
+or the root tracer factory.
+
+This allows you to log any arbitrary observations (think token count prompt, whatnot) that may not be part of state/results.
 
-You can make observations on the state by calling out to the `log_artifact` method on the `__tracer` context manager.
 For instance:
 
 .. code-block:: python
 
-    from burr.visibility import TracingFactory, ArtifactLogger
+    from burr.visibility import TracingFactory
     from burr.core import action
 
-    @action(reads=['input_var'], writes=['output_var'])
-    def my_action(
-        state: State,
-        __tracer: TracingFactory,
-        __logger: ArtifactLogger
-        ) -> State:
-        with __tracer('process_data'):
-            initial_data = _process_data(state['input_var'])
-            with __tracer('validate_data'):
-                validation_results = _validate(initial_data)
-                t.log_artifact(validation_results=validation_results)
-        with __tracer('transform_data', dependencies=['process_data'])
-            transformed_data = _transform(initial_data)
-            __logger.log_artifact(transformed_data_size=len(transformed_data))
-
-        return state.update(output_var=transformed_data)
-
-The output can be any "json-dumpable" object (or pydantic model). This will be stored along with the span and can be used for debugging or analysis.
+    @action(reads=['prompt'], writes=['response'])
+    def my_action(state: State, __tracer: TracingFactory) -> State:
+        __tracer.log_attribute(prompt_length=len(state["prompt"]), prompt=state["prompt"])
+        with __tracer('create_prompt') as t:
+            modified_prompt = _modify_prompt(state["prompt"])
+            t.log_attribute(modified_prompt=modified_prompt)
+            with __tracer('check_prompt_safety') as t:
+                if is_unsafe:=_is_unsafe(modified_prompt):
+                    modified_prompt = _fix(modified_prompt)
+                t.log_attribute(fixed_prompt=modified_prompt, is_unsafe=is_unsafe)
+        with __tracer('call_llm', dependencies=['create_prompt']):
+            response = _query(prompt=modified_prompt)
+            t.log_attribute(response=response.message, tokens=response.tokens)
+        return state.update({'response': response.message})
+
+The above will log quite a few attributes, prompt length, response tokens, etc... The observation can be any serializable object.
+
+Note that we are currently building out the capability to wrap a class and "auto-log" standard attributes.
 
 You can read more in the :ref:`reference documentation <visibility>`.