Tutorial 9: Creating and using a custom agent tool

Overview

This tutorial shows how to create a custom tool for Enterprise h2oGPTe agents that automatically applies company brand colors to plots and other outputs. Custom tools enable extending agent functionality by integrating private APIs or organization-specific logic.

This screenshot illustrates how the agent’s behavior changes when the custom tool is integrated. The screenshot on the left shows the default behavior, where the agent responds based solely on the prompt and general knowledge. The screenshot on the right demonstrates how the custom tool enhances responses by dynamically applying company branding based on the tool’s data. Notice the prompt includes explicit conditions that instruct the agent to call the tool when brand color information is needed, enabling more tailored and consistent outputs.

Objective

• Learn how to create and configure a custom agent tool that fetches brand colors based on company names.
• Enable the agent to apply official branding automatically in generated content.

Prerequisites

• Completion of Tutorial 1: A quick introduction to Enterprise h2oGPTe.

Step 1: Create a collection

1. In the Enterprise h2oGPTe navigation menu, click Collections.
2. Click + New collection.
3. In the Collection name box, enter:

   Tutorial 9
4. (Optional) In the Description box, enter:

   This AI Agent comes equipped with a custom tool for looking up brand colors. When you ask for plots, charts, and documents about one or more companies it will automatically brand appropriately.
   
5. Under Default chat settings, set Default generation approach to Agent Only.
   
6. Click + Create.

Step 2: Upload custom tool code

1. Open your new collection.
2. Click Add documents.
3. Choose either Upload documents or Upload plain text.
4. Set Ingest mode to Agent Only.
5. Upload the brand_colors.py as either text or document.
   note

   When uploading your custom tool as plain text:

   • Set Ingest mode to Agent Only.
   • Do not include any file extensions when naming the file.

Step 3: Configure the prompt to use the tool

1. Start a new chat session in the collection.
2. Click the gear icon for Settings.
3. Navigate to Prompts, then click Clone to create a new prompt template.
4. Rename the prompt to Branding Tool.
5. Edit the System prompt to instruct the agent on how to use the tool:

   “You are h2oGPTe, an expert AI developed by H2O.ai.
   
   Additional Tools:
   
   Brand Colors
   When an organization is mentioned, use the brand_colors.py tool to fetch primary and secondary hex color codes. Apply these colors in all outputs like plots or documents.
   
   Usage: python brand_colors.py "OrganizationName"
   
   Call the tool for each organization as needed. If not needed, continue normally.”
6. Switch to Collections and click Apply current settings as collection defaults. Then click Apply again to confirm.
7. Once you have applied the settings, enable the Agent by either clicking the Agent icon near the message input or toggling it on in the Configuration tab.

Step 4: Test your custom tool

To test your new custom tool, try the following queries:

• Plot stock prices of major US airlines over the last year.
• Create a word cloud from T-Mobile’s website.
• Plot major telecommunications stocks.
• Build a predictive model with data from AT&T and explain results.

Example tool code (brand_colors.py)

You can upload the following code as a document or as plain text:

import argparse
import json
import re
import requests
from bs4 import BeautifulSoup
import time
import random
import sys

# Database of known brand colors
BRAND_COLORS = {
    "t-mobile": "#E20074",
    "tmobile": "#E20074",
    "apple": "#000000",
    "google": "#4285F4",
    "microsoft": "#00A4EF",
    "amazon": "#FF9900",
    "facebook": "#1877F2",
    "twitter": "#1DA1F2",
    "x": "#000000",
    "netflix": "#E50914",
    "spotify": "#1DB954",
    "coca-cola": "#F40009",
    "pepsi": "#004B93",
    "starbucks": "#00704A",
    "nike": "#000000",
    "adidas": "#000000",
    "samsung": "#1428A0",
    "ibm": "#0530AD",
    "intel": "#0071C5",
    "hp": "#0096D6",
    "dell": "#007DB8",
    "oracle": "#F80000",
    "salesforce": "#00A1E0",
    "adobe": "#FF0000",
    "mcdonalds": "#FFC72C",
    "burger king": "#D62300",
    "walmart": "#0071CE",
    "target": "#CC0000",
    "disney": "#0107A5",
    "verizon": "#CD040B",
    "at&t": "#00A8E0",
    "sprint": "#FFCE00",
    "h2o.ai": "#FFE600",
    "h2o": "#FFE600",
    # Adding Airbnb with correct colors
    "airbnb": "#FF5A5F",
    # Adding more companies
    "uber": "#000000",
    "lyft": "#FF00BF",
    "slack": "#4A154B",
    "dropbox": "#0061FF",
    "linkedin": "#0A66C2",
    "instagram": "#E4405F",
    "pinterest": "#E60023",
    "reddit": "#FF4500",
    "twitch": "#9146FF",
    "youtube": "#FF0000",
    # Adding Zoom with correct colors
    "zoom": "#0B5CFF",
}

def extract_hex_colors(text):
    """Extract hex color codes from text"""
    # Match both 6-digit and 3-digit hex color codes
    hex_pattern = r'#[0-9a-fA-F]{6}|#[0-9a-fA-F]{3}'
    hex_colors = re.findall(hex_pattern, text)

    # Normalize to 6-digit hex and remove duplicates while preserving order
    unique_colors = []
    for color in hex_colors:
        # Convert 3-digit hex to 6-digit
        if len(color) == 4:  # #RGB format
            color = '#' + color[1] + color[1] + color[2] + color[2] + color[3] + color[3]

        if color.lower() not in [c.lower() for c in unique_colors]:
            unique_colors.append(color.lower())

    return unique_colors

def fetch_url_with_retry(url, headers, max_retries=3):
    """Fetch URL with retry logic"""
    for attempt in range(max_retries):
        try:
            response = requests.get(url, headers=headers, timeout=10)
            if response.status_code == 200:
                return response
            time.sleep(1 + random.random())  # Add jitter to avoid rate limiting
        except Exception as e:
            print(f"Attempt {attempt+1} failed: {e}")
            if attempt < max_retries - 1:
                time.sleep(2 + random.random() * 2)  # Exponential backoff with jitter
    return None

def search_brand_color_websites(company_name):
    """Search specific brand color websites for the company"""
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
    }

    # List of brand color websites to check
    brand_sites = [
        f"https://www.brandcolorcode.com/{company_name.lower().replace(' ', '-')}",
        f"https://www.schemecolor.com/{company_name.lower().replace(' ', '-')}-logo-colors.php",
        f"https://www.designpieces.com/palette/{company_name.lower().replace(' ', '-')}-color-palette-hex-and-rgb/",
        f"https://encycolorpedia.com/companies/{company_name.lower().replace(' ', '-')}"
    ]

    for url in brand_sites:
        response = fetch_url_with_retry(url, headers)
        if response:
            soup = BeautifulSoup(response.text, 'html.parser')
            text_content = soup.get_text()
            hex_colors = extract_hex_colors(text_content)

            if len(hex_colors) >= 2:
                return hex_colors[0], hex_colors[1]
            elif len(hex_colors) == 1:
                return hex_colors[0], None

    return None, None

def get_brand_colors(company_name, default_primary="#1f77b4"):
    """
    Get brand colors for a given company.
    Returns primary and secondary colors in hex format.
    """
    # Normalize company name for lookup
    normalized_name = company_name.lower().strip()

    if normalized_name in BRAND_COLORS:
        print(f"Found {company_name} in brand database!")
        return BRAND_COLORS[normalized_name],

    # Try common variations (e.g., "t-mobile" vs "tmobile")
    normalized_no_dash = normalized_name.replace('-', '')
    if normalized_no_dash in BRAND_COLORS:
        print(f"Found {company_name} (without dashes) in brand database!")
        return BRAND_COLORS[normalized_no_dash]

    normalized_with_dash = re.sub(r'([a-z])([a-z])', r'\1-\2', normalized_name)
    if normalized_with_dash in BRAND_COLORS:
        print(f"Found {company_name} (with dashes) in brand database!")
        return BRAND_COLORS[normalized_with_dash]

    # If not in database, try searching brand color websites
    print(f"Searching brand color websites for {company_name}...")
    primary, secondary = search_brand_color_websites(company_name)

    if primary:
        print(f"Found colors on brand color websites!")
        return primary

    # As a last resort, try a Google search
    try:
        print(f"Trying web search for {company_name} brand colors...")
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }

        # Try a more specific search query
        search_query = f"{company_name} official brand color hex code"
        search_url = f"https://www.google.com/search?q={search_query}"

        response = fetch_url_with_retry(search_url, headers)
        if response:
            soup = BeautifulSoup(response.text, 'html.parser')
            text_content = soup.get_text()
            hex_colors = extract_hex_colors(text_content)

            if len(hex_colors) >= 1:
                print(f"Found colors via web search!")
                return hex_colors[0]

        print(f"Could not find brand colors for {company_name}. Using defaults.")
        return default_primary

    except Exception as e:
        print(f"Error: {e}")
        return default_primary


def main():
    if len(sys.argv) < 2:
        print("Usage: python file_name.py 'company name' ['company name' ...]")
        sys.exit(1)

    results = []
    for company in sys.argv[1:]:  # Iterate over each company name provided
        print(f"Searching for {company} brand colors...")
        primary_color = get_brand_colors(company)
        results.append({
            "company": company,
            "primary_color": primary_color,
        })

    # Output all results as JSON for easy parsing
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()

---

Summary

You have successfully created and configured a custom tool in Enterprise h2oGPTe that allows agents to automatically apply brand colors for mentioned organizations, enhancing output branding and consistency.

---

Feedback

• Submit and view feedback for this page
• Send feedback about Enterprise h2oGPTe to cloud-feedback@h2o.ai