[PYD-877] Log OpenAI streaming response at the end instead of opening a span and attaching context in a generator that may not finish

[alexmojaki]

In the long term I'd like to be able to replace the log with a span that's ended instantly but has a proper start and end timestamp so that it shows a nice duration line in the UI, but that's lower priority.

[linear]

PYD-877 Bug when using OpenAI streaming - failed to attach context

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

📢 Thoughts on this report? Let us know!

[alexmojaki]

@willbakst in #110 (comment) you said

We've already implemented the majority of it in our library following how you instrument openai

so I looked and I can see the same problem in https://github.com/Mirascope/mirascope/blob/5f6ca6038e0041b84dc8939517327d7ad2bcad8b/mirascope/logfire/logfire.py#L174 so you should probably imitate this PR.

To demonstrate the problem in mirascope terms, this:

from mirascope.logfire import with_logfire
from mirascope.openai import OpenAICall, OpenAICallParams

import logfire

logfire.configure()


@with_logfire
class RecipeRecommender(OpenAICall):
    prompt_template = 'Recommend recipes that use {ingredient} as an ingredient'

    ingredient: str

    call_params = OpenAICallParams(model='gpt-3.5-turbo-0125')


def main():
    stream = RecipeRecommender(ingredient='apples').stream()
    for chunk in stream:
        print(chunk)
        break


main()

gives the following error before this PR:

Failed to detach context
Traceback (most recent call last):
  File "/home/alex/work/platform/.venv/lib/python3.12/site-packages/logfire/_internal/integrations/openai.py", line 89, in __stream__
    yield chunk
GeneratorExit

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/alex/work/platform/.venv/lib/python3.12/site-packages/opentelemetry/context/__init__.py", line 158, in detach
    _RUNTIME_CONTEXT.detach(token)  # type: ignore
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/alex/work/platform/.venv/lib/python3.12/site-packages/opentelemetry/context/contextvars_context.py", line 50, in detach
    self._current_context.reset(token)  # type: ignore
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: <Token var=<ContextVar name='current_context' default={} at 0x7f444b347f10> at 0x7f4425b50480> was created in a different Context

so I'm guessing you'll also get similar errors if you stream from a different LLM provider such as Anthropic.

[willbakst]

@alexmojaki thank you for bringing this to my attention!

[cloudflare-workers-and-pages]

Deploying logfire-docs with    Cloudflare Pages

Latest commit:	737b1ba
Status:	✅  Deploy successful!
Preview URL:	https://d1dcd642.logfire-docs.pages.dev
Branch Preview URL:	https://alex-openai-generators.logfire-docs.pages.dev

View logs