From 1ded26c98b651ea53d24470aae7f864b250fc6ce Mon Sep 17 00:00:00 2001 From: elijahbenizzy Date: Sun, 4 Aug 2024 15:42:14 -0700 Subject: [PATCH] Adds documentation for tracing --- burr/visibility/tracing.py | 13 +++++ docs/concepts/additional-visibility.rst | 69 +++++++++++++------------ 2 files changed, 49 insertions(+), 33 deletions(-) diff --git a/burr/visibility/tracing.py b/burr/visibility/tracing.py index 7e6f6ee05..01a7667e4 100644 --- a/burr/visibility/tracing.py +++ b/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, diff --git a/docs/concepts/additional-visibility.rst b/docs/concepts/additional-visibility.rst index 90e3cedfc..f7d2eabe8 100644 --- a/docs/concepts/additional-visibility.rst +++ b/docs/concepts/additional-visibility.rst @@ -17,7 +17,7 @@ the `OpenTelemetry `_ 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 ` or :py:meth:`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 `.