creates a common abstraction for providers

workers/site/providers/google.ts

@@ -1,97 +1,75 @@
 import { OpenAI } from "openai";
 import ChatSdk from "../lib/chat-sdk";
 import { StreamParams } from "../services/ChatService";
+import { BaseChatProvider, CommonProviderParams } from "./chat-stream-provider";
+
+export class GoogleChatProvider extends BaseChatProvider {
+  getOpenAIClient(param: CommonProviderParams): OpenAI {
+    return new OpenAI({
+      baseURL: "https://generativelanguage.googleapis.com/v1beta/openai",
+      apiKey: param.env.GEMINI_API_KEY,
+    });
+  }
+
+  getStreamParams(param: CommonProviderParams, safeMessages: any[]): any {
+    return {
+      model: param.model,
+      messages: safeMessages,
+      stream: true,
+    };
+  }
+
+  async processChunk(chunk: any, dataCallback: (data: any) => void): Promise<boolean> {
+    if (chunk.choices?.[0]?.finish_reason === "stop") {
+      dataCallback({
+        type: "chat",
+        data: {
+          choices: [
+            {
+              delta: { content: chunk.choices[0].delta.content || "" },
+              finish_reason: "stop",
+              index: chunk.choices[0].index,
+            },
+          ],
+        },
+      });
+      return true; // Break the stream
+    } else {
+      dataCallback({
+        type: "chat",
+        data: {
+          choices: [
+            {
+              delta: { content: chunk.choices?.[0]?.delta?.content || "" },
+              finish_reason: null,
+              index: chunk.choices?.[0]?.index || 0,
+            },
+          ],
+        },
+      });
+      return false; // Continue the stream
+    }
+  }
+}
 
 export class GoogleChatSdk {
+  private static provider = new GoogleChatProvider();
+
   static async handleGoogleStream(
     param: StreamParams,
     dataCallback: (data) => void,
   ) {
-    const {
-      preprocessedContext,
-      messages,
-      env,
-      maxTokens,
-      systemPrompt,
-      model,
-    } = param;
-
-    const assistantPrompt = ChatSdk.buildAssistantPrompt({
-      maxTokens: maxTokens,
-    });
-
-    const safeMessages = ChatSdk.buildMessageChain(messages, {
-      systemPrompt: systemPrompt,
-      model,
-      assistantPrompt,
-      toolResults: preprocessedContext,
-    });
-
-    const openai = new OpenAI({
-      baseURL: "https://generativelanguage.googleapis.com/v1beta/openai",
-      apiKey: param.env.GEMINI_API_KEY,
-    });
-
-    return GoogleChatSdk.streamGoogleResponse(
-      safeMessages,
+    return this.provider.handleStream(
       {
-        model: param.model,
+        systemPrompt: param.systemPrompt,
+        preprocessedContext: param.preprocessedContext,
         maxTokens: param.maxTokens,
-        openai: openai,
+        messages: param.messages,
+        model: param.model,
+        env: param.env,
+        disableWebhookGeneration: param.disableWebhookGeneration,
       },
       dataCallback,
     );
   }
-  private static async streamGoogleResponse(
-    messages: any[],
-    opts: {
-      model: string;
-      maxTokens: number | unknown | undefined;
-      openai: OpenAI;
-    },
-    dataCallback: (data: any) => void,
-  ) {
-    const chatReq = JSON.stringify({
-      model: opts.model,
-      messages: messages,
-      stream: true,
-    });
-
-    const googleStream = await opts.openai.chat.completions.create(
-      JSON.parse(chatReq),
-    );
-
-    for await (const chunk of googleStream) {
-      console.log(JSON.stringify(chunk));
-
-      if (chunk.choices?.[0]?.finishReason === "stop") {
-        dataCallback({
-          type: "chat",
-          data: {
-            choices: [
-              {
-                delta: { content: chunk.choices[0].delta.content || "" },
-                finish_reason: "stop",
-                index: chunk.choices[0].index,
-              },
-            ],
-          },
-        });
-        break;
-      } else {
-        dataCallback({
-          type: "chat",
-          data: {
-            choices: [
-              {
-                delta: { content: chunk.choices?.[0]?.delta?.content || "" },
-                finish_reason: null,
-                index: chunk.choices?.[0]?.index || 0,
-              },
-            ],
-          },
-        });
-      }
-    }
-  }
 }