feat: allow passing other llms when using chat()

README.md

@@ -5,7 +5,7 @@
 > [!NOTE]  
 > **This project is in GA Stage.**
 >
-> The Upstash Professional Support fully covers this project. It receives regular updates, and bug fixes. 
+> The Upstash Professional Support fully covers this project. It receives regular updates, and bug fixes.
 > The Upstash team is committed to maintaining and improving its functionality.
 
 **QStash** is an HTTP based messaging and scheduling solution for serverless and

src/client/http.ts

@@ -38,6 +38,7 @@ export type UpstashRequest = {
    * @default true
    */
   parseResponseAsJson?: boolean;
+  baseUrl?: string;
 };
 export type UpstashResponse<TResult> = TResult & { error?: string };
 
@@ -150,6 +151,7 @@ export class HttpClient implements Requester {
       await reader.cancel();
     }
   }
+
   private requestWithBackoff = async (
     request: UpstashRequest
   ): Promise<{
@@ -183,16 +185,17 @@ export class HttpClient implements Requester {
   private processRequest = (request: UpstashRequest): [string, RequestOptions] => {
     //@ts-expect-error caused by undici and bunjs type overlap
     const headers = new Headers(request.headers);
-    headers.set("Authorization", this.authorization);
-
+    if (!headers.has("Authorization")) {
+      headers.set("Authorization", this.authorization);
+    }
     const requestOptions: RequestOptions = {
       method: request.method,
       headers,
       body: request.body,
       keepalive: request.keepalive,
     };
 
-    const url = new URL([this.baseUrl, ...request.path].join("/"));
+    const url = new URL([request.baseUrl ?? this.baseUrl, ...request.path].join("/"));
     if (request.query) {
       for (const [key, value] of Object.entries(request.query)) {
         if (value !== undefined) {

src/client/llm/chat.test.ts

@@ -26,6 +26,7 @@ describe("Test Qstash chat", () => {
     "should respond to prompt",
     async () => {
       const response = await client.chat().prompt({
+        provider: "upstash",
         model: "meta-llama/Meta-Llama-3-8B-Instruct",
         system: "from now on, foo is whale",
         user: "what exactly is foo?",
@@ -44,6 +45,7 @@ describe("Test Qstash chat", () => {
     "should respond to create",
     async () => {
       const response = await client.chat().create({
+        provider: "upstash",
         model: "meta-llama/Meta-Llama-3-8B-Instruct",
         messages: [
           {
@@ -70,6 +72,7 @@ describe("Test Qstash chat", () => {
     "should stream prompt",
     async () => {
       const response = await client.chat().prompt({
+        provider: "upstash",
         model: "meta-llama/Meta-Llama-3-8B-Instruct",
         system: "from now on, foo is whale",
         user: "what exactly is foo?",
@@ -86,6 +89,7 @@ describe("Test Qstash chat", () => {
     "should stream create",
     async () => {
       const response = await client.chat().create({
+        provider: "upstash",
         model: "meta-llama/Meta-Llama-3-8B-Instruct",
         messages: [
           {
@@ -162,3 +166,99 @@ describe("Test Qstash chat", () => {
     expect(result.messageId).toBeTruthy();
   });
 });
+
+describe("Test Qstash chat with third party LLMs", () => {
+  const client = new Client({ token: process.env.QSTASH_TOKEN! });
+
+  test(
+    "should respond to prompt",
+    async () => {
+      const response = await client.chat().prompt({
+        provider: "openai",
+        model: "gpt-3.5-turbo",
+        llmToken: process.env.OPENAI_API_KEY!,
+        system: "from now on, foo is whale",
+        user: "what exactly is foo?",
+        temperature: 0.5,
+      });
+
+      expect(response instanceof ReadableStream).toBeFalse();
+      expect(response.choices.length).toBe(1);
+      expect(response.choices[0].message.content.includes("whale")).toBeTrue();
+      expect(response.choices[0].message.role).toBe("assistant");
+    },
+    { timeout: 30_000, retry: 3 }
+  );
+
+  test(
+    "should respond to create",
+    async () => {
+      const response = await client.chat().create({
+        provider: "openai",
+        model: "gpt-3.5-turbo",
+        llmToken: process.env.OPENAI_API_KEY!,
+        messages: [
+          {
+            role: "system",
+            content: "from now on, foo is whale",
+          },
+          {
+            role: "user",
+            content: "what exactly is foo?",
+          },
+        ],
+        temperature: 0.5,
+      });
+
+      expect(response instanceof ReadableStream).toBeFalse();
+      expect(response.choices.length).toBe(1);
+      expect(response.choices[0].message.content.includes("whale")).toBeTrue();
+      expect(response.choices[0].message.role).toBe("assistant");
+    },
+    { timeout: 30_000, retry: 3 }
+  );
+
+  test(
+    "should stream prompt",
+    async () => {
+      const response = await client.chat().prompt({
+        provider: "openai",
+        model: "gpt-3.5-turbo",
+        llmToken: process.env.OPENAI_API_KEY!,
+        system: "from now on, foo is whale",
+        user: "what exactly is foo?",
+        stream: true,
+        temperature: 0.5,
+      });
+
+      await checkStream(response, ["whale"]);
+    },
+    { timeout: 30_000, retry: 3 }
+  );
+
+  test(
+    "should stream create",
+    async () => {
+      const response = await client.chat().create({
+        provider: "openai",
+        model: "gpt-3.5-turbo",
+        llmToken: process.env.OPENAI_API_KEY!,
+        messages: [
+          {
+            role: "system",
+            content: "from now on, foo is whale",
+          },
+          {
+            role: "user",
+            content: "what exactly is foo?",
+          },
+        ],
+        stream: true,
+        temperature: 0.5,
+      });
+
+      await checkStream(response, ["whale"]);
+    },
+    { timeout: 30_000, retry: 3 }
+  );
+});

src/client/llm/chat.ts

@@ -1,8 +1,9 @@
 import type { Requester } from "../http";
+import { PROVIDER_MAP } from "./constants";
 import type {
   ChatRequest,
   ChatCompletion,
-  PromptRequest,
+  PromptChatRequest,
   ChatCompletionMessage,
   StreamParameter,
   StreamEnabled,
@@ -17,15 +18,15 @@ export class Chat {
   }
 
   private static toChatRequest<TStream extends StreamParameter>(
-    request: PromptRequest<TStream>
+    request: PromptChatRequest<TStream>
   ): ChatRequest<TStream> {
     const messages: ChatCompletionMessage[] = [];
+
     messages.push(
       { role: "system", content: request.system },
       { role: "user", content: request.user }
     );
 
-    // @ts-expect-error ts can't resolve the type
     const chatRequest: ChatRequest<TStream> = { ...request, messages };
     return chatRequest;
   }
@@ -44,6 +45,9 @@ export class Chat {
   ): Promise<
     TStream extends StreamEnabled ? AsyncIterable<ChatCompletionChunk> : ChatCompletion
   > => {
+    if (request.provider === "openai" || request.provider === "togetherai")
+      return this.createThirdParty<TStream>(request);
+
     const body = JSON.stringify(request);
 
     if ("stream" in request && request.stream) {
@@ -69,6 +73,64 @@ export class Chat {
     });
   };
 
+  /**
+   * Calls the Upstash completions api given a ChatRequest.
+   *
+   * Returns a ChatCompletion or a stream of ChatCompletionChunks
+   * if stream is enabled.
+   *
+   * @param request ChatRequest with messages
+   * @returns Chat completion or stream
+   */
+  private createThirdParty = async <TStream extends StreamParameter>(
+    request: ChatRequest<TStream>
+  ): Promise<
+    TStream extends StreamEnabled ? AsyncIterable<ChatCompletionChunk> : ChatCompletion
+  > => {
+    if (request.provider === "openai" || request.provider === "togetherai") {
+      const baseUrl = PROVIDER_MAP[request.provider];
+
+      const llmToken = request.llmToken;
+      //@ts-expect-error We need to delete the prop, otherwise openai throws an error
+      delete request.llmToken;
+      //@ts-expect-error We need to delete the prop, otherwise openai throws an error
+      delete request.system;
+      //@ts-expect-error We need to delete the prop, otherwise openai throws an error
+      delete request.provider;
+
+      const body = JSON.stringify(request);
+
+      if ("stream" in request && request.stream) {
+        // @ts-expect-error when req.stream, we return ChatCompletion
+        return this.http.requestStream({
+          path: ["v1", "chat", "completions"],
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+            Connection: "keep-alive",
+            Accept: "text/event-stream",
+            "Cache-Control": "no-cache",
+            Authorization: `Bearer ${llmToken}`,
+          },
+          body,
+          baseUrl,
+        });
+      }
+
+      return this.http.request({
+        path: ["v1", "chat", "completions"],
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          Authorization: `Bearer ${llmToken}`,
+        },
+        body,
+        baseUrl,
+      });
+    }
+    throw new Error("Could not find any third party provider");
+  };
+
   /**
    * Calls the Upstash completions api given a PromptRequest.
    *
@@ -81,7 +143,7 @@ export class Chat {
    * @returns Chat completion or stream
    */
   prompt = async <TStream extends StreamParameter>(
-    request: PromptRequest<TStream>
+    request: PromptChatRequest<TStream>
   ): Promise<
     TStream extends StreamEnabled ? AsyncIterable<ChatCompletionChunk> : ChatCompletion
   > => {

src/client/llm/constants.ts

@@ -0,0 +1,6 @@
+import type { LlmProviderBaseUrl, LlmProvider } from "./types";
+
+export const PROVIDER_MAP: Record<LlmProvider, LlmProviderBaseUrl> = {
+  openai: "https://api.openai.com",
+  togetherai: "https://api.together.xyz",
+};