Vertex AI via API key

[shandheap]

Description

Google launched Vertex AI Studio which allows users to use Vertex AI via api key. Currently Pydantic AI only supports using api keys through the Generative Language API, whereas Vertex AI is only supported via service account: https://ai.pydantic.dev/models/google/

I've written a custom model wrapper that uses the Vertex AI api endpoint with the api key but it would be ideal if Pydantic AI supports this out of the box since the implementation is vulnerable to breaking changes in Pydantic AI and requires additional overhead to maintain.

References

Vertex AI docs for usage with api key: https://cloud.google.com/vertex-ai/generative-ai/docs/start/quickstart?usertype=apikey

[DouweM]

@shandheap Thanks for the report. Can you please share the custom model wrapper you've built?

Does it just require API_KEY and GOOGLE_GENAI_USE_VERTEXAI to be set? I assume the latter is identical to the vertexai=True we already support on GoogleProvider, so I'm curious of just setting that + the API_KEY env var would be sufficient.

[shandheap]

@DouweM Thanks for the quick response!

Does it just require API_KEY and GOOGLE_GENAI_USE_VERTEXAI to be set? I assume the latter is identical to the vertexai=True we already support on GoogleProvider, so I'm curious of just setting that + the API_KEY env var would be sufficient.

Thanks so much for this info, I was able to remove all my custom model implementation by simply switching the client to the genai.Client from Google. Here's my updated code:

# Set environment variable to enable Vertex AI routing
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"

client = genai.Client(
    http_options=HttpOptions(
        api_version="v1",
        timeout=60,
    )
)

# Create provider with custom client
provider = GoogleProvider(client=client)

# GoogleModel instance with the specified model name and provider
model = GoogleModel(self.model, provider=provider)

I was able to drop in model into the Pydantic AI Agent and it works now without needing the complex custom model setup I had before.

[DouweM]

@shandheap Did it not work with GoogleProvider(vertexai=True)? I'm not sure why the custom client is necessary.

[shandheap]

@DouweM When I tried using GoogleProvider(vertexai=True) it didn't work but I realize that this might be due to the fact that we're using an older version of pydantic-ai (1.0.8).

Anyways it seems to be working fine now, thanks for your help! Pydantic AI has been such a fantastic library