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/vertexai@1.70.0	None	0	0 B

View full report↗︎