Avoid hardcoding streamable methods

openai · Mar 5, 2023 · 90dbe7b · 90dbe7b
1 parent 937bcde
commit 90dbe7b
Showing 1 changed file with 22 additions and 1 deletion.
diff --git a/scripts/override_stream.ts b/scripts/override_stream.ts
@@ -151,7 +151,28 @@ project.addSourceFilesAtPaths(process.argv[process.argv.length - 1]);
 const sourceFile = project.getSourceFileOrThrow("api.ts");
 const declarations = sourceFile.getExportedDeclarations();
 
-for (const method of ["createCompletion", "createChatCompletion"]) {
+const methods = new Set(
+ sourceFile
+ .getInterfaces()
+ .filter((declaration) => {
+ return declaration
+ .getDescendantsOfKind(ts.SyntaxKind.PropertySignature)
+ .some((a) => a.getName().includes(`'stream'`));
+ })
+ .map((decl) =>
+ decl
+ .findReferencesAsNodes()
+ .map((refs) =>
+ refs
+ .getFirstAncestorByKind(ts.SyntaxKind.MethodDeclaration)
+ ?.getName()
+ )
+ )
+ .flat(2)
+ .filter((x): x is string => !!x)
+);
+
+for (const method of methods) {
  transformObject(method, declarations.get("OpenAIApiFp")![0]);
  transformObject(method, declarations.get("OpenAIApiFactory")![0]);
  transformClass(method, declarations.get("OpenAIApi")![0]);