How to Set Custom Trace ID with a Decorator and Include langfuse_observation_id?

[octadion]

I’m trying to set a custom trace ID using the @observe() decorator and pass langfuse_observation_id into the trace context. Here's a snippet of my code:

@observe()
async def _agenerate(
    self,
    messages: str | list[Any],
    streaming: bool = True,
    callbacks: list[BaseLLMCallback] | None = None,
    **kwargs: Any,
) -> str:
    extra_body = kwargs.pop('extra_body', self.extra_body)
    if extra_body and 'metadata' in extra_body:
        trace_user_id = extra_body['metadata'].get('trace_user_id')
        trace_id = extra_body['metadata'].get('trace_id')

        langfuse_context.update_current_trace(
            user_id=trace_user_id,
        )
response = await self.async_client.chat.completions.create(  # type: ignore
        model=model,
        messages=messages,  # type: ignore
        stream=streaming,
        # extra_body=extra_body,
        **kwargs,
    )

I want to know how to correctly pass the langfuse_observation_id in this context. What’s the best approach?

Thanks!

[dosubot[bot]]

To set a custom trace ID and pass langfuse_observation_id using the @observe() decorator, you can pass langfuse_observation_id as a keyword argument to your function. Here's how you can modify your code:

@observe()
async def _agenerate(
    self,
    messages: str | list[Any],
    streaming: bool = True,
    callbacks: list[BaseLLMCallback] | None = None,
    langfuse_observation_id: str = None,
    **kwargs: Any,
) -> str:
    extra_body = kwargs.pop('extra_body', self.extra_body)
    if extra_body and 'metadata' in extra_body:
        trace_user_id = extra_body['metadata'].get('trace_user_id')
        trace_id = extra_body['metadata'].get('trace_id')

        langfuse_context.update_current_trace(
            user_id=trace_user_id,
        )

    # Pass langfuse_observation_id to the async client call if needed
    response = await self.async_client.chat.completions.create(
        model=model,
        messages=messages,
        stream=streaming,
        langfuse_observation_id=langfuse_observation_id,
        **kwargs,
    )

By adding langfuse_observation_id as a parameter, you can pass it when calling _agenerate, ensuring it is included in the trace context ^[1][2].

To continue talking to Dosu, mention @dosu.

[octadion]

Unfortunately, I can't use langfuse_observation_id as a parameter in AsyncCompletions.create() because it results in an error:

TypeError: AsyncCompletions.create() got an unexpected keyword argument 'langfuse_observation_id'

Since the create() method doesn't accept langfuse_observation_id @dosu

[dosubot[bot]]

To include langfuse_observation_id in the trace context without passing it directly to AsyncCompletions.create(), you can pass it as a keyword argument to the function decorated with @observe(). Here's how you can adjust your code:

@observe()
async def _agenerate(
    self,
    messages: str | list[Any],
    streaming: bool = True,
    callbacks: list[BaseLLMCallback] | None = None,
    **kwargs: Any,
) -> str:
    langfuse_observation_id = kwargs.pop('langfuse_observation_id', None)
    extra_body = kwargs.pop('extra_body', self.extra_body)
    
    if extra_body and 'metadata' in extra_body:
        trace_user_id = extra_body['metadata'].get('trace_user_id')
        trace_id = extra_body['metadata'].get('trace_id')

        langfuse_context.update_current_trace(
            user_id=trace_user_id,
        )

    # Proceed with the async client call
    response = await self.async_client.chat.completions.create(
        model=model,
        messages=messages,
        stream=streaming,
        **kwargs,
    )

# Call _agenerate with langfuse_observation_id
await _agenerate(
    messages="your message",
    langfuse_observation_id="your-custom-id"
)

This way, langfuse_observation_id is used within the function's context without being passed to AsyncCompletions.create() ^[1][2].

To continue talking to Dosu, mention @dosu.

[marcklingen]

Hi @octadion, what is this AsyncCompletions class that disallows passing additional kwargs?

This is the recommended way to set custom ids on traces created via the decorator: https://langfuse.com/docs/sdk/python/decorators#set-custom-ids

[marcklingen]

Got it, did you try using the native instrumentation? Thereby you can manage the trace context yourself via a decorator but get auto-instrumentation of the actual llm call: https://langfuse.com/docs/integrations/openai/python/get-started#use-traces

[octadion]

@marcklingen Thanks for the suggestion! but what I'm specifically trying to do is set a custom trace ID (langfuse_observation_id) for the trace context. I need the ability to pass this custom ID when calling _agenerate, so it is included in the trace.

[marcklingen]

As you marked this as resolved, how did you end up implementing this?

[octadion]

@marcklingen As you suggested earlier, I am using the Langfuse decorator with the langfuse_observation_id parameter. However, instead of sending the observation_id in completions.create, I'm passing it in the previous function that calls the observed function.

I encountered another issue, though. I need to trace usage cost tokens, and I followed the documentation here: https://langfuse.com/docs/model-usage-and-cost. But when I use as_type="generation", my user ID doesn't get logged:

@observe()
async def _agenerate(
    self,
    messages: str | list[Any],
    streaming: bool = True,
    callbacks: list[BaseLLMCallback] | None = None,
    **kwargs: Any,
) -> str:
    model = self.model
    if not model:
        raise ValueError(_MODEL_REQUIRED_MSG)
    extra_body = kwargs.pop('extra_body', self.extra_body)
    if extra_body and 'metadata' in extra_body:
        trace_user_id = extra_body['metadata'].get('trace_user_id')

        langfuse_context.update_current_trace(
            user_id=trace_user_id
        )
    kwargs_clone = kwargs.copy()
    langfuse_context.update_current_observation(
        input=messages,
        model=model,
        metadata=kwargs_clone
    )
    response = await self.async_client.chat.completions.create(  # type: ignore
        model=model,
        messages=messages,  # type: ignore
        stream=streaming,
        # extra_body=extra_body,
        **kwargs,
    )
    if streaming:
        full_response = ""
        while True:
            try:
                chunk = await response.__anext__()  # type: ignore
                if not chunk or not chunk.choices:
                    continue

                delta = (
                    chunk.choices[0].delta.content
                    if chunk.choices[0].delta and chunk.choices[0].delta.content
                    else ""
                )  # type: ignore

                full_response += delta
                if callbacks:
                    for callback in callbacks:
                        callback.on_llm_new_token(delta)
                if chunk.choices[0].finish_reason == "stop":  # type: ignore
                    break
            except StopIteration:
                break
        return full_response

    return response.choices[0].message.content or ""

i tried passing user_id in update_current_observation too. Any suggestions on how I can resolve this?

[marcklingen]

Thanks for sharing, this is easy to fix:

From docs

Important: Make sure the as_type="generation" decorated function is called inside another @observe()-decorated function for it to have a top-level trace.