feat(weave): Add integration for Google Genai SDK

[soumik12345]

Continuing from #2078

This PR adds the autopatch integration with Google GenerativeAI SDK.

Tasks:

• Non async, non streaming
• Non async, streaming
• Async, non streaming
• Async, streaming
• Add tests
• Add docs

Examples

Non async, non streaming

import os
import google.generativeai as genai
import weave

genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
weave.init(project_name="google_ai_studio-test")
model = genai.GenerativeModel("gemini-1.5-flash")
response = model.generate_content("Write a story about an AI and magic")

Non async, streaming

import os
import google.generativeai as genai
import weave

genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
weave.init(project_name="google_ai_studio-test")
model = genai.GenerativeModel("gemini-1.5-flash")
response = model.generate_content("Write a story about an AI and magic", stream=True)
chunks = [chunk.text for chunk in response]

Async, non streaming

import os
import asyncio
import google.generativeai as genai
import weave

genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
weave.init(project_name="google_ai_studio-test")
model = genai.GenerativeModel("gemini-1.5-flash")

async def async_generate():
    response = await model.generate_content_async(
        "Write a story about an AI and magic"
    )
    return response

response = asyncio.run(async_generate())

Async Streaming

import os
import asyncio
import google.generativeai as genai
import weave

genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
weave.init(project_name="google_ai_studio-test")
model = genai.GenerativeModel("gemini-1.5-flash")

async def get_response():
    async for chunk in await model.generate_content_async(
        "Write a story about an AI and magic", stream=True
    ):
        if chunk.text:
            print(chunk.text)
        print("_" * 80)


response = asyncio.run(get_response())

[circle-job-mirror]

Preview this PR with FeatureBee: https://beta.wandb.ai/?betaVersion=6800472bd896aa5680ab08c39c6459561f56b551

[andrewtruong]

Is this if correct? What if the acc is done and it's the same as the value (e.g. you return just 1 token)

[soumik12345]

If i don't add this if, the first token is being accumulated twice, both in the weave trace as well as in the output of the function. I agree that this is not the best way to tackle this issue and I'm open to suggestions to improve the logic here.

[andrewtruong]

Are you saying that adding the accumulator modifies the output of the function? It shouldn't do that -- we should never modify the user's output

[soumik12345]

Yeah, I placed the check to prevent modification of output.

[andrewtruong]

sorry, I don't understand.

Here's what I think I'm reading:

1. If the acc is fresh, use the first value of the iterator
2. While the acc is not finished, use the next values
3. If acc != value, then aggregate the candidate completions

So open questions are:

1. How does the above modify the user's output? It shouldn't. We should only be "listening" to the stream
2. Why do we need the acc != value check? I guess there's an interaction with _done, but I'm not sure how

[soumik12345]

Done.

[andrewtruong]

Why are you putting a type ignore here?

Is this safe? I think it assumes that whatever is yielded by the contextmanager is iter-able

[soumik12345]

If I don't put type ignore here, its failing the mypy checks.

[andrewtruong]

The point of mypy is to check if it's type-safe, and the ignore here suggests that it's not. Can you please remove the ignore and fix?

[tcapelle]

We added this back then when the class was extended to support context manager for Anthropic, I pushed a hacky fix