@livekit/agents 0.0.0-20260120144724

package/src/tokenize/basic/sentence.ts

@@ -0,0 +1,89 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Split the text into sentences.
+ */
+export const splitSentences = (
+  text: string,
+  minLength = 20,
+  retainFormat: boolean = false,
+): [string, number, number][] => {
+  const alphabets = /([A-Za-z])/g;
+  const prefixes = /(Mr|St|Mrs|Ms|Dr)[.]/g;
+  const suffixes = /(Inc|Ltd|Jr|Sr|Co)/g;
+  const starters =
+    /(Mr|Mrs|Ms|Dr|Prof|Capt|Cpt|Lt|He\s|She\s|It\s|They\s|Their\s|Our\s|We\s|But\s|However\s|That\s|This\s|Wherever)/g;
+  const acronyms = /([A-Z][.][A-Z][.](?:[A-Z][.])?)/g;
+  const websites = /[.](com|net|org|io|gov|edu|me)/g;
+  const digits = /([0-9])/g;
+  const dots = /\.{2,}/g;
+
+  if (retainFormat) {
+    text = text.replaceAll('\n', '<nel><stop>');
+  } else {
+    text = text.replaceAll('\n', ' ');
+  }
+
+  text = text.replaceAll(prefixes, '$1<prd>');
+  text = text.replaceAll(websites, '<prd>$2');
+  text = text.replaceAll(new RegExp(`${digits.source}[.]${digits.source}`, 'g'), '$1<prd>$2');
+  text = text.replaceAll(dots, (match) => '<prd>'.repeat(match.length));
+  text = text.replaceAll('Ph.D.', 'Ph<prd>D<prd>');
+  text = text.replaceAll(new RegExp(`\\s${alphabets.source}[.] `, 'g'), ' $1<prd> ');
+  text = text.replaceAll(new RegExp(`${acronyms.source} ${starters.source}`, 'g'), '$1<stop> $2');
+  text = text.replaceAll(
+    new RegExp(`${alphabets.source}[.]${alphabets.source}[.]${alphabets.source}[.]`, 'g'),
+    '$1<prd>$2<prd>$3<prd>',
+  );
+  text = text.replaceAll(
+    new RegExp(`${alphabets.source}[.]${alphabets.source}[.]`, 'g'),
+    '$1<prd>$2<prd>',
+  );
+  text = text.replaceAll(
+    new RegExp(` ${suffixes.source}[.] ${starters.source}`, 'g'),
+    '$1<stop> $2',
+  );
+  text = text.replaceAll(new RegExp(` ${suffixes.source}[.]`, 'g'), '$1<prd>');
+  text = text.replaceAll(new RegExp(` ${alphabets.source}[.]`, 'g'), '$1<prd>');
+  text = text.replaceAll('.”', '”.');
+  text = text.replaceAll('."', '".');
+  text = text.replaceAll('!"', '"!');
+  text = text.replaceAll('?"', '"?');
+  text = text.replaceAll('.', '.<stop>');
+  text = text.replaceAll('?', '?<stop>');
+  text = text.replaceAll('!', '!<stop>');
+  text = text.replaceAll('<prd>', '.');
+
+  if (retainFormat) {
+    text = text.replaceAll('<nel>', '\n');
+  }
+
+  const split = text.split('<stop>');
+  text = text.replaceAll('<stop>', '');
+
+  const sentences: [string, number, number][] = [];
+  let buf = '';
+  let start = 0;
+  let end = 0;
+  const prePad = retainFormat ? '' : ' ';
+  for (const match of split) {
+    const sentence = retainFormat ? match : match.trim();
+    if (!sentence) continue;
+
+    buf += prePad + sentence;
+    end += match.length;
+    if (buf.length > minLength) {
+      sentences.push([buf.slice(prePad.length), start, end]);
+      start = end;
+      buf = '';
+    }
+  }
+
+  if (buf) {
+    sentences.push([buf.slice(prePad.length), start, text.length - 1]);
+  }
+
+  return sentences;
+};

package/src/tokenize/basic/word.ts

@@ -0,0 +1,27 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { PUNCTUATIONS } from '../tokenizer.js';
+
+/**
+ * Split the text into words.
+ */
+export const splitWords = (text: string, ignorePunctuation = true): [string, number, number][] => {
+  const re = /\S+/g;
+  const words: [string, number, number][] = [];
+
+  let arr;
+  while ((arr = re.exec(text)) !== null) {
+    let word = arr[0];
+    const start = arr.index;
+    const end = start + word.length;
+
+    if (ignorePunctuation) {
+      word = word.replace(new RegExp(`[${PUNCTUATIONS.join('')}]`, 'g'), '');
+    }
+
+    words.push([word, start, end]);
+  }
+
+  return words;
+};

