Merge pull request #129 from erictik:porxy-fetch

Porxy-fetch
erictik · Jun 16, 2023 · 8d4dcf9 · 8d4dcf9
2 parents 8584561 + 18d16e6
commit 8d4dcf9
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 38 deletions.
diff --git a/src/interfaces/config.ts b/src/interfaces/config.ts
@@ -1,3 +1,8 @@
+import WebSocket from "isomorphic-ws";
+
+export type FetchFn = typeof fetch;
+export type WebSocketFn = typeof WebSocket;
+
 export interface MJConfig {
   ChannelId: string;
   SalaiToken: string;
@@ -10,7 +15,8 @@ export interface MJConfig {
   HuggingFaceToken?: string;
   DiscordBaseUrl: string;
   WsBaseUrl: string;
-  ProxyUrl?: string;
+  fetch: FetchFn;
+  WebSocket: WebSocketFn;
 }
 export interface MJConfigParam {
   SalaiToken: string;
@@ -24,7 +30,8 @@ export interface MJConfigParam {
   SessionId?: string;
   DiscordBaseUrl?: string;
   WsBaseUrl?: string;
-  ProxyUrl?: string;
+  fetch?: FetchFn;
+  WebSocket?: WebSocketFn;
 }
 
 export const DefaultMJConfig: MJConfig = {
@@ -36,4 +43,6 @@ export const DefaultMJConfig: MJConfig = {
   MaxWait: 200,
   DiscordBaseUrl: "https://discord.com",
   WsBaseUrl: "wss://gateway.discord.gg?v=9&encoding=json&compress=gzip-stream",
+  fetch: fetch,
+  WebSocket: WebSocket,
 };
diff --git a/src/midjourne.api.ts b/src/midjourne.api.ts
@@ -10,10 +10,7 @@ interface CustomRequestInit extends RequestInit {
 export class MidjourneyApi {
   private apiQueue = CreateQueue(1);
   UpId = Date.now() % 10; // upload id
-  constructor(public config: MJConfig) {
-    if (this.config.ProxyUrl && this.config.ProxyUrl !== "") {
-    }
-  }
+  constructor(public config: MJConfig) {}
   // limit the number of concurrent interactions
   protected async safeIteractions(payload: any) {
     return this.apiQueue.addTask(
@@ -34,19 +31,14 @@ export class MidjourneyApi {
         "Content-Type": "application/json",
         Authorization: this.config.SalaiToken,
       };
-      console.log("api.DiscordBaseUrl", this.config.DiscordBaseUrl);
-
-      let fetchUrl = `${this.config.ProxyUrl}?url=${encodeURIComponent(
-        `${this.config.DiscordBaseUrl}/api/v9/interactions`
-      )}`;
-
-      fetchUrl = `${this.config.DiscordBaseUrl}/api/v9/interactions`;
-      console.log("api.fetchUrl", fetchUrl);
-      const response = await fetch(fetchUrl, {
-        method: "POST",
-        body: JSON.stringify(payload),
-        headers: headers,
-      });
+      const response = await this.config.fetch(
+        `${this.config.DiscordBaseUrl}/api/v9/interactions`,
+        {
+          method: "POST",
+          body: JSON.stringify(payload),
+          headers: headers,
+        }
+      );
       callback && callback(response.status);
       //discord api rate limit
       await sleep(950);
@@ -278,7 +270,7 @@ export class MidjourneyApi {
     let file_size;
 
     if (fileUrl.startsWith("http")) {
-      const response = await fetch(fileUrl);
+      const response = await this.config.fetch(fileUrl);
       fileData = await response.arrayBuffer();
       mimeType = response.headers.get("content-type");
       filename = path.basename(fileUrl) || "image.png";
@@ -321,7 +313,7 @@ export class MidjourneyApi {
       `${this.config.DiscordBaseUrl}/api/v9/channels/${this.config.ChannelId}/attachments`
     );
     const body = { files };
-    const response = await fetch(url.toString(), {
+    const response = await this.config.fetch(url.toString(), {
       headers,
       method: "POST",
       body: JSON.stringify(body),
@@ -342,7 +334,7 @@ export class MidjourneyApi {
   ): Promise<void> {
     const body = new Uint8Array(data);
     const headers = { "content-type": contentType };
-    const response = await fetch(slot.upload_url, {
+    const response = await this.config.fetch(slot.upload_url, {
       method: "PUT",
       headers,
       body,

diff --git a/src/midjourney.message.ts b/src/midjourney.message.ts
@@ -20,8 +20,6 @@ export class MidjourneyMessage {
       ...DefaultMJConfig,
       ...defaults,
     };
-    if (this.config.ProxyUrl && this.config.ProxyUrl !== "") {
-    }
   }
   protected log(...args: any[]) {
     this.config.Debug && console.log(...args, new Date().toISOString());
@@ -164,7 +162,7 @@ export class MidjourneyMessage {
       "Content-Type": "application/json",
       Authorization: this.config.SalaiToken,
     };
-    const response = await fetch(
+    const response = await this.config.fetch(
       `${this.config.DiscordBaseUrl}/api/v10/channels/${this.config.ChannelId}/messages?limit=${limit}`,
       {
         headers,

diff --git a/src/verify.human.ts b/src/verify.human.ts
@@ -1,9 +1,11 @@
 import { HfInference } from "@huggingface/inference";
+import { MJConfig } from "./interfaces";
 export class VerifyHuman {
   private inference: HfInference;
 
-  constructor(HuggingFaceToken: string) {
-    if (HuggingFaceToken === "") {
+  constructor(public config: MJConfig) {
+    const { HuggingFaceToken } = config;
+    if (HuggingFaceToken === "" || HuggingFaceToken) {
       throw new Error("HuggingFaceToken is required");
     }
     this.inference = new HfInference(HuggingFaceToken);
@@ -12,7 +14,7 @@ export class VerifyHuman {
   async verify(imageUri: string, categories: string[]) {
     console.log("verify----start", imageUri, categories);
     const imageCates = await this.inference.imageClassification({
-      data: await (await fetch(imageUri)).blob(),
+      data: await (await this.config.fetch(imageUri)).blob(),
       model: "google/vit-base-patch16-224",
     });
     console.log("verify----response", { imageCates });

diff --git a/src/ws.message.ts b/src/ws.message.ts
@@ -10,8 +10,6 @@ import {
 import { MidjourneyApi } from "./midjourne.api";
 import { VerifyHuman } from "./verify.human";
 import WebSocket from "isomorphic-ws";
-// import { HttpsProxyAgent } from "https-proxy-agent";
-
 export class WsMessage {
   ws: WebSocket;
   MJBotId = "936929561302675456";
@@ -21,14 +19,9 @@ export class WsMessage {
   private waitMjEvents: Map<string, WaitMjEvent> = new Map();
   private reconnectTime: boolean[] = [];
   private heartbeatInterval = 0;
-  // agent?: HttpsProxyAgent<string>;
 
   constructor(public config: MJConfig, public MJApi: MidjourneyApi) {
-    if (this.config.ProxyUrl && this.config.ProxyUrl !== "") {
-      // this.agent = new HttpsProxyAgent(this.config.ProxyUrl);
-    }
-    // const agent = this.agent;
-    this.ws = new WebSocket(this.config.WsBaseUrl);
+    this.ws = new this.config.WebSocket(this.config.WsBaseUrl);
     this.ws.addEventListener("open", this.open.bind(this));
   }
 
@@ -53,7 +46,7 @@ export class WsMessage {
   private reconnect() {
     if (this.closed) return;
     // const agent = this.agent;
-    this.ws = new WebSocket(this.config.WsBaseUrl);
+    this.ws = new this.config.WebSocket(this.config.WsBaseUrl);
     this.ws.addEventListener("open", this.open.bind(this));
   }
   // After opening ws
@@ -218,7 +211,7 @@ export class WsMessage {
     const uri = embeds[0].image.url;
     const categories = components[0].components;
     const classify = categories.map((c: any) => c.label);
-    const verifyClient = new VerifyHuman(HuggingFaceToken);
+    const verifyClient = new VerifyHuman(this.config);
     const category = await verifyClient.verify(uri, classify);
     if (category) {
       const custom_id = categories.find(