Trace LLM, Retriever and Tool Spans

In cases where teams want to manually instrument their own spans and are not using auto-instrumentation, this section instrumentation attributes of the main span types.

This is a link to the semantic conventions of open inference for reference:

openinference/spec/semantic_conventions.md at main · Arize-ai/openinferenceGitHub

General Attributes

These are attributes that can work on any span.

from openinference.semconv.trace import SpanAttributes

def do_work():
    with tracer.start_as_current_span("span-name") as span:
        span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, "CHAIN") # see here for a list of span kinds: https://github.com/Arize-ai/openinference/blob/main/python/openinference-semantic-conventions/src/openinference/semconv/trace/__init__.py#L271
        span.set_attribute(SpanAttributes.TAG_TAGS, str("['tag1','tag2']")) # List of tags to give the span a category
        span.set_attribute(SpanAttributes.INPUT_VALUE, "<INPUT>") # The input value to an operation
        span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "text/plain") # either text/plain or application/json
        span.set_attribute(SpanAttributes.OUTPUT_VALUE, "<OUTPUT>") # The output value of an operation
        span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "text/plain") # either text/plain or application/json 
        span.set_attribute(SpanAttributes.METADATA, "<ADDITIONAL_METADATA>") # additional key value pairs you want to store
        span.set_attribute(SpanAttributes.IMAGE_URL, "<IMAGE_URL>") # An http or base64 image url
        span.set_attribute("exception.message", "<EXCEPTION_MESSAGE>")
        span.set_attribute("exception.stacktrace", "<EXCEPTION_STACKTRACE>")
        span.set_attribute("exception.type", "<EXCEPTION_TYPE>") # e.g. NullPointerException
        
        
        # do some work that 'span' will track
        print("doing some work...")
        # When the 'with' block goes out of scope, 'span' is closed for you

LLM

from openinference.semconv.trace import SpanAttributes

def llm_call():
    with tracer.start_as_current_span("span-name") as span:
        span.set_attribute(SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES, "<prompt_template_variables>") # JSON of key value pairs representing prompt vars: to be applied to prompt template
        span.set_attribute(SpanAttributes.LLM_PROMPT_TEMPLATE, "<prompt_template>") # Template used to generate prompts as Python f-strings
        span.set_attribute(SpanAttributes.LLM_PROMPT_TEMPLATE_VERSION, "<input_messages>") # The version of the prompt template, "v1.0"
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, "<prompt_tokens>") # The number of tokens in the prompt
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, "<completion_tokens>") # The number of tokens in the completion
        span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, "<tokens_total>") # Total number of tokens, including both prompt and completion.
        span.set_attribute(SpanAttributes.LLM_FUNCTION_CALL, "<function_call_results>") # For models and APIs that support function calling. Records attributes such as the function name and arguments to the called function. This is the result JSON from a model representing the function(s) "to call" 
        span.set_attribute(SpanAttributes.LLM_INVOCATION_PARAMETERS, "<invocation_parameters>") # These are the invocation Object recording details of a function call in models or APIs, "{model_name: 'gpt-3', temperature: 0.7}"
        span.set_attribute(SpanAttributes.LLM_INPUT_MESSAGES, "<input_messages>") # List of messages sent to the LLM in a chat API request, [{"message.role": "user", "message.content": "hello"}]
        span.set_attribute(SpanAttributes.LLM_OUTPUT_MESSAGES, "<output_messages>") # Messages received from a chat API, [{"message.role": "user", "message.content": "hello"}]
        span.set_attribute(SpanAttributes.LLM_MODEL_NAME, "<input_messages") # The name of the language model being utilized
        

EMBEDDING

from openinference.semconv.trace import SpanAttributes, EmbeddingAttributes

def get_embeddings():
    with tracer.start_as_current_span("span-name") as span:
        span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.EMBEDDING.value)
        span.set_attribute(SpanAttributes.EMBEDDING_MODEL_NAME, "<RETURNED_EMBEDDING_VECTOR>") # The name of the embedding model.                
        for i in range(number_of_embeddings):
            span.set_attribute(f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.{i}.{EmbeddingAttributes.EMBEDDING_TEXT}", <TEXT HERE>)
            span.set_attribute(f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.{i}.{EmbeddingAttributes.EMBEDDING_VECTOR}", <EMBEDDING VECTOR HERE>)
        # do some work that 'span' will track
        print("doing some work...")
        # When the 'with' block goes out of scope, 'span' is closed for you

RETRIEVER

Use this span type to log spans for documents retrieved as part of a RAG pipeline. Here is a simple example:

from openinference.semconv.trace import SpanAttributes

def get_embeddings():
    with tracer.start_as_current_span("span-name") as span:
        span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.RETRIEVER.value)
        span.set_attribute(SpanAttributes.DOCUMENT_ID, "<DOCUMENT_ID>") # Unique identifier for a document               
        span.set_attribute(SpanAttributes.DOCUMENT_SCORE, "<DOCUMENT_SCORE>") # Score representing the relevance of a document
        span.set_attribute(SpanAttributes.DOCUMENT_CONTENT, "<DOCUMENT_CONTENT>") # The content of a retrieved document
        span.set_attribute(SpanAttributes.DOCUMENT_METADATA, str(<DOCUMENT_METADATA_JSON>)) # Metadata associated with a document
        # do some work that 'span' will track
        print("doing some work...")
        # When the 'with' block goes out of scope, 'span' is closed for you

If you are making calls to external APIs to retrieve documents, you can do something like the below:

with tracer.start_as_current_span("Vector Search") as vector_search_span:
    vector_search_span.set_attribute(SpanAttributes.INPUT_VALUE, query_text)
    vector_search_span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.RETRIEVER.value)
    try:
        documents = search_client.search(
            search_text=query_text,
        )
        # Example documents returned from the vector search
        # documents = [
        #         {'document.id': 0, 'document.score': 0.1, 'document.content': "This is a test"}, 
        #         {'document.id': 1, 'document.score': 0.2, 'document.content': "This is a test2"} 
        #     ]
        for i, document in enumerate(documents):
            vector_search_span.set_attribute(f"retrieval.documents.{i}.document.id", document["document.id"])
            vector_search_span.set_attribute(f"retrieval.documents.{i}.document.score", document["document.score"])
            vector_search_span.set_attribute(f"retrieval.documents.{i}.document.content", document["document.content"])
                      
        vector_search_span.set_status(Status(StatusCode.OK, "Vector search successful"))
    except Exception as e:
        vector_search_span.set_status(Status(StatusCode.ERROR, "Vector search failed"))
        pass

