Trace LLM, Retriever and Tool Spans
Last updated
Was this helpful?
Last updated
Was this helpful?
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:
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 youfrom 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
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 youUse 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 youIf 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"))
passfrom 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 youThe 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.
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)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 youThe 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.
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)Before diving into the code, ensure you have correctly:
Before diving into the code, ensure you have correctly:
