-
Notifications
You must be signed in to change notification settings - Fork 323
/
AssemblyAIReader.ts
144 lines (132 loc) · 5 KB
/
AssemblyAIReader.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import { getEnv } from "@llamaindex/env";
import type {
BaseServiceParams,
SubtitleFormat,
TranscribeParams,
TranscriptParagraph,
TranscriptSentence,
} from "assemblyai";
import { AssemblyAI } from "assemblyai";
import { Document } from "../Node.js";
import type { BaseReader } from "./type.js";
type AssemblyAIOptions = Partial<BaseServiceParams>;
/**
* Base class for AssemblyAI Readers.
*/
abstract class AssemblyAIReader implements BaseReader {
protected client: AssemblyAI;
/**
* Creates a new AssemblyAI Reader.
* @param assemblyAIOptions The options to configure the AssemblyAI Reader.
* Configure the `assemblyAIOptions.apiKey` with your AssemblyAI API key, or configure it as the `ASSEMBLYAI_API_KEY` environment variable.
*/
constructor(assemblyAIOptions?: AssemblyAIOptions) {
let options = assemblyAIOptions;
if (!options) {
options = {};
}
if (!options.apiKey) {
options.apiKey = getEnv("ASSEMBLYAI_API_KEY");
}
if (!options.apiKey) {
throw new Error(
"No AssemblyAI API key provided. Pass an `apiKey` option, or configure the `ASSEMBLYAI_API_KEY` environment variable.",
);
}
this.client = new AssemblyAI(options as BaseServiceParams);
}
abstract loadData(params: TranscribeParams | string): Promise<Document[]>;
protected async transcribeOrGetTranscript(params: TranscribeParams | string) {
if (typeof params === "string") {
return await this.client.transcripts.get(params);
} else {
return await this.client.transcripts.transcribe(params);
}
}
protected async getTranscriptId(params: TranscribeParams | string) {
if (typeof params === "string") {
return params;
} else {
return (await this.client.transcripts.transcribe(params)).id;
}
}
}
/**
* Transcribe audio and read the transcript as a document using AssemblyAI.
*/
class AudioTranscriptReader extends AssemblyAIReader {
/**
* Transcribe audio or get a transcript and load the transcript as a document using AssemblyAI.
* @param params Parameters to transcribe an audio file or get an existing transcript.
* @returns A promise that resolves to a single document containing the transcript text.
*/
async loadData(params: TranscribeParams | string): Promise<Document[]> {
const transcript = await this.transcribeOrGetTranscript(params);
return [new Document({ text: transcript.text || undefined })];
}
}
/**
* Transcribe audio and return a document for each paragraph.
*/
class AudioTranscriptParagraphsReader extends AssemblyAIReader {
/**
* Transcribe audio or get a transcript, and returns a document for each paragraph.
* @param params The parameters to transcribe audio or get an existing transcript.
* @returns A promise that resolves to an array of documents, each containing a paragraph of the transcript.
*/
async loadData(params: TranscribeParams | string): Promise<Document[]> {
const transcriptId = await this.getTranscriptId(params);
const paragraphsResponse =
await this.client.transcripts.paragraphs(transcriptId);
return paragraphsResponse.paragraphs.map(
(p: TranscriptParagraph) => new Document({ text: p.text }),
);
}
}
/**
* Transcribe audio and return a document for each sentence.
*/
class AudioTranscriptSentencesReader extends AssemblyAIReader {
/**
* Transcribe audio or get a transcript, and returns a document for each sentence.
* @param params The parameters to transcribe audio or get an existing transcript.
* @returns A promise that resolves to an array of documents, each containing a sentence of the transcript.
*/
async loadData(params: TranscribeParams | string): Promise<Document[]> {
const transcriptId = await this.getTranscriptId(params);
const sentencesResponse =
await this.client.transcripts.sentences(transcriptId);
return sentencesResponse.sentences.map(
(p: TranscriptSentence) => new Document({ text: p.text }),
);
}
}
/**
* Transcribe audio a transcript and read subtitles for the transcript as `srt` or `vtt` format.
*/
class AudioSubtitlesReader extends AssemblyAIReader {
/**
* Transcribe audio or get a transcript and reads subtitles for the transcript as `srt` or `vtt` format.
* @param params The parameters to transcribe audio or get an existing transcript.
* @param subtitleFormat The format of the subtitles, either `srt` or `vtt`.
* @returns A promise that resolves a document containing the subtitles as the page content.
*/
async loadData(
params: TranscribeParams | string,
subtitleFormat: SubtitleFormat = "srt",
): Promise<Document[]> {
const transcriptId = await this.getTranscriptId(params);
const subtitles = await this.client.transcripts.subtitles(
transcriptId,
subtitleFormat,
);
return [new Document({ text: subtitles })];
}
}
export {
AudioSubtitlesReader,
AudioTranscriptParagraphsReader,
AudioTranscriptReader,
AudioTranscriptSentencesReader,
};
export type { AssemblyAIOptions, SubtitleFormat, TranscribeParams };