Ideas on how to implement the Tracing API in the Modelmetry Python SDK

[lazharichir]

This tracks how to implement a clean, intuitive, yet flexible Tracing API to the Modelmetry Python SDK.

By instrumenting their code with our tracing APIs, developers can generate data (traces, spans, events, findings) that can be visualized and analyzed using the Modelmetry frontend. This will empower users to gain insights into various LLM metrics based on their specific use cases and model outputs.

The main benefits include:

• Improved Observability: Users can gain deeper visibility into the behavior of their LLM features.
• Enhanced Debugging: Tracing data can pinpoint issues and bottlenecks within LLM workflows.
• Detailed Analysis: Users can analyze various metrics to understand LLM performance characteristics by using Modelmetry evaluators for traces and spans (e.g., PII detection, accuracy, helpfulness, tone).

Key Concepts

• Trace: A collection of spans that represent the lifecycle of a higher level operation.
• Span: A named, timed granular operation representing a piece of the LLM workflow.
• Finding: A custom data point that can be attached to a span or a trace (numerical, labeled, or boolean value).
• Event: A timestamped data point that can be attached to a span.
• Ingest Endpoint: The backend endpoint where traces, spans, events, and findings, are all sent for storage (and later visualization).

Example Use Case

For example, a company has a customer service chatbot. One of their customers interacts with the chatbot asking for some help. The developers need to a few steps to get this right:

1. Get the user's message and send it to the LLM alongside available tools for the LLM to use (e.g., GetOrderInformation, RequestRefund, Escalate).
2. The LLM needs to use GetOrderInformation which you call and send the response to the LLM so it can write a proper response to the user.
3. The LLM processes the request and sends the response back.
4. You check the response to ensure it's safe.
5. You send the response back to the user.

All of this happened within 3 to 5 seconds. You can trace this operation and see how long each step took, what the LLM did, and what the response was. You can also attach events to each step to see what the LLM was thinking at each step.

The entire operation could have been a single trace named chatbot.turn with spans llm.generate, llm.tools.get_user_message, llm.generate_final, and llm.response_check. Each span could have events attached to them to show what the LLM was thinking at each step. Some findings could be attached to the spans (e.g., input_tokens, output_tokens, response_safe_score).

Technical Design

The Python SDK will be extended with APIs for creating and managing traces, spans, events, and findings.

The backend –and the OpenAPI spec– are already defined. So our implementation needs to elegantly wrap the following entities:

• IngestSignalsV1RequestBody – the data structure that the backend expects.
• The embedded CreateEventParams, CreateFindingParams, CreateSpanParams, and CreateTraceParams classes.

The Tracing API must allow users to:

• Define spans to mark specific sections of their LLM code.
• Attach events with relevant data (e.g., model inputs, outputs, errors) to specific spans.
• Emit custom findings/events associated with LLM operations (i.e., associated with a Span or a Trace).
• The SDK will handle packaging and sending traces to the Modelmetry backend service.

Implementation Details

• The API design for trace creation and management should be intuitive and user-friendly.
• The Tracing SDK functions should allow for two flows:
  • synchronous flushing: in an AWS Lambda or Cloud Run request where no background jobs can be executed once the request has been handled so we need a client.flushAll() call at the end of the request to send all data collected to the backend.
  • asynchronous flushing: in an API server that's running 24/7 and can flush traces in the background after a certain number of traces have been collected or after a certain amount of time has passed, whichever comes first.
• Clear documentation and examples will be provided to guide users on implementing tracing within their LLM code.
• Consider supporting context propagation across asynchronous operations and distributed systems (if applicable).

Pseudocode just to illustrate the API on a high level

# Function to handle incoming chat messages
def handle_chat_message(user_message):
    trace = client.new_trace("chatbot.turn", {
      "user_id": "2948",
      "user_name": "Alice",
      "user_email": "[email protected]",
    })
    
    try:
        # Call the LLM first and offer tools
        initial_prompt = f"The user said: {user_message}\n\nYou can use the following tools:\n- GetOrderInformation(order_id)\nWhat do you need to help the user?"
        
        # In a real scenario, this would be a list of available tools
        tools = []

        span = trace.new_span("llm.generate")
        try:
            response = openai.Completion.create(prompt=initial_prompt)
            span.add_finding("output_tokens", 50)
        except Exception as err:
            span.mark_as_error(err)
            span.new_event("The completion failed: " + str(err))
            raise
        finally:
            span.end()
        
        llm_response = response.choices[0].text.strip()
        tool_chosen = response.choices[0].tool_choice  # Assuming the LLM chooses a tool
        
        # Assuming the LLM asks for an order ID
        if tool_chosen == "GetOrderInformation":
            order_id = "12345"  # In a real scenario, extract this from user_message or further context
            
            # Call the tool
            span = trace.new_span("llm.tools.get_order_information")
            try:
                tool_response = get_order_information(order_id)
                span.add_finding("tool_used", True)
            except Exception as err:
                span.mark_as_error(err)
                span.new_event("The tool call failed: " + str(err))
                raise
            finally:
                span.end()
            
            # Add the tool's response to the context for the next/final LLM call
            context_with_tool_data = (
                f"The user said: {user_message}\n"
                f"The LLM requested order information for order_id: {order_id}\n"
                f"The tool responded with: {tool_response}\n"
                "Now, generate a final response for the user."
            )
            
            # Call the LLM with the tool's output to generate the final response
            span = trace.new_span("llm.generate")
            try:
                final_response = openai.Completion.create(prompt=context_with_tool_data )
                span.add_finding("output_tokens", final_response.usage.completion_tokens)
            except Exception as err:
                span.mark_as_error(err)
                span.new_event("The completion failed: " + str(err))
                raise
            finally:
                span.end()
            
            final_llm_response = final_response.choices[0].text.strip()

            # Check the final response for safety with external API
            span = trace.new_span("check_response_safety")
            try:
                response_safe_score = check_response_safety(final_llm_response)
                span.add_finding("response_safe_score", response_safe_score)
            except Exception as err:
                span.mark_as_error(err)
                span.new_event("The safety check failed: " + str(err))
                raise
            finally:
                span.end()
            
            # Send the final response to the user
            return final_llm_response
    finally:
        trace.end()
        client.flush()

# Mock function representing the actual tool call
def get_order_information(order_id):
    # This is a mock function. In reality, this would query a database or an external service.
    return {
        "order_id": order_id,
        "status": "shipped",
        "estimated_delivery": "2023-07-25"
    }

# Example usage
user_message = "Can you tell me the status of my order?"
response_to_user = handle_chat_message(user_message)
print(response_to_user)

Create a finding for a span

span.new_finding("input_tokens", 238)
span.new_finding("output_tokens", 172, {
  "model": "gpt-4o-mini",
})

Create an event for a span

span.mark_as_error(err)
span.new_event("the completion failed: " + errMessage)

Testing

• Unit tests will be implemented to ensure the functionality of the tracing APIs.
• Integration tests will be conducted to verify that traces are correctly sent to the Modelmetry backend and can be visualized in the frontend.

Future Considerations

• Ensure this Tracing API allows automatic tracing for common LLM operations/libraries (e.g., LangChain, OpenAI SDK) in the future by simply adding adapters to the SDK on top of this API.