feat(weave): Implement VertexAI integration

[soumik12345]

Description

This PR adds the autopatch integration with VertexAI Generative Models API.

Supported cases

Sync generation

Without Streaming

import vertexai
import weave
from vertexai.generative_models import GenerativeModel

weave.init(project_name="google_ai_studio-test")
vertexai.init(project="wandb-growth", location="us-central1")
model = GenerativeModel("gemini-1.5-flash-002")
response = model.generate_content(
    "What's a good name for a flower shop specialising in selling dried flower bouquets?"
)

Sample trace

With Streaming

import vertexai
import weave
from vertexai.generative_models import GenerativeModel

weave.init(project_name="google_ai_studio-test")
vertexai.init(project="wandb-growth", location="us-central1")
model = GenerativeModel("gemini-1.5-flash-002")
response = model.generate_content(
    "What's a good name for a flower shop specialising in selling dried flower bouquets?",
    stream=True,
)

Sample trace

Async Generation

Without Streaming

import vertexai
import weave
from vertexai.generative_models import GenerativeModel

weave.init(project_name="google_ai_studio-test")
vertexai.init(project="wandb-growth", location="us-central1")
model = GenerativeModel("gemini-1.5-flash-002")

async def async_generate():
    response = await model.generate_content_async(
        "What's a good name for a flower shop that specializes in selling bouquets of dried flowers?"
    )
    return response

response = asyncio.run(async_generate())

Sample trace

With Streaming

import vertexai
import weave
from vertexai.generative_models import GenerativeModel

weave.init(project_name="google_ai_studio-test")
vertexai.init(project="wandb-growth", location="us-central1")
model = GenerativeModel("gemini-1.5-flash-002")

async def get_response():
    chunks = []
    async for chunk in await model.generate_content_async(
        "What's a good name for a flower shop that specializes in selling bouquets of dried flowers?",
        stream=True,
    ):
        if chunk.text:
            chunks.append(chunk.text)
    return chunks

response = asyncio.run(get_response())

Sample trace

[socket-security]

New dependencies detected. Learn more about Socket for GitHub ↗︎

Package	New capabilities	Transitives	Size	Publisher
pypi/[email protected]	eval, filesystem	0	45.7 kB	vertex_ai

View full report↗︎

[andrewtruong]

can you make a more relevant assert using a mock? This is just checking if anything populated at all.

[soumik12345]

I added cassettes for test_content_generation and test_content_generation_stream. However, I'm unable to generate them for test_content_generation_async and test_content_generation_async_stream.

[andrewtruong]

Do cassettes work here? I thought this used the google grpc stuff under the hood?

In the case where they don't work, you (or gpt!) will need to mock out the request/response manually

[soumik12345]

@andrewtruong Do you mean that we need to write the cassettes manually for those specific functions? Or is there another way?

[andrewtruong]

It's unfortunate we can't use VCRpy here. Instead of mocking, you can put a 5x retry on the test. It's not great, but given how fast these things move it might be the best option

[andrewtruong]

Your pytest markers should be called flaky -- see example: https://github.com/wandb/weave/blame/master/tests/integrations/google_ai_studio/google_ai_studio_test.py#L50

[andrewtruong]

Approving with the understanding that you have thoroughly checked the integration since these tests are skipped