VERTEX_AI_SEARCH_WIDGET.md

Component documentation for the VertexAISearchWidget search functionality.

Overview

The VertexAISearchWidget component provides integration with Google Vertex AI Search (formerly Enterprise Search). It handles authentication, token management, and provides a search interface for knowledge bases with automatic refresh and error handling.

Component Location

  • File: src/components/VertexAISearchWidget.tsx
  • Type: Search Integration Component
  • Framework: React with TypeScript

Features

  • Google Vertex AI Integration: Direct integration with Google’s enterprise search
  • Authentication Management: Automatic token handling and refresh
  • Search Interface: Clean, accessible search trigger button
  • Error Handling: Graceful degradation when search is unavailable
  • Token Refresh: Automatic token renewal before expiration
  • Configuration Management: Dynamic search configuration support
  • Loading States: Visual feedback during authentication and setup

Props Interface

interface VertexSearchWidgetProps {
  configId: string;                    // Vertex AI Search configuration ID
  buttonText?: string | null;          // Button label text
  className?: string;                  // Additional CSS classes
  iconOnly?: boolean;                  // Icon-only display mode
}

Usage Examples

Basic Implementation

import VertexSearchWidget from '@/components/VertexAISearchWidget';

function ChatInterface() {
  return (
    <div className="chat-tools">
      <VertexSearchWidget
        configId="projects/my-project/locations/global/collections/default_collection/engines/my-search-engine"
        buttonText="Search Knowledge Base"
        iconOnly={false}
      />
    </div>
  );
}

Icon-Only Mode

<VertexSearchWidget
  configId={searchConfig.id}
  iconOnly={true}
  className="toolbar-icon"
/>

With Custom Styling

<VertexSearchWidget
  configId={assistantConfig.searchEngineId}
  buttonText="Search Docs"
  className="search-button border-2 border-blue-500"
/>

Authentication Flow

Token Fetch Process

const fetchSearchToken = async (user: User) => {
  try {
    // Get Firebase ID token
    const idToken = await user.getIdToken();
    
    // Call API endpoint for Google access token
    const response = await fetch('/api/aisearch', {
      method: 'GET',
      headers: {
        'Authorization': `Bearer ${idToken}`
      }
    });
    
    if (!response.ok) {
      console.error('Vertex AI Search authentication error');
      setError('search-disabled');
      return;
    }
    
    const data = await response.json();
    
    // Store and apply token
    if (data.accessToken) {
      latestToken.current = data.accessToken;
      if (widgetRef.current) {
        widgetRef.current.authToken = data.accessToken;
      }
      
      // Set up auto-refresh
      scheduleTokenRefresh(data.expiresIn, user);
    }
  } catch (error) {
    console.error('Token fetch error:', error);
    setError('search-disabled');
  }
};

Token Refresh Scheduling

const scheduleTokenRefresh = (expiresIn: number, user: User) => {
  if (tokenRefreshTimer) {
    clearTimeout(tokenRefreshTimer);
  }
  
  // Refresh 5 minutes before expiration
  const refreshTime = (expiresIn || 3600) * 1000 - 5 * 60 * 1000;
  
  const refreshTimer = setTimeout(() => {
    tokenRequestInProgress.current = false;
    fetchSearchToken(user);
  }, refreshTime);
  
  setTokenRefreshTimer(refreshTimer);
};

Widget Integration

Custom Element Declaration

declare global {
  namespace JSX {
    interface IntrinsicElements {
      'gen-search-widget': React.DetailedHTMLProps<React.HTMLAttributes<HTMLElement> & {
        configId?: string;
        location?: string;
        triggerId?: string;
        authToken?: string;
        ref?: React.Ref<HTMLElement>;
      }, HTMLElement>;
    }
  }
}

Widget Setup

<gen-search-widget
  ref={(el: VertexSearchWidgetElement | null) => {
    if (el !== widgetRef.current) {
      widgetRef.current = el;
      
      // Apply token if available
      if (el && latestToken.current) {
        // Remove configId, set token, then re-add configId
        const currentConfigId = el.getAttribute('configId');
        el.removeAttribute('configId');
        el.authToken = latestToken.current;
        
        setTimeout(() => {
          if (el && currentConfigId) {
            el.setAttribute('configId', currentConfigId);
          }
        }, 100);
      }
    }
  }}
  configId={configId}
  location="eu"
  triggerId="searchWidgetTrigger"
/>

Configuration Management

Dynamic Config Updates