package/src/tokenize/index.ts

@@ -0,0 +1,16 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import * as basic from './basic/index.js';
+
+export {
+  type TokenData,
+  SentenceTokenizer,
+  SentenceStream,
+  WordTokenizer,
+  WordStream,
+} from './tokenizer.js';
+
+export { BufferedSentenceStream, BufferedTokenStream, BufferedWordStream } from './token_stream.js';
+
+export { basic };

package/src/tokenize/token_stream.ts

@@ -0,0 +1,180 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { AsyncIterableQueue, shortuuid } from '../utils.js';
+import type { TokenData } from './tokenizer.js';
+import { SentenceStream, WordStream } from './tokenizer.js';
+
+type TokenizeFunc = (x: string) => string[] | [string, number, number][];
+
+export class BufferedTokenStream implements AsyncIterableIterator<TokenData> {
+  protected queue = new AsyncIterableQueue<TokenData>();
+  protected closed = false;
+
+  #func: TokenizeFunc;
+  #minTokenLength: number;
+  #minContextLength: number;
+  #bufTokens: string[] = [];
+  #inBuf = '';
+  #outBuf = '';
+  #currentSegmentId: string;
+
+  constructor(func: TokenizeFunc, minTokenLength: number, minContextLength: number) {
+    this.#func = func;
+    this.#minTokenLength = minTokenLength;
+    this.#minContextLength = minContextLength;
+
+    this.#currentSegmentId = shortuuid();
+  }
+
+  /** Push a string of text into the token stream */
+  pushText(text: string) {
+    if (this.closed) {
+      throw new Error('Stream is closed');
+    }
+
+    this.#inBuf += text;
+    if (this.#inBuf.length < this.#minContextLength) return;
+
+    while (true) {
+      const tokens = this.#func(this.#inBuf);
+      if (tokens.length <= 1) break;
+
+      if (this.#outBuf) this.#outBuf += ' ';
+
+      const tok = tokens.shift()!;
+      let tokText: string;
+      if (Array.isArray(tok)) {
+        tokText = tok[0];
+      } else {
+        tokText = tok;
+      }
+
+      this.#outBuf += tokText;
+
+      if (this.#outBuf.length >= this.#minTokenLength) {
+        this.queue.put({ token: this.#outBuf, segmentId: this.#currentSegmentId });
+        this.#outBuf = '';
+      }
+
+      if (typeof tok! !== 'string') {
+        this.#inBuf = this.#inBuf.slice(tok![2]);
+      } else {
+        this.#inBuf = this.#inBuf
+          .slice(Math.max(0, this.#inBuf.indexOf(tok)) + tok.length)
+          .trimStart();
+      }
+    }
+  }
+
+  /** Flush the stream, causing it to process all pending text */
+  flush() {
+    if (this.closed) {
+      throw new Error('Stream is closed');
+    }
+
+    if (this.#inBuf || this.#outBuf) {
+      const tokens = this.#func(this.#inBuf);
+      if (tokens) {
+        if (this.#outBuf) this.#outBuf += ' ';
+
+        if (Array.isArray(tokens[0])) {
+          this.#outBuf += tokens.map((tok) => tok[0]).join(' ');
+        } else {
+          this.#outBuf += tokens.join(' ');
+        }
+      }
+
+      if (this.#outBuf) {
+        this.queue.put({ token: this.#outBuf, segmentId: this.#currentSegmentId });
+      }
+
+      this.#currentSegmentId = shortuuid();
+    }
+
+    this.#inBuf = '';
+    this.#outBuf = '';
+  }
+
+  /** Mark the input as ended and forbid additional pushes */
+  endInput() {
+    if (this.closed) {
+      throw new Error('Stream is closed');
+    }
+    this.flush();
+    this.close();
+  }
+
+  next(): Promise<IteratorResult<TokenData>> {
+    return this.queue.next();
+  }
+
+  /** Close both the input and output of the token stream */
+  close() {
+    this.queue.close();
+    this.closed = true;
+  }
+
+  [Symbol.asyncIterator](): BufferedTokenStream {
+    return this;
+  }
+}
+
+export class BufferedSentenceStream extends SentenceStream {
+  #stream: BufferedTokenStream;
+
+  constructor(func: TokenizeFunc, minTokenLength: number, minContextLength: number) {
+    super();
+    this.#stream = new BufferedTokenStream(func, minTokenLength, minContextLength);
+  }
+
+  pushText(text: string) {
+    this.#stream.pushText(text);
+  }
+
+  flush() {
+    this.#stream.flush();
+  }
+
+  close() {
+    super.close();
+    this.#stream.close();
+  }
+
+  endInput() {
+    this.#stream.endInput();
+  }
+
+  next(): Promise<IteratorResult<TokenData>> {
+    return this.#stream.next();
+  }
+}
+
+export class BufferedWordStream extends WordStream {
+  #stream: BufferedTokenStream;
+
+  constructor(func: TokenizeFunc, minTokenLength: number, minContextLength: number) {
+    super();
+    this.#stream = new BufferedTokenStream(func, minTokenLength, minContextLength);
+  }
+
+  pushText(text: string) {
+    this.#stream.pushText(text);
+  }
+
+  flush() {
+    this.#stream.flush();
+  }
+
+  endInput() {
+    this.#stream.endInput();
+  }
+
+  close() {
+    this.#stream.close();
+  }
+
+  next(): Promise<IteratorResult<TokenData>> {
+    return this.#stream.next();
+  }
+}

package/src/tokenize/tokenizer.test.ts

@@ -0,0 +1,255 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { describe, expect, it } from 'vitest';
+import { SentenceTokenizer, WordTokenizer, hyphenateWord } from './basic/index.js';
+import { splitParagraphs } from './basic/paragraph.js';
+
+const TEXT =
+  'Hi! ' +
+  'LiveKit is a platform for live audio and video applications and services. ' +
+  'R.T.C stands for Real-Time Communication... again R.T.C. ' +
+  'Mr. Theo is testing the sentence tokenizer. ' +
+  'This is a test. Another test. ' +
+  'A short sentence. ' +
+  'A longer sentence that is longer than the previous sentence. ' +
+  'f(x) = x * 2.54 + 42. ' +
+  'Hey! Hi! Hello! ';
+
+const EXPECTED_MIN_20 = [
+  'Hi! LiveKit is a platform for live audio and video applications and services.',
+  'R.T.C stands for Real-Time Communication... again R.T.C.',
+  'Mr. Theo is testing the sentence tokenizer.',
+  'This is a test. Another test.',
+  'A short sentence. A longer sentence that is longer than the previous sentence.',
+  'f(x) = x * 2.54 + 42.',
+  'Hey! Hi! Hello!',
+];
+
+const WORDS_TEXT = 'This is a test. Blabla another test! multiple consecutive spaces:     done';
+const WORDS_EXPECTED = [
+  'This',
+  'is',
+  'a',
+  'test',
+  'Blabla',
+  'another',
+  'test',
+  'multiple',
+  'consecutive',
+  'spaces',
+  'done',
+];
+
+const WORDS_PUNCT_TEXT =
+  'This is <phoneme alphabet="cmu-arpabet" ph="AE K CH UW AH L IY">actually</phoneme> tricky to handle.';
+const WORDS_PUNCT_EXPECTED = [
+  'This',
+  'is',
+  '<phoneme',
+  'alphabet="cmu-arpabet"',
+  'ph="AE',
+  'K',
+  'CH',
+  'UW',
+  'AH',
+  'L',
+  'IY">actually</phoneme>',
+  'tricky',
+  'to',
+  'handle.',
+];
+
+const HYPHENATOR_TEXT = ['Segment', 'expected', 'communication', 'window', 'welcome', 'bedroom'];
+const HYPHENATOR_EXPECTED = [
+  ['Seg', 'ment'],
+  ['ex', 'pect', 'ed'],
+  ['com', 'mu', 'ni', 'ca', 'tion'],
+  ['win', 'dow'],
+  ['wel', 'come'],
+  ['bed', 'room'],
+];
+
+const PARAGRAPH_TEST_CASES: [string, [string, number, number][]][] = [
+  ['Single paragraph.', [['Single paragraph.', 0, 17]]],
+  [
+    'Paragraph 1.\n\nParagraph 2.',
+    [
+      ['Paragraph 1.', 0, 12],
+      ['Paragraph 2.', 14, 26],
+    ],
+  ],
+  [
+    'Para 1.\n\nPara 2.\n\nPara 3.',
+    [
+      ['Para 1.', 0, 7],
+      ['Para 2.', 9, 16],
+      ['Para 3.', 18, 25],
+    ],
+  ],
+  ['\n\nParagraph with leading newlines.', [['Paragraph with leading newlines.', 2, 34]]],
+  ['Paragraph with trailing newlines.\n\n', [['Paragraph with trailing newlines.', 0, 33]]],
+  [
+    '\n\n  Paragraph with leading and trailing spaces.  \n\n',
+    [['Paragraph with leading and trailing spaces.', 4, 47]],
+  ],
+  [
+    'Para 1.\n\n\n\nPara 2.', // Multiple newlines between paragraphs
+    [
+      ['Para 1.', 0, 7],
+      ['Para 2.', 11, 18],
+    ],
+  ],
+  [
+    'Para 1.\n \n \nPara 2.', // Newlines with spaces between paragraphs
+    [
+      ['Para 1.', 0, 7],
+      ['Para 2.', 12, 19],
+    ],
+  ],
+  [
+    '', // Empty string
+    [],
+  ],
+  [
+    '\n\n\n', // Only newlines
+    [],
+  ],
+  [
+    'Line 1\nLine 2\nLine 3', // Single paragraph with newlines
+    [['Line 1\nLine 2\nLine 3', 0, 20]],
+  ],
+];
+
+describe('tokenizer', () => {
+  describe('SentenceTokenizer', () => {
+    const tokenizer = new SentenceTokenizer();
+
+    it('should tokenize sentences correctly', () => {
+      expect(tokenizer.tokenize(TEXT).every((x, i) => EXPECTED_MIN_20[i] === x)).toBeTruthy();
+    });
+
+    it('should stream tokenize sentences correctly', async () => {
+      const pattern = [1, 2, 4];
+      let text = TEXT;
+      const chunks = [];
+      const patternIter = Array(Math.ceil(text.length / pattern.reduce((sum, num) => sum + num, 0)))
+        .fill(pattern)
+        .flat()
+        [Symbol.iterator]();
+
+      for (const size of patternIter) {
+        if (!text) break;
+        chunks.push(text.slice(undefined, size));
+        text = text.slice(size);
+      }
+      const stream = tokenizer.stream();
+      for (const chunk of chunks) {
+        stream.pushText(chunk);
+      }
+      stream.endInput();
+      stream.close();
+
+      for (const x of EXPECTED_MIN_20) {
+        await stream.next().then((value) => {
+          if (value.value) {
+            expect(value.value.token).toStrictEqual(x);
+          }
+        });
+      }
+    });
+  });
+  describe('WordTokenizer', () => {
+    const tokenizer = new WordTokenizer();
+
+    it('should tokenize words correctly', () => {
+      expect(tokenizer.tokenize(WORDS_TEXT).every((x, i) => WORDS_EXPECTED[i] === x)).toBeTruthy();
+    });
+
+    it('should stream tokenize words correctly', async () => {
+      const pattern = [1, 2, 4];
+      let text = WORDS_TEXT;
+      const chunks = [];
+      const patternIter = Array(Math.ceil(text.length / pattern.reduce((sum, num) => sum + num, 0)))
+        .fill(pattern)
+        .flat()
+        [Symbol.iterator]();
+
+      for (const size of patternIter) {
+        if (!text) break;
+        chunks.push(text.slice(undefined, size));
+        text = text.slice(size);
+      }
+      const stream = tokenizer.stream();
+      for (const chunk of chunks) {
+        stream.pushText(chunk);
+      }
+      stream.endInput();
+      stream.close();
+
+      for (const x of WORDS_EXPECTED) {
+        await stream.next().then((value) => {
+          if (value.value) {
+            expect(value.value.token).toStrictEqual(x);
+          }
+        });
+      }
+    });
+
+    describe('punctuation handling', () => {
+      const tokenizerPunct = new WordTokenizer(false);
+
+      it('should tokenize words correctly', () => {
+        expect(
+          tokenizerPunct.tokenize(WORDS_PUNCT_TEXT).every((x, i) => WORDS_PUNCT_EXPECTED[i] === x),
+        ).toBeTruthy();
+      });
+
+      it('should stream tokenize words correctly', async () => {
+        const pattern = [1, 2, 4];
+        let text = WORDS_PUNCT_TEXT;
+        const chunks = [];
+        const patternIter = Array(
+          Math.ceil(text.length / pattern.reduce((sum, num) => sum + num, 0)),
+        )
+          .fill(pattern)
+          .flat()
+          [Symbol.iterator]();
+
+        for (const size of patternIter) {
+          if (!text) break;
+          chunks.push(text.slice(undefined, size));
+          text = text.slice(size);
+        }
+        const stream = tokenizerPunct.stream();
+        for (const chunk of chunks) {
+          stream.pushText(chunk);
+        }
+        stream.endInput();
+        stream.close();
+
+        for (const x of WORDS_PUNCT_EXPECTED) {
+          await stream.next().then((value) => {
+            if (value.value) {
+              expect(value.value.token).toStrictEqual(x);
+            }
+          });
+        }
+      });
+    });
+  });
+  describe('hyphenateWord', () => {
+    it('should hyphenate correctly', () => {
+      HYPHENATOR_TEXT.forEach((x, i) => {
+        expect(hyphenateWord(x)).toStrictEqual(HYPHENATOR_EXPECTED[i]);
+      });
+    });
+  });
+  describe('splitParagraphs', () => {
+    it('should tokenize paragraphs correctly', () => {
+      PARAGRAPH_TEST_CASES.forEach(([a, b]) => {
+        expect(splitParagraphs(a)).toStrictEqual(b);
+      });
+    });
+  });
+});