🐛 Bug Report: Gemini (google generativeai) instrumentation sets the input message attributes incorrectly

[dinmukhamedm]

Which component is this bug for?

All Packages

📜 Description

The attribute for gen_ai.prompt.N.* is not set correctly. Yes, I know we will soon/eventually move to sending these as events, but for now, this is an error.

👟 Reproduction steps

import google.generativeai as genai
import os
from traceloop.sdk import Traceloop

Traceloop.init()

genai.configure(api_key=os.environ['GEMINI_API_KEY'])
model = genai.GenerativeModel('gemini-1.5-flash')

def gemini_response(user_message: str):
    response = model.generate_content(user_message)
    print(response.text)

gemini_response('The opposite of hot is')

👍 Expected behavior

The following prompt and completion attributes must be set (by analogy with openai and anthropic):

gen_ai.prompt.0.content: "The opposite of hot is\n"
gen_ai.prompt.0.role: "user"
gen_ai.completion.0.content: "cold\n"
gen_ai.completion.0.role: "assistant"

👎 Actual Behavior with Screenshots

Instead, these attributes are set:

gen_ai.completion.0.content: "cold\n"
gen_ai.prompt.0.user: "The opposite of hot is\n"

🤖 Python Version

3.12

📃 Provide any additional context for the Bug.

No response

👀 Have you spent some time to check if this bug has been raised before?

• I checked and didn't find similar issue

Are you willing to submit PR?

Yes I am willing to submit a PR!

[dosubot]

To address the issue with the incorrect setting of input message attributes for gen_ai.prompt.N.* in the Gemini instrumentation, you can modify the _set_input_attributes function to ensure that the attributes are set correctly based on the request type. Here's a suggested approach:

1. Modify the _set_input_attributes function to correctly set the gen_ai.prompt.N.* attributes. Ensure that the prompt and role are set according to the expected behavior.

Here's an example of how you can adjust the code:

def _set_input_attributes(span, llm_request_type, kwargs):
    if llm_request_type == LLMRequestTypeValues.COMPLETION:
        set_span_attribute(
            span, f"{SpanAttributes.LLM_PROMPTS}.0.content", kwargs.get("prompt")
        )
        set_span_attribute(
            span, f"{SpanAttributes.LLM_PROMPTS}.0.role", "user"
        )
    elif llm_request_type == LLMRequestTypeValues.CHAT:
        for i, message in enumerate(kwargs.get("messages")):
            set_span_attribute(
                span, f"{SpanAttributes.LLM_PROMPTS}.{i}.content", message.get("content")
            )
            set_span_attribute(
                span, f"{SpanAttributes.LLM_PROMPTS}.{i}.role", message.get("role")
            )

2. Ensure the completion attributes are set correctly in the _set_span_completions function, similar to how the prompt attributes are set.

This approach will align the attribute setting with the expected behavior, ensuring that both the content and role are correctly recorded for prompts and completions ^[1][2].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Other}

[nirga]

You're right! Appreciate the willingness to fix this :)