Creating chat agent with Vertex AI

Most LLMs do not have up-to-date information, but you can enable them to access the internet using a mechanism called “Tool”. Vertex AI allows you to build chat agents that can access custom tools and perform specific tasks. In this article, we will learn how to build a custom chat agent. https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/overview
In this article, we will discuss how to deploy an agent on Vertex AI, covering the following steps

1. Setting up the environment
2. Creating search tool using Tavily
3. Creating Agent with Gemini and integrating with Tavily search
4. Deploying the agent
5. Using the agent
6. Summary

1. Setting up environment

First, you need to install VertexAI package with Langchain and Reasoning Engine. Ensure that you have all the IAM permissions set up correctly. If it’s your personal project, you should already have the necessary permissions.

pip install google-cloud-aiplatform[reasoningengine,langchain]

Next you need to setup Tavily API (https://app.tavily.com/home). Create an account and then generate an API key that will be required for searching via the agent.

2. Creating a search tool with Tavily

Once you have the API key, you can access Tavily search using either API calls or Python SDK. In this example, we will using the API. Before we get into the details, here is a big gotcha!

NOTE:

Docstrings are very important for the creating any new tool. Ensure that all the args defined, return format needs to be provided as an example as the agent will use this example to process the future questions. If your tool acts weirdly, this is where I’d recommend you look first.

def tavily_search_method(search_query: str = "where is sydney?"):
    """
    Retrieves results using Tavily search

    Uses the Tavily API (https://api.tavily.com/search) to obtain
    exchange rate data.

    Args:
        search_query: whatever you want to search

    Returns:
        dict: {'query': 'What is the weather in Sydney?', 'follow_up_questions': None, 'answer': 'The 
        current weather in Sydney is 12.3°C with clear skies. The wind is blowing at 11.2 kph from the SSW    direction. The humidity level is at 88%, and there is no precipitation expected.', 'images': None, 'results': [{'title': 'Sydney, New South Wales, Australia Monthly Weather | AccuWeather', 'url': 'https://www.accuweather.com/en/au/sydney/22889/june-weather/22889', 'content': 'Get the monthly weather forecast for Sydney, New South Wales, Australia, including daily high/low, historical averages, to help you plan ahead.', 'score': 0.98987, 'raw_content': None}, {'title': 'Weather in Sydney', 'url': 'https://www.weatherapi.com/', 'content': "{'location': {'name': 'Sydney', 'region': 'New South Wales', 'country': 'Australia', 'lat': -33.88, 'lon': 151.22, 'tz_id': 'Australia/Sydney', 'localtime_epoch': 1719135568, 'localtime': '2024-06-23 19:39'}, 'current': {'last_updated_epoch': 1719135000, 'last_updated': '2024-06-23 19:30', 'temp_c': 12.3, 'temp_f': 54.1, 'is_day': 0, 'condition': {'text': 'Clear', 'icon': '//cdn.weatherapi.com/weather/64x64/night/113.png', 'code': 1000}, 'wind_mph': 6.9, 'wind_kph': 11.2, 'wind_degree': 210, 'wind_dir': 'SSW', 'pressure_mb': 1021.0, 'pressure_in': 30.15, 'precip_mm': 0.0, 'precip_in': 0.0, 'humidity': 88, 'cloud': 0, 'feelslike_c': 10.8, 'feelslike_f': 51.5, 'windchill_c': 12.3, 'windchill_f': 54.2, 'heatindex_c': 13.5, 'heatindex_f': 56.4, 'dewpoint_c': 9.2, 'dewpoint_f': 48.5, 'vis_km': 10.0, 'vis_miles': 6.0, 'uv': 1.0, 'gust_mph': 14.7, 'gust_kph': 23.7}}", 'score': 0.98315, 'raw_content': None}, {'title': 'Weather in Sydney in June 2024', 'url': 'https://world-weather.info/forecast/australia/sydney/june-2024/', 'content': 'Detailed ⚡ Sydney Weather Forecast for June 2024 - day/night 🌡️ temperatures, precipitations - World-Weather.info. Add the current city. Search. Weather; Archive; Widgets °F. World; Australia; Weather in Sydney; Weather in Sydney in June 2024. ... 23 +57° +52° 24 +61° +48° 25 ...', 'score': 0.98173, 'raw_content': None}, {'title': 'Sydney weather in June 2024 | Sydney 14 day weather', 'url': 'https://www.weather25.com/oceania/australia/new-south-wales/sydney?page=month&month=June', 'content': "Our weather forecast can give you a great sense of what weather to expect in Sydney in June 2024. If you're planning to visit Sydney in the near future, we highly recommend that you review the 14 day weather forecast for Sydney before you arrive. Temperatures. 17 ° / 10 °. Rainy Days.", 'score': 0.97698, 'raw_content': None}, {'title': 'Sydney, NSW Weather Forecast June 2024: Daily Highs/Lows & Rain Trends', 'url': 'https://www.weathertab.com/en/c/e/06/commonwealth-of-australia/state-of-new-south-wales/sydney/', 'content': 'Explore comprehensive June 2024 weather forecasts for Sydney, including daily high and low temperatures, precipitation risks, and monthly temperature trends. Featuring detailed day-by-day forecasts, dynamic graphs of daily rain probabilities, and temperature trends to help you plan ahead. ... 23 61°F 46°F 16°C 8°C 09% 24 61°F 47°F 16°C 8 ...', 'score': 0.97489, 'raw_content': None}], 'response_time': 2.53}
    """
    # Initialise Tavily client
    api_key="" # The key you created in the setup stage

    url = "https://api.tavily.com/search"
    headers = {
        "Content-Type": "application/json"
    }
    payload = {
        "api_key": api_key,
        "query": search_query,
        "include_answer": True
    }

    response = requests.post(url, json=payload, headers=headers)

    if response.status_code == 200:
        return response.json()
    else:
        return f"An error occurred: {response.status_code}, {response.text}"

3. Creating the agent and integrating with Tavily search

This is the easy part. There is a simple function call that VertexAI provides to create the agent. Feel free to play around with the arguments and variables. Here, we can add a list of all tools that the agent should access. In this case, we will integrate the agent with the search method created in the previous step.

model_name = "gemini-1.0-pro"
safety_settings = {
        HarmCategory.HARM_CATEGORY_UNSPECIFIED:HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT:HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
        HarmCategory.HARM_CATEGORY_HATE_SPEECH:HarmBlockThreshold.BLOCK_ONLY_HIGH,
        HarmCategory.HARM_CATEGORY_HARASSMENT:HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT:HarmBlockThreshold.BLOCK_NONE,
    }
model_kwargs = {
        "temperature": 0.28,
        "max_output_tokens": 1000,
        "top_p": 0.95,
        "top_k": 40,
        "safety_settings": safety_settings,
    }
agent = reasoning_engines.LangchainAgent(
          model=model_name, 
          tools=[tavily_search_method],
          model_kwargs=model_kwargs,
       )

By default, the agent will try to chat with you without the use of the tool, but when the agent does not have given information, it will use the provided tools. To test the agent, we can ask it a few questions

response = agent.query(input="What is the exchange rate from US dollars to Swedish currency?")
print(response)

P.S. This question is picked up from the linked documentation by yet our search can provide the correct results. You could also ask a few more questions, like:

response = agent.query(input="What is the weather in Sydney?")
print(response)

4. Deploying the agent

Deploying the agent is straightforward. Just provide a reference to the agent we created, specify the requirements, and name the agent.

remote_app = reasoning_engines.ReasoningEngine.create(
        agent,
        requirements=[
            "google-cloud-aiplatform[reasoningengine,langchain]",
        ],
        display_name=agent_name,
    )
print(remote_app) # Prints the information about the deployed agent

5. Using the agent

Once the agent is deployed, you can find the full resource path in the log, or by using the reasoning engine API list method. The deployed agent is of the format projects/PROJECT_ID/locations/LOCATION/reasoningEngines/RESOURCE_ID

print(reasoning_engines.ReasoningEngine.list())

remote_app = reasoning_engines.ReasoningEngine(<Full Engine Path>)

response = remote_app.query(input="What is the weather in Sydney on 24th June? Should I go to the office?")
print(response)

{'input': 'What is the weather in Sydney on 24th June? Should I go to the office?', 'output': 'The weather in Sydney on 24th June 2024 is expected to be pleasant, with a high of around 70°F and a low of about 52°F. The forecast indicates clear conditions with no precipitation expected. \n\nTherefore, based on the weather, going to the office on 24th June seems like a good idea. \n'}

If your LLM did not have access to the internet, it couldn’t fetch the weather from the coming days (This was originally written on 23 June 2024). The agent will then use this information to answer your full question i.e. should I go to the office or not.

Now, you can easily setup the an app, or website for your chat UI!

6. Summary

In this blog, we demonstrated how to create a custom chat agent using Vertex AI and Langchain. By setting up the environment, creating a Tavily search tool, integrating it with the Gemini model, and deploying the agent, you can build an intelligent chat system that accesses real-time information. This approach enhances the capabilities of LLMs, providing up-to-date responses and improving user interactions. Reach out if you need help with setting up an LLM backed Agent.

For the code sample, checkout this github repository https://github.com/yashmehta10/vertex_ai_agent