Arize-ai / openinference

OpenTelemetry Instrumentation for AI Observability
https://arize-ai.github.io/openinference/
Apache License 2.0
201 stars 34 forks source link

[bug] Ensure that both langchain and llamaIndex dump prompt template variables as kv pairs #191

Closed mikeldking closed 8 months ago

mikeldking commented 8 months ago

In order to re-play a prompt and apply variables consistently, we must dump the variables in a consistent format so that it can be applied to the template:

E.x.

prompt_template: tell me a {type} joke variables: { type: "funny" }

dosubot[bot] commented 8 months ago

🤖

Hey @mikeldking, nice to see you diving deep again! How's everything on your side?

Based on the current implementation of OpenInference, template variables are indeed dumped as key-value pairs. This is done in the on_event_end method of the OpenInferenceTraceCallbackHandler class, specifically when the event type is CBEventType.LLM. The code iterates over the preceding events in the stack and checks if they are of type CBEventType.TEMPLATING. If they are, the code pops the event from the stack and updates the attributes of the LLM span with the template attributes extracted from the payload of the preceding event.

Here is the relevant code snippet:

if event_type is CBEventType.LLM:
    # Need to update the template variables attribute by extracting their values
    # from the template events, which are sibling events preceding the LLM event.
    # Note that the stack should be sorted in chronological order, and we are
    # going backwards in time on the stack.
    while dfs_stack:
        _, preceding_event_data = dfs_stack[-1]
        if (
            preceding_event_data.parent_id != event_data.parent_id
            or preceding_event_data.event_type is not CBEventType.TEMPLATING
        ):
            break
        dfs_stack.pop()  # remove template event from stack
        for payload in preceding_event_data.payloads:
            # Add template attributes to the LLM span to which they belong.
            try:
                attributes.update(_template_attributes(payload))
            except Exception:
                logger.exception(
                    f"Failed to convert payload to semantic attributes. "
                    f"event_type={preceding_event_data.event_type}, payload={payload}",
                )

This code is responsible for updating the template variables attribute by extracting their values from the template events, which are sibling events preceding the LLM event. The stack should be sorted in chronological order, and the code is going backwards in time on the stack. If the preceding event is not a template event or does not have the same parent ID as the current event, the loop breaks. Otherwise, the template event is removed from the stack and its payload is used to update the attributes of the LLM span. If there is an error during this process, an exception is logged.

You can find this code in the python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_callback.py file.

As for the structure of the payload object in the CBEventType.TEMPLATING events, it is a dictionary (Dict[str, Any]) that contains various event-specific data. The template variables are stored in the payload object under the key EventPayload.TEMPLATE_VARS. This is inferred from the _template_attributes function:

def _template_attributes(payload: Dict[str, Any]) -> Iterator[Tuple[str, Any]]:
    """
    Yields template attributes if present
    """
    if template := payload.get(EventPayload.TEMPLATE):
        yield LLM_PROMPT_TEMPLATE, template
    if template_vars := payload.get(EventPayload.TEMPLATE_VARS):
        yield LLM_PROMPT_TEMPLATE_VARIABLES, json.dumps(template_vars)

In this function, template_vars are extracted from the payload object and then converted to a JSON string. This suggests that template_vars is a dictionary or another JSON-serializable object.

However, without the actual definitions of EventPayload.TEMPLATE and EventPayload.TEMPLATE_VARS, or more context about the CBEventType.TEMPLATING events, it's impossible to provide a more precise answer. You can find this code in the python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_callback.py file.

I hope this helps! If you have any more questions or need further clarification, feel free to ask.

Sources

#### About Dosu This response is meant to be useful and save you time. It is **not meant to be a precise solution**, but rather a starting point for your own research. Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant. If you want to continue the conversation, start your reply with @dosu-bot.
RogerHYang commented 8 months ago

for langchain, we have to put these two pieces together manually

Screenshot 2024-02-08 at 10 19 04 AM