useEffect(() => {
  // Handle configId changes
  if (configId !== previousConfigId.current) {
    previousConfigId.current = configId;
    
    // Reset widget by recreating it
    if (widgetRef.current) {
      const widgetElement = widgetRef.current;
      widgetElement.removeAttribute('configId');
      
      // Apply token if available
      if (latestToken.current) {
        widgetElement.authToken = latestToken.current;
      }
      
      // Add new configId after delay
      setTimeout(() => {
        if (widgetElement) {
          widgetElement.setAttribute('configId', configId);
        }
      }, 100);
    }
  }
}, [configId]);

Component States

Unauthenticated State

if (!isAuthenticated) {
  return (
    <div className="flex flex-col items-center justify-center p-2 border rounded-lg">
      <p className="mb-2 text-center text-sm">Sign in to access search</p>
      <Button 
        variant="outline"
        size="sm"
        onClick={() => FirebaseService.loginWithGoogle()}
      >
        Sign in with Google
      </Button>
    </div>
  );
}

Loading State

if (isLoading && !latestToken.current) {
  return (
    <div className="flex items-center justify-center p-2">
      <div className="animate-spin rounded-full h-4 w-4 border-b-2 border-gray-900" />
      <span className="ml-2 text-sm">Loading search...</span>
    </div>
  );
}

Error State

if (error === 'search-disabled') {
  return (
    <Button
      variant="ghost"
      size="icon"
      disabled={true}
      className="opacity-50"
      title="Search unavailable - authentication error"
    >
      <Search className="h-4 w-4" />
    </Button>
  );
}

Script Loading

Widget Script Integration

useEffect(() => {
  if (!scriptLoaded.current) {
    const script = document.createElement('script');
    script.src = 'https://cloud.google.com/ai/gen-app-builder/client?hl=en_US';
    script.async = true;
    script.onload = () => {
      scriptLoaded.current = true;
    };
    script.onerror = (error) => {
      console.error('Error loading widget script:', error);
    };
    document.body.appendChild(script);
    
    return () => {
      if (document.body.contains(script)) {
        document.body.removeChild(script);
      }
    };
  }
}, []);

Firebase Authentication Integration

Auth State Monitoring

useEffect(() => {
  const unsubscribe = FirebaseService.onAuthStateChange((user) => {
    setIsAuthenticated(!!user);
    setCurrentUser(user);
    
    // Reset setup flag on auth change
    hasSetupWidget.current = false;
    
    if (user) {
      // Schedule immediate token fetch
      setTimeout(() => {
        if (!tokenRequestInProgress.current) {
          fetchSearchToken(user);
        }
      }, 100);
    }
  });
  
  return () => unsubscribe();
}, []);

Error Recovery

Authentication Errors

if (!response.ok) {
  const errorText = await response.text();
  console.error(`Vertex AI Search authentication error: ${response.status} ${errorText}`);
  console.error('To fix locally, run: gcloud auth application-default login');
  setError('search-disabled');
  return; // Don't retry to avoid spamming logs
}

Token Request Safety

// Prevent multiple concurrent requests
if (tokenRequestInProgress.current) {
  debugLog('Token request already in progress, skipping');
  return;
}

tokenRequestInProgress.current = true;

Debugging Support

Debug Logging

const debugLog = (message: string) => {
  console.log(`[VertexSearch] ${message}`);
};

// Example usage:
debugLog(`ConfigId changed from ${previousConfigId.current} to ${configId}`);
debugLog('Token stored in latestToken ref');
debugLog('Widget ref exists, applying token');

Performance Optimizations

  • Token Caching: Stores tokens to avoid unnecessary API calls
  • Request Deduplication: Prevents multiple concurrent token requests
  • Lazy Loading: Only loads widget script when needed
  • Efficient Updates: Minimal DOM manipulation for config changes

Security Considerations

  • Token Expiration: Automatic refresh before expiration
  • Secure Storage: Tokens stored in component refs, not global state
  • API Validation: Backend validates Firebase ID tokens
  • Error Isolation: Authentication errors don’t crash the component

Dependencies

Core Libraries

  • React hooks for state management
  • Firebase Authentication for user management

UI Components

  • @/components/ui/button for interface elements
  • Lucide icons for search icon

External Services

  • Google Vertex AI Search widget script
  • Firebase Authentication service
  • Backend API for token exchange

Local Development

For local development, ensure Google Cloud credentials are configured:

gcloud auth application-default login
  • Chat interfaces that include search functionality
  • Knowledge base management systems
  • Authentication components
  • Google Cloud service integrations