fix(weave): Fixes Async integration timing

[tssweeney]

Bug Observation: Many integrations (notably OpenAI) when using async methods would report very long runtimes - specifically inside of an eval loop. Initial triage showed that the start event was getting emitted long before the actual function started executing.

Bug Root Cause: The root cause of this issue is the widespread usage of the following pattern inside the integration. Specifically: wrapping an async function with an outer async function, and binding the op to the outer function:

async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
    # Possible pre-processing
    res = await fn(*args, **kwargs)
    # Possible post-processing
    return res

wrapped_op = weave.op(_async_wrapper)

Explanation: This is particularly bad because the weave tracing engine thinks that the user-defined method is _async_wrapper, however the function if interest is actually fn. Therefore, we start our wall clock at the entry of _async_wrapper, which immediately pauses. This means the time reported includes all the scheduling delay of the event loop!

Fix: This is an anti-pattern in our integrations. We should move the pre- and post- processing logic to different mechanics. There are a few classes of changes:

1. Cases where no additional logic was even needed and this pattern spread via copy-pasta:
   • weave/integrations/anthropic/anthropic_sdk.py
   • weave/integrations/cerebras/cerebras_sdk.py
   • weave/integrations/google_ai_studio/google_ai_studio_sdk.py
   • weave/integrations/instructor/instructor_iterable_utils.py
   • weave/integrations/vertexai/vertexai_sdk.py
2. Cases where post processing was used to change the output to weave. In these cases we should simply use the new postprocess_output method
   • weave/integrations/cohere/cohere_sdk.py
3. Harder cases (OpenAI). In this case we do some really weird logic where we inject a param to get the costs out. This required two changes: 1) do the param injections outside the op; 2) use a context var to determine skipping logic as opposed to inputs. Note: this is definitely not pretty... but it works
   • weave/integrations/openai/openai_sdk.py

Finally, weave/trace/op_extensions/accumulator.py was needed to be enhanced such that on_output can handle Coroutines now (:

Testing: All the integration tests served as a pretty good suite

[circle-job-mirror]

Preview this PR with FeatureBee: https://beta.wandb.ai/?betaVersion=0dcde2d3a65b7edb29845dec8d15eb4532c5380e

[gtarpenning]

how were we using this before?

[gtarpenning]

the comment is now confusing, also single quotes is weird while we are here

[gtarpenning]

I think @andrewtruong should look at the accumulator bit

[andrewtruong]

What does skip_last mean? Maybe a comment here would be helpful

[andrewtruong]

Same here, I'm not sure what skip_last means. Implicitly it has something to do with streaming?

[gtarpenning]

this is because openai appends token information to the end. we need that info, so we always try to grab it. if the user wants it, we dont want to skip it, but if the user hasn't requested it, we need to skip it so the output looks as expected