AI Sales Copilot 개발 기록 - (2) Streaming Queue

{
  "sectionKey": "HOOK_POINTS",
  "sectionStatus": "STREAMING",
  "content": {
    "title": "Seamless Communication",
    "description": "Operating in the similar timezone as Happy Client, Tsunagaru Inc. ensures effective and timely communication, facilitating smooth project execution and collaboration.",
    "factors": ["Speed of Implementation", "Technical Support"],
  }
}

type SectionKey = 'HOOK_POINTS' | 'OFFERS' | 'AGENDA' | 'QUESTIONS' | 'RISKS' | 'TERMS';
type SectionStatus = 'NOT_STARTED' | 'STREAMING' | 'COMPLETE' | 'UNHEALTHY';
type SectionQueue = ReplaySubject<Delta>; // Delta는 데이터 조각의 명칭

class StreamingQueue {
  /**
   * sectionKey별로 섹션의 queue를 저장하는 맵
   */
  private sectionQueueMap: Map<SectionKey, SectionQueue> = new Map();

  /**
   * 전체 섹션을 순서에 따라 재구성된 형태로 저장하는 상위 레벨의 queue
   */
  private processedSnapshotQueue: ReplaySubject<SectionQueue> = new ReplaySubject();

  /**
   * 인스턴스 생성 시점에 sectionKey 배열을 인자로 받아 sectionQueueMap과 processedSnapshotQueue를 초기화한다.
   */
  constructor({ sectionKeys }: { sectionKeys: SectionKey[] }) {
    sectionKeys.forEach((sectionKey) => {
      const sectionQueue = new ReplaySubject<Delta>();
      this.sectionQueueMap.set(sectionKey, sectionQueue);
      this.processedSnapshotQueue.next(sectionQueue);
    });
  }

  /**
   * 해당되는 sectionKey에 매칭되는 queue를 업데이트한다.
   * 새로운 데이터 조각이 들어온 시점에 실행한다.
   */
  public enqueue({
    sectionKey,
    delta,
  }: {
    sectionKey: SectionKey;
    delta: Delta;
  }): void {
    const targetsectionQueue = this.sectionQueueMap.get(sectionKey);

    if (!targetsectionQueue) {
      throw new Error(`no matched section queue for ${sectionKey}`);
    }

    targetsectionQueue.next(delta);
  }

  /**
   * 해당되는 sectionKey에 매칭되는 queue를 종료 처리한다.
   * 섹션의 스트리밍 상태가 완료(또는 에러)로 바뀌면 실행한다.
   */
  public completeSectionQueue(sectionKey: SectionKey): void {
    const targetSectionQueue = this.sectionQueueMap.get(sectionKey);

    if (!targetSectionQueue) {
      throw new Error(`no matched section queue for ${sectionKey}`);
    }

    targetSectionQueue.complete();
  }

  /**
   * 사용자 화면에 컨텐츠를 렌더링하기 위해 queue를 실행한다.
   * processedSnapshotQueue > sectionQueue 순서로 파이프 스트림이 흘러간다.
   */
  public async subscribe(
    callback: (delta: Delta) => Promise<void>,
  ): Promise<void> {
    from(this.processedSnapshotQueue)
      .pipe(
        concatMap((sectionQueue) => // 병합 연산자 중 순서가 보장되는 concatMap을 사용한다.
          sectionQueue.pipe(
            bufferCount(1), // processedSnapshotQueue 에서 sectionQueue 를 1개씩 처리하기 위해 사용한다.
            concatMap((deltas) =>
              from(deltas).pipe(
                bufferCount(1), // sectionQueue 에서 delta 를 1개씩 처리하기 위해 사용한다.
                delay(100), // 유저에게 스트리밍 효과를 주기 위해 100ms 대기한다.
                concatMap(([delta]) => callback(delta)), // 화면에 렌더링하기 위한 로직을 실행한다.
              ),
            ),
          ),
        ),
      )
      .subscribe();
  }

  /**
   * 전체 queue를 종료 처리한다.
   * 모든 섹션의 스트리밍 상태가 완료(또는 에러)로 바뀌면 실행한다.
   */
  public unsubscribe(): void {
    this.sectionQueueMap.forEach((sectionQueue) => {
      sectionQueue.complete();
      sectionQueue.unsubscribe();
    });
    this.processedSnapshotQueue.complete();
    this.processedSnapshotQueue.unsubscribe();
  }
}

type StreamingContextValue = {
  snapshot: Delta;
}

const StreamingContext = React.createContext<StreamingContextValue | null>(null);

function StreamingContextProvider({ children }: Props) {
  const [snapshot, setSnapshot] = useState<Delta>(new Delta()); // UI에서 참조할 snapshot 상태
  const streamingQueue = React.useRef<StreamingQueue | null>(null);

  /**
   * queue 구독
   */
  function subscribeStreamingQueue() {
    if (!streamingQueue.current) return;

    streamingQueue.current.subscribe(async (delta) => {
      // 델타로부터 업데이트된 섹션을 추출
      const section = getUpdatedSection(delta, snapshot);

      if (['STREAMING', 'COMPLETE'].includes(section.status)) {
        // 렌더링할 스냅샷을 새로운 델타로 교체
        setSnapshot(delta);
      }

      if (section.status === 'COMPLETE') {
        // sectionQueue 완료 처리
        streamingQueue.current.completeSectionQueue(section.key);
      }

      if (section.status === 'UNHEALTHY') {
        // queue 구독 해제 및 초기화
        streamingQueue.current.unsubscribe();
        streamingQueue.current = null;
      }
    });
  }

  /**
   * 이벤트 리스너 등록
   */
  function addDeltaEventListener() {
    if (!streamingQueue.current) return;

    // 델타 이벤트를 전달하는 엔진의 tune 인스턴스 가져오기
    const tune = getTune();
    
    tune.addEventListener('VERSION_UPDATED', (delta) => {
      // 델타로부터 업데이트된 섹션을 추출
      const section = getUpdatedSection(delta, snapshot);

      // 새로운 데이터가 들어왔을 때 queue에 추가
      streamingQueue.current.enqueue({
        sectionKey: section.key,
        delta,
      });

      if (['UNHEALTHY', 'COMPLETE'].includes(delta.streamingStatus)) {
        // queue 구독 해제 및 초기화
        streamingQueue.current.unsubscribe();
        streamingQueue.current = null;
      }
    });
  }

  /**
   * 마운트 후 최초 한번 다음 로직들을 실행한다.
   */
  useEffect(() => {
    streamingQueue.current = new StreamingQueue({ sectionKeys });

    subscribeStreamingQueue();
    addDeltaEventListener();
  }, []);

  return (
    <StreamingContext.Provider value={{ snapshot }}>
      {children}
    </StreamingContext.Provider>
  );
}