ContextQMD
Back to Llama Index

Instrumentation: Basic Usage

docs/examples/instrumentation/basic_usage.ipynb

0.14.216.1 KB
Original Source

Instrumentation: Basic Usage

The instrumentation module can be used for observability and monitoring of your llama-index application. It is comprised of the following core abstractions:

  • Event — represents a single moment in time that a certain occurrence took place within the execution of the application’s code.
  • EventHandler — listen to the occurrences of Event's and execute code logic at these moments in time.
  • Span — represents the execution flow of a particular part in the application’s code and thus contains Event's.
  • SpanHandler — is responsible for the entering, exiting, and dropping (i.e., early exiting due to error) of Span's.
  • Dispatcher — emits Event's as well as signals to enter/exit/drop a Span to the appropriate handlers.

In this notebook, we demonstrate the basic usage pattern of instrumentation:

  1. Define your custom EventHandler
  2. Define your custom SpanHandler which handles an associated Span type
  3. Attach your EventHandler and SpanHandler to the dispatcher of choice (here, we'll attach it to the root dispatcher).
python
import nest_asyncio

nest_asyncio.apply()

Custom Event Handlers

python
from llama_index.core.instrumentation.event_handlers import BaseEventHandler

Defining your custom EventHandler involves subclassing the BaseEventHandler. Doing so, requires defining logic for the abstract method handle().

python
class MyEventHandler(BaseEventHandler):
    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "MyEventHandler"

    def handle(self, event) -> None:
        """Logic for handling event."""
        # THIS IS WHERE YOU ADD YOUR LOGIC TO HANDLE EVENTS
        print(event.dict())
        print("")
        with open("log.txt", "a") as f:
            f.write(str(event))
            f.write("\n")

Custom Span Handlers

SpanHandler also involve subclassing a base class, in this case BaseSpanHandler. However, since SpanHandler's work with an associated Span type, you will need to create this as well if you want to handle a new Span type.

python
from llama_index.core.instrumentation.span import BaseSpan
python
import inspect
from typing import Any, Dict, Optional
from llama_index.core.bridge.pydantic import Field
from llama_index.core.instrumentation.span.base import BaseSpan
from llama_index.core.instrumentation.span_handlers import BaseSpanHandler


class MyCustomSpan(BaseSpan):
    custom_field_1: Any = Field(...)
    custom_field_2: Any = Field(...)


class MyCustomSpanHandler(BaseSpanHandler[MyCustomSpan]):
    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "MyCustomSpanHandler"

    def new_span(
        self,
        id_: str,
        bound_args: inspect.BoundArguments,
        instance: Optional[Any] = None,
        parent_span_id: Optional[str] = None,
        tags: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Optional[MyCustomSpan]:
        """Create a span."""
        # logic for creating a new MyCustomSpan
        pass

    def prepare_to_exit_span(
        self,
        id_: str,
        bound_args: inspect.BoundArguments,
        instance: Optional[Any] = None,
        result: Optional[Any] = None,
        **kwargs: Any,
    ) -> Any:
        """Logic for preparing to exit a span."""
        pass

    def prepare_to_drop_span(
        self,
        id_: str,
        bound_args: inspect.BoundArguments,
        instance: Optional[Any] = None,
        err: Optional[BaseException] = None,
        **kwargs: Any,
    ) -> Any:
        """Logic for preparing to drop a span."""
        pass

For this notebook, we'll use SimpleSpanHandler that works with the SimpleSpan type.

python
from llama_index.core.instrumentation.span_handlers import SimpleSpanHandler

Dispatcher

Now that we have our EventHandler and our SpanHandler, we can attach it to a Dispatcher that will emit Event's and signals to start/exit/drop a Span to the appropriate handlers. Those that are familiar with Logger from the logging Python module, might notice that Dispatcher adopts a similar interface. What's more is that Dispatcher also utilizes a similar hierarchy and propagation scheme as Logger. Specifically, a dispatcher will emit Event's to its handlers and by default propagate these events to its parent Dispatcher for it to send to its own handlers.

python
import llama_index.core.instrumentation as instrument

dispatcher = instrument.get_dispatcher()  # modify root dispatcher
python
span_handler = SimpleSpanHandler()

dispatcher.add_event_handler(MyEventHandler())
dispatcher.add_span_handler(span_handler)

You can also get dispatcher's by name. Purely for the sake of demonstration, in the cells below we get the dispatcher that is defined in the base.base_query_engine submodule of llama_index.core.

python
qe_dispatcher = instrument.get_dispatcher(
    "llama_index.core.base.base_query_engine"
)
python
qe_dispatcher
python
qe_dispatcher.parent

Test It Out

python
!mkdir -p 'data/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'
python
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex

documents = SimpleDirectoryReader(input_dir="./data").load_data()
index = VectorStoreIndex.from_documents(documents)
python
query_engine = index.as_query_engine()

Sync

python
query_result = query_engine.query("Who is Paul?")

Async

Dispatcher also works on async methods.

python
query_result = await query_engine.aquery("Who is Paul?")

Streaming

Dispatcher also works on methods that support streaming!

python
chat_engine = index.as_chat_engine()
python
streaming_response = chat_engine.stream_chat("Tell me a joke.")
python
for token in streaming_response.response_gen:
    print(token, end="")

Printing Basic Trace Trees with SimpleSpanHandler

python
span_handler.print_trace_trees()