AI Engineer

4 minute read

Migrating LangChain Applications from ChatVertexAI to Google’s Vertex AI: A Comprehensive Guide

Introduction

If you’re using LangChain with Google’s Vertex AI, you might have noticed that the ChatVertexAI class from the langchain_community.chat_models.vertexai module has been marked as deprecated. This deprecation means it’s time to upgrade your code to use the new implementation from the langchain_google_vertexai package. In this guide, we’ll walk through why this change happened, how to migrate your code, and what benefits you’ll gain from the new implementation.

Understanding the Deprecation

Starting from version 0.0.12 of LangChain, the ChatVertexAI class in the community package has been marked as deprecated with the following notice:

# Deprecated since version 0.0.12: Use :class:`~langchain_google_vertexai.ChatVertexAI` instead. 
# It will not be removed until langchain-community==1.0.

This means that while your existing code will continue to work for now, it’s recommended to update to the new implementation to ensure future compatibility and to take advantage of any new features and improvements.

Why Migrate?

There are several compelling reasons to migrate to the new implementation:

  1. Future-proofing: The old implementation will eventually be removed in version 1.0 of langchain-community.
  2. Improved maintenance: The new implementation is maintained directly by the LangChain team in collaboration with Google.
  3. New features: The new implementation may include features and optimizations not available in the deprecated version.
  4. Better integration: The dedicated package provides better integration with other Google Vertex AI services.

Migration Steps

Step 1: Install the Required Package

First, you need to install the langchain-google-vertexai package:

pip install langchain-google-vertexai

Step 2: Update Your Import Statements

Change your import statements from:

from langchain_community.chat_models.vertexai import ChatVertexAI

To:

from langchain_google_vertexai import ChatVertexAI

Step 3: Update Your Code

In most cases, the API remains similar, so minimal changes are required. Here’s a simple example of before and after migration:

Before:

from langchain_community.chat_models.vertexai import ChatVertexAI
from langchain.schema import HumanMessage

chat = ChatVertexAI(
    model_name="gemini-pro",
    project="your-project-id",
    location="us-central1",
    temperature=0.7
)

response = chat.invoke([HumanMessage(content="Tell me a joke about programming.")])
print(response.content)

After:

from langchain_google_vertexai import ChatVertexAI
from langchain.schema import HumanMessage

chat = ChatVertexAI(
    model_name="gemini-pro",
    project="your-project-id",
    location="us-central1",
    temperature=0.7
)

response = chat.invoke([HumanMessage(content="Tell me a joke about programming.")])
print(response.content)

As you can see, the only change required is in the import statement. The rest of the code remains the same.

Handling Advanced Configurations

If you were using advanced configurations with the deprecated ChatVertexAI, you might need to review the documentation for the new implementation to ensure all parameters are correctly mapped.

Example with Advanced Configuration

Here’s an example that shows how to migrate code with more advanced configurations:

Before:

from langchain_community.chat_models.vertexai import ChatVertexAI
from langchain.callbacks import StreamingStdOutCallbackHandler
from langchain.schema import HumanMessage, SystemMessage

chat = ChatVertexAI(
    model_name="gemini-pro",
    project="your-project-id",
    location="us-central1",
    temperature=0.7,
    max_output_tokens=1024,
    top_p=0.95,
    top_k=40,
    streaming=True,
    callbacks=[StreamingStdOutCallbackHandler()]
)

messages = [
    SystemMessage(content="You are a helpful AI assistant."),
    HumanMessage(content="Explain quantum computing in simple terms.")
]

response = chat.invoke(messages)

After:

from langchain_google_vertexai import ChatVertexAI
from langchain.callbacks import StreamingStdOutCallbackHandler
from langchain.schema import HumanMessage, SystemMessage

chat = ChatVertexAI(
    model_name="gemini-pro",
    project="your-project-id",
    location="us-central1",
    temperature=0.7,
    max_output_tokens=1024,
    top_p=0.95,
    top_k=40,
    streaming=True,
    callbacks=[StreamingStdOutCallbackHandler()]
)

messages = [
    SystemMessage(content="You are a helpful AI assistant."),
    HumanMessage(content="Explain quantum computing in simple terms.")
]

response = chat.invoke(messages)

Again, the only change is in the import statement, as the new implementation maintains compatibility with the API of the deprecated version.

Working with Tools and Function Calling

If you were using tools or function calling with the deprecated ChatVertexAI, the migration process is still straightforward:

from langchain_google_vertexai import ChatVertexAI
from langchain.schema import HumanMessage

# Define tools
tools = [
    {
        "name": "get_weather",
        "description": "Get the current weather in a given location",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The city and state, e.g. San Francisco, CA"
                }
            },
            "required": ["location"]
        }
    }
]

# Create model with tools
chat = ChatVertexAI(
    model_name="gemini-pro",
    project="your-project-id",
    location="us-central1"
)

# Bind tools to the model
chat_with_tools = chat.bind_tools(tools)

# Use the model with tools
response = chat_with_tools.invoke([
    HumanMessage(content="What's the weather like in Boston?")
])

print(response)

Handling Streaming Responses

If you’re using streaming responses, the migration is also straightforward:

from langchain_google_vertexai import ChatVertexAI
from langchain.schema import HumanMessage

chat = ChatVertexAI(
    model_name="gemini-pro",
    project="your-project-id",
    location="us-central1",
    streaming=True
)

messages = [HumanMessage(content="Write a short poem about artificial intelligence.")]

# Stream the response
for chunk in chat.stream(messages):
    print(chunk.content, end="", flush=True)

Testing Your Migration

After migrating your code, it’s essential to test thoroughly to ensure everything works as expected. Pay special attention to:

  1. Response formats: Ensure the responses from the new implementation match what you expect
  2. Error handling: Verify that errors are handled correctly
  3. Performance: Check if there are any performance differences

Here’s a simple test script to verify the migration:

from langchain_google_vertexai import ChatVertexAI
from langchain.schema import HumanMessage

def test_chat_vertexai():
    chat = ChatVertexAI(
        model_name="gemini-pro",
        project="your-project-id",
        location="us-central1"
    )
    
    # Test basic functionality
    response = chat.invoke([HumanMessage(content="Hello, how are you?")])
    print(f"Basic response: {response.content}")
    
    # Test streaming
    chat_streaming = ChatVertexAI(
        model_name="gemini-pro",
        project="your-project-id",
        location="us-central1",
        streaming=True
    )
    
    print("Streaming response: ", end="")
    for chunk in chat_streaming.stream([HumanMessage(content="Count to 5.")]):
        print(chunk.content, end="", flush=True)
    print()

if __name__ == "__main__":
    test_chat_vertexai()

Conclusion

Migrating from the deprecated ChatVertexAI to the new implementation in langchain_google_vertexai is a straightforward process that involves updating your import statements and reviewing your configuration parameters. The new implementation maintains compatibility with the API of the deprecated version, making the migration process smooth.

By migrating to the new implementation, you ensure that your code will continue to work with future versions of LangChain and take advantage of any new features and improvements in the integration with Google’s Vertex AI.

Remember that the deprecated implementation will not be removed until langchain-community reaches version 1.0, so you have time to plan and execute your migration. However, it’s recommended to update your code sooner rather than later to avoid potential issues in the future.

Resources

Happy coding with LangChain and Google’s Vertex AI!

This post was originally written in my native language and then translated using an LLM. I apologize if there are any grammatical inconsistencies.

Categories:

Updated:

You may also enjoy