TOOL

from openinference.semconv.trace import SpanAttributes,ToolCallAttributes

def tool_call():
    with tracer.start_as_current_span("span-name") as span:
        span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.TOOL.value)
        span.set_attribute(ToolCallAttributes.TOOL_CALL_FUNCTION_NAME, "<NAME_OF_YOUR_TOOL>") # The name of the tool being utilized
        span.set_attribute(ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON, str(<JSON_OBJ_OF_FUNCTION_PARAMS>)) # The arguments for the function being invoked by a tool call
        # do some work that 'span' will track
        print("doing some work...")
        # When the 'with' block goes out of scope, 'span' is closed for you

The example below demonstrates how to manually trace a tool function along with a chat completion response. You'll see how to create spans for both the tool and LLM to capture their input, output, and key events.

Before diving into the code, ensure you have configured your tracer correctly:

import json

#from your_tracer import tracer

def run_tool(tool_function, tool_args):
    #first set context for current span 
    with tracer.start_as_current_span(
        name="Tool - some tool",
        attributes={
            # Set these attributes before calling the tool, in case an exception is raised by the tool
            **{
                "openinference.span.kind": "TOOL",
                "input.value": question,
                "message.tool_calls.0.tool_call.function.name": tool_function.__name__,
                "message.tool_calls.0.tool_call.function.arguments": json.dumps(
                    tool_args
                ),
            },
        },
    ) as tool_span:
        #run tool; output is formatted prompt for chat completion
        resulting_prompt = tool_function(input=tool_args)
        # optional - set the resulting prompt as the tool span output
        tool_span.set_attribute(
            "message.tool_calls.0.tool_call.function.output", resulting_prompt
        )

        # This LLM span nests under the tool span in the trace
        with tracer.start_as_current_span(
            name="Tool - llm response",
            # Set these attributes before calling the LLM
            attributes={
                "openinference.span.kind": "LLM",
                "input.value": resulting_prompt,
            },
        ) as llm_span:
            llm_response = openai_client.chat.completions.create(
                    model=model_version,
                    messages=[current_user_message],
                    temperature=TEMPERATURE,
                )
            llm_span.set_attribute("output.value", llm_response)

RERANKER

from openinference.semconv.trace import SpanAttributes
def tool_call():
    with tracer.start_as_current_span("span-name") as span:
        span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, OpenInferenceSpanKindValues.RERANKER.value)
        span.set_attribute(SpanAttributes.RERANKER_INPUT_DOCUMENTS, str(<LIST_OF_DOCUMENTS>)) # List of documents as input to the reranker
        span.set_attribute(SpanAttributes.RERANKER_OUTPUT_DOCUMENTS, str(<LIST_OF_DOCUMENTS>)) # List of documents as outputs of the reranker
        span.set_attribute(SpanAttributes.RERANKER_QUERY, "<RERANKER_QUERY>") # Query parameter of the reranker
        span.set_attribute(SpanAttributes.RERANKER_MODEL_NAME, "<MODEL_NAME>") # Name of the reranker model
        span.set_attribute(SpanAttributes.RERANKER_TOP_K, "<RERANKER_TOP_K>") # Top K parameter of the reranker        
        # do some work that 'span' will track
        print("doing some work...")
        # When the 'with' block goes out of scope, 'span' is closed for you

Instrumenting tool spans

The example below demonstrates how to manually trace a tool function along with a chat completion response. You'll see how to create spans for both the tool and LLM to capture their input, output, and key events.

Before diving into the code, ensure you have configured your tracer correctly:

import json

#from your_tracer import tracer

def run_tool(tool_function, tool_args):
    #first set context for current span 
    with tracer.start_as_current_span(
        name="Tool - some tool",
        attributes={
            # Set these attributes before calling the tool, in case an exception is raised by the tool
            **{
                "openinference.span.kind": "TOOL",
                "input.value": question,
                "message.tool_calls.0.tool_call.function.name": tool_function.__name__,
                "message.tool_calls.0.tool_call.function.arguments": json.dumps(
                    tool_args
                ),
            },
        },
    ) as tool_span:
        #run tool; output is formatted prompt for chat completion
        resulting_prompt = tool_function(input=tool_args)
        # optional - set the resulting prompt as the tool span output
        tool_span.set_attribute(
            "message.tool_calls.0.tool_call.function.output", resulting_prompt
        )

        # This LLM span nests under the tool span in the trace
        with tracer.start_as_current_span(
            name="Tool - llm response",
            # Set these attributes before calling the LLM
            attributes={
                "openinference.span.kind": "LLM",
                "input.value": resulting_prompt,
            },
        ) as llm_span:
            llm_response = openai_client.chat.completions.create(
                    model=model_version,
                    messages=[current_user_message],
                    temperature=TEMPERATURE,
                )
            llm_span.set_attribute("output.value", llm_response)