✨ Support stream chat completion

README.md

@@ -29,8 +29,11 @@ pnpm i kimichat-js
 
 ## 使用方法
 
+> 请注意, 流式调用需要调用方传入自定义的回调方法, 用于处理数据流
+
 ```typescript
 import { KimiChat } from "kimichat-js";
+import { PassThrough } from "stream";
 
 const kimi = new KimiChat("Your API Key");
 
@@ -43,13 +46,52 @@ const { data: messages } = await kimi.chatCompletions({
     },
   ],
 });
+
+// Stream Chat Completion example
+kimi.streamChatCompletions({
+  messages: [
+    {
+      role: "user",
+      content: "hello",
+    },
+  ],
+  callback: () => {
+    const bufs = [] as Buffer[];
+    const pt = new PassThrough()
+      .on("error", (err: any) => {
+        // Handling error
+      })
+      .on("data", (buf: Buffer) => {
+        // Handling buf data
+        bufs.push(buf);
+      })
+      .on("end", () => {
+        // Handle buf data on end
+        Buffer.concat(bufs).toString("utf8");
+      });
+
+    return pt;
+  },
+});
 ```
 
-> 更多方法请参考 examples
+对于 stream 格式返回数据格式如下, 类型引用 `StreamChatCompletionData`
+
+```json
+data: {"id":"cmpl-1305b94c570f447fbde3180560736287","object":"chat.completion.chunk","created":1698999575,"model":"moonshot-v1-8k","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}
+
+data: {"id":"cmpl-1305b94c570f447fbde3180560736287","object":"chat.completion.chunk","created":1698999575,"model":"moonshot-v1-8k","choices":[{"index":0,"delta":{"content":"你好"},"finish_reason":null}]}
+
+...
 
-## Todos
+data: {"id":"cmpl-1305b94c570f447fbde3180560736287","object":"chat.completion.chunk","created":1698999575,"model":"moonshot-v1-8k","choices":[{"index":0,"delta":{"content":"。"},"finish_reason":null}]}
 
-[ ] Chat Completion 支持 stream 请求
+data: {"id":"cmpl-1305b94c570f447fbde3180560736287","object":"chat.completion.chunk","created":1698999575,"model":"moonshot-v1-8k","choices":[{"index":0,"delta":{},"finish_reason":"stop","usage":{"prompt_tokens":19,"completion_tokens":13,"total_tokens":32}}]}
+
+data: [DONE]
+```
+
+> 更多方法请参考 examples
 
 ## 引用
 

examples/examples.md

@@ -1,5 +1,6 @@
 ```typescript
 import { KimiChat } from "../lib";
+import { PassThrough } from "stream";
 
 const kimi = new KimiChat("Your API Key");
 
@@ -13,6 +14,34 @@ const { data: messages } = await kimi.chatCompletions({
   ],
 });
 
+// Stream Chat Completion
+kimi.streamChatCompletions({
+  messages: [
+    {
+      role: "user",
+      content: "hello",
+    },
+  ],
+  callback: () => {
+    const bufs = [] as Buffer[];
+    const pt = new PassThrough()
+      .on("error", (err: any) => {
+        // @ts-ignore
+      })
+      .on("data", (buf: Buffer) => {
+        console.log(buf.toString());
+        // data: {"id":"cmpl-1305b94c570f447fbde3180560736287","object":"chat.completion.chunk","created":1698999575,"model":"moonshot-v1-8k","choices":[{"index":0,"delta":{"content":"你好"},"finish_reason":null}]}
+
+        bufs.push(buf);
+      })
+      .on("end", () => {
+        console.log(Buffer.concat(bufs).toString("utf8")).toBeTruthy();
+      });
+
+    return pt;
+  },
+});
+
 // Estimate Token
 const { data: tokenCounts } = await kimi.estimateTokens({
   messages: [

lib/__tests__/kimi.spec.ts

@@ -3,6 +3,7 @@ dotenv.config()
 
 import fs from "fs"
 import { KimiChat } from "..";
+import { PassThrough } from "stream";
 
 const clearFiles = async () => {
   // @ts-ignore
@@ -23,6 +24,9 @@ describe('Test KimiChat SDK', () => {
     tmpDir = fs.mkdtempSync("/tmp/kimichat")
     tepFileName = `${tmpDir}/temp.md`
     fs.writeFileSync(tepFileName, "Hello kimi")
+
+    // @ts-ignore
+    kimi = new KimiChat(process.env.MOONSHOT_API_KEY)
   })
 
   afterAll(() => {
@@ -36,9 +40,6 @@ describe('Test KimiChat SDK', () => {
   })
 
   test("Constructor", () => {
-    // @ts-ignore
-    kimi = new KimiChat(process.env.MOONSHOT_API_KEY)
-
     expect(kimi).toBeInstanceOf(KimiChat)
   })
 
@@ -70,6 +71,31 @@ describe('Test KimiChat SDK', () => {
     })
   })
 
+  test("streamChatCompletions", async () => {
+    await kimi.streamChatCompletions({
+      messages:[{
+        role: "user",
+        content: "hello"
+      }],
+      callback: () => {
+        const bufs = [] as Buffer[]
+        const pt = new PassThrough()
+          .on('error', (err: any) => {
+            // @ts-ignore
+          })
+          .on('data', (buf: Buffer) => {
+            expect(buf).toBeTruthy()
+            bufs.push(buf)
+          })
+          .on('end', () => {
+            expect(Buffer.concat(bufs).toString('utf8')).toBeTruthy()
+          })
+
+        return pt
+      }
+    })
+  })
+
   test("estimateTokens", async () => {
     const res = await kimi.estimateTokens({
       messages: [{

lib/index.ts

@@ -1,6 +1,7 @@
-import { Dispatcher, request } from 'undici';
+import { Dispatcher, request, stream } from 'undici';
 import FormData from 'form-data';
 import fs from 'fs';
+const { PassThrough } = require('node:stream')
 
 import { BASE_URL } from "./constants";
 import { ChatCompletion, Model, FileContent, Message, ModelResource, KimiFile, DeletedFile, Tokens, KimiResponse, IKimiChat } from "./types";
@@ -43,15 +44,13 @@ export class KimiChat implements IKimiChat {
     temperature = 0.3,
     topN = 1.0,
     n = 1,
-    stream = false
   }: {
     messages: Message[]
     model?: Model
     maxTokens?: number
     temperature?: number
     topN?: number
     n?: number
-    stream?: boolean
   }) {
     const res = await request(`${this.baseUrl}/chat/completions`, {
       method: "POST",
@@ -73,6 +72,41 @@ export class KimiChat implements IKimiChat {
     return this.handleResponse<ChatCompletion>(res)
   }
 
+  public async streamChatCompletions({
+    messages,
+    model = "moonshot-v1-8k",
+    maxTokens = 1024,
+    temperature = 0.3,
+    topN = 1.0,
+    n = 1,
+    callback
+  }: {
+    messages: Message[]
+    callback: Dispatcher.StreamFactory
+    model?: Model
+    maxTokens?: number
+    temperature?: number
+    topN?: number
+    n?: number
+  }) {
+    return stream(`${this.baseUrl}/chat/completions`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "Authorization": `Bearer ${this.apiKey}`
+      },
+      body: JSON.stringify({
+        messages,
+        model,
+        max_tokens: maxTokens,
+        temperature,
+        top_p: topN,
+        n,
+        stream: true,
+      })
+    }, callback)
+  }
+
   public async getFiles() {
     const res = await request(`${this.baseUrl}/files`, {
       method: "GET",

lib/types.ts

@@ -1,3 +1,5 @@
+import { Dispatcher } from "undici"
+
 export type Permission = {
   created: number
   id: string
@@ -100,6 +102,24 @@ export interface ErrorResponse {
   }
 }
 
+export type StreamChoice = {
+  index: number
+  delta: {
+    role?: string
+    content?: string
+  }
+  finish_reason?: string
+  usage?: Usage
+}
+
+export type StreamChatCompletionData = {
+  id: string;
+  object: string;
+  created: number;
+  model: string;
+  choices: StreamChoice[];
+}
+
 export interface IKimiChat {
   chatCompletions: (data: {
     messages: Message[];
@@ -108,8 +128,16 @@ export interface IKimiChat {
     temperature?: number;
     topN?: number;
     n?: number;
-    stream?: boolean;
   }) => Promise<KimiResponse<ChatCompletion>>;
+  streamChatCompletions: (data: {
+    messages: Message[];
+    callback: Dispatcher.StreamFactory
+    model?: Model;
+    maxTokens?: number;
+    temperature?: number;
+    topN?: number;
+    n?: number;
+  }) => Promise<Dispatcher.StreamData>
   getModels: () => Promise<KimiResponse<ModelResource[]>>;
   getFileContent: (id: string) => Promise<KimiResponse<FileContent>>;
   deleteFile: (id: string) => Promise<KimiResponse<DeletedFile>>;