sam-kd/ai_sam/ref/services/liveManager.ts

import { GoogleGenAI, LiveServerMessage, Modality, FunctionDeclaration, Type } from '@google/genai';
import { decode, decodeAudioData, createPcmBlob } from '../utils/audio';
import { SYSTEM_INSTRUCTION } from '../constants';

export type ToolHandler = (name: string, args: any) => Promise<any>;

export class LiveManager {
  private ai: GoogleGenAI;
  private inputAudioContext: AudioContext | null = null;
  private outputAudioContext: AudioContext | null = null;
  private sessionPromise: Promise<any> | null = null;
  private nextStartTime = 0;
  private sources = new Set<AudioBufferSourceNode>();
  private toolHandler: ToolHandler;
  private onStatusChange: (status: string) => void;
  private onAudioLevel: (level: number) => void; // For visualizer
  public outputAnalyser: AnalyserNode | null = null;

  constructor(apiKey: string, toolHandler: ToolHandler, onStatusChange: (status: string) => void, onAudioLevel: (level: number) => void) {
    this.ai = new GoogleGenAI({ apiKey });
    this.toolHandler = toolHandler;
    this.onStatusChange = onStatusChange;
    this.onAudioLevel = onAudioLevel;
  }

  async connect() {
    this.onStatusChange('connecting');

    // Audio Context Setup
    this.inputAudioContext = new (window.AudioContext || (window as any).webkitAudioContext)({ sampleRate: 16000 });
    this.outputAudioContext = new (window.AudioContext || (window as any).webkitAudioContext)({ sampleRate: 24000 });
    
    // Setup Analyser for Visualizer
    this.outputAnalyser = this.outputAudioContext.createAnalyser();
    this.outputAnalyser.fftSize = 256;
    this.outputAnalyser.connect(this.outputAudioContext.destination);

    // Tools Definition
    const navigateToPage: FunctionDeclaration = {
      name: 'navigateToPage',
      parameters: {
        type: Type.OBJECT,
        description: 'Navigate the user to a specific page in the application menu.',
        properties: {
          pageId: {
            type: Type.STRING,
            description: 'The ID of the page to navigate to (e.g., dashboard, files, analytics, settings, team, calendar).',
          },
        },
        required: ['pageId'],
      },
    };

    const searchDocuments: FunctionDeclaration = {
      name: 'searchDocuments',
      parameters: {
        type: Type.OBJECT,
        description: 'Search for work documents or files based on a query.',
        properties: {
          query: {
            type: Type.STRING,
            description: 'The search keywords to filter files by.',
          },
        },
        required: ['query'],
      },
    };

    try {
      // Connect to Gemini
      this.sessionPromise = this.ai.live.connect({
        model: 'gemini-2.5-flash-native-audio-preview-09-2025',
        config: {
          responseModalities: [Modality.AUDIO],
          systemInstruction: SYSTEM_INSTRUCTION,
          tools: [{ functionDeclarations: [navigateToPage, searchDocuments] }],
        },
        callbacks: {
          onopen: this.handleOpen.bind(this),
          onmessage: this.handleMessage.bind(this),
          onclose: () => this.onStatusChange('disconnected'),
          onerror: (err) => {
            console.error(err);
            this.onStatusChange('error');
          },
        },
      });
    } catch (error) {
      console.error("Connection failed", error);
      this.onStatusChange('error');
    }
  }

  private async handleOpen() {
    this.onStatusChange('connected');
    
    // Start Microphone Stream
    try {
      const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
      if (!this.inputAudioContext) return;

      const source = this.inputAudioContext.createMediaStreamSource(stream);
      const scriptProcessor = this.inputAudioContext.createScriptProcessor(4096, 1, 1);
      
      scriptProcessor.onaudioprocess = (e) => {
        const inputData = e.inputBuffer.getChannelData(0);
        // Calculate volume for visualizer (input)
        let sum = 0;
        for (let i = 0; i < inputData.length; i++) sum += inputData[i] * inputData[i];
        this.onAudioLevel(Math.sqrt(sum / inputData.length));

        const pcmBlob = createPcmBlob(inputData);
        if (this.sessionPromise) {
            this.sessionPromise.then((session) => {
                session.sendRealtimeInput({ media: pcmBlob });
            });
        }
      };

      source.connect(scriptProcessor);
      scriptProcessor.connect(this.inputAudioContext.destination);
    } catch (err) {
      console.error("Microphone access denied", err);
      this.onStatusChange('error');
    }
  }

  private async handleMessage(message: LiveServerMessage) {
    // 1. Handle Audio
    const base64Audio = message.serverContent?.modelTurn?.parts?.[0]?.inlineData?.data;
    if (base64Audio && this.outputAudioContext && this.outputAnalyser) {
        // Ensure playback context is running
        if(this.outputAudioContext.state === 'suspended') {
            await this.outputAudioContext.resume();
        }

      const audioBuffer = await decodeAudioData(
        decode(base64Audio),
        this.outputAudioContext,
        24000
      );

      this.nextStartTime = Math.max(this.outputAudioContext.currentTime, this.nextStartTime);
      
      const source = this.outputAudioContext.createBufferSource();
      source.buffer = audioBuffer;
      source.connect(this.outputAnalyser); // Connect to analyser instead of destination directly
      source.start(this.nextStartTime);
      
      this.nextStartTime += audioBuffer.duration;
      this.sources.add(source);
      
      source.onended = () => this.sources.delete(source);
    }

    // 2. Handle Tool Calls
    if (message.toolCall) {
      for (const fc of message.toolCall.functionCalls) {
        console.log(`Tool Call: ${fc.name}`, fc.args);
        let result;
        try {
           result = await this.toolHandler(fc.name, fc.args);
        } catch(e) {
           result = { error: 'Failed to execute tool' };
        }
        
        // Send response back
        if (this.sessionPromise) {
            this.sessionPromise.then(session => {
                session.sendToolResponse({
                    functionResponses: {
                        id: fc.id,
                        name: fc.name,
                        response: { result },
                    }
                });
            });
        }
      }
    }

    // 3. Handle Interruption
    if (message.serverContent?.interrupted) {
      this.sources.forEach(s => s.stop());
      this.sources.clear();
      this.nextStartTime = 0;
    }
  }

  disconnect() {
      // Since there is no direct close method on the session object exposed in the quick snippet, 
      // we usually rely on closing the underlying websocket or just refreshing contexts.
      // However, the example shows session.close is not directly on the session object but implies clean up.
      // We will close contexts to stop processing.
      this.inputAudioContext?.close();
      this.outputAudioContext?.close();
      this.inputAudioContext = null;
      this.outputAudioContext = null;
      this.onStatusChange('disconnected');
  }
}