feat: use streaming in playground requests

Fixes #23

Signed-off-by: Jeff MAURY <[email protected]>

packages/backend/package.json

@@ -38,6 +38,7 @@
     "watch": "vite --mode development build -w"
   },
   "dependencies": {
+    "openai": "^4.26.0",
     "simple-git": "^3.22.0"
   },
   "devDependencies": {

packages/backend/src/managers/playground.ts

@@ -33,6 +33,9 @@ import type { QueryState } from '@shared/src/models/IPlaygroundQueryState';
 import { MSG_NEW_PLAYGROUND_QUERIES_STATE, MSG_PLAYGROUNDS_STATE_UPDATE } from '@shared/Messages';
 import type { PlaygroundState, PlaygroundStatus } from '@shared/src/models/IPlaygroundState';
 import type { ContainerRegistry } from '../registries/ContainerRegistry';
+import OpenAI from 'openai';
+import { Completion } from 'openai/resources';
+import { Stream } from 'openai/streaming';
 
 // TODO: this should not be hardcoded
 const PLAYGROUND_IMAGE = 'quay.io/bootsy/playground:v0';
@@ -207,54 +210,31 @@ export class PlayGroundManager {
       prompt: prompt,
     };
 
-    const post_data = JSON.stringify({
+    const client = new OpenAI({ baseURL: `http://localhost:${state.container.port}/v1` });
+
+    const response = await client.completions.create({
       model: modelInfo.file,
-      prompt: prompt,
+      prompt,
       temperature: 0.7,
       max_tokens: 1024,
+      stream: true,
     });
 
-    const post_options: http.RequestOptions = {
-      host: 'localhost',
-      port: '' + state.container.port,
-      path: '/v1/completions',
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/json',
-      },
-    };
+    this.queries.set(query.id, query);
+    this.sendQueriesState();
 
-    const post_req = http.request(post_options, res => {
-      res.setEncoding('utf8');
-      const chunks = [];
-      res.on('data', data => chunks.push(data));
-      res.on('end', () => {
-        const resBody = chunks.join();
-        if (res.headers['content-type'] === 'application/json') {
-          const result = JSON.parse(resBody);
-          const q = this.queries.get(query.id);
-          if (!q) {
-            throw new Error('query not found in state');
-          }
-          q.error = undefined;
-          q.response = result as ModelResponse;
-          this.queries.set(query.id, q);
-          this.sendQueriesState();
+    (async () => {
+      for await (const chunk of response) {
+        query.error = undefined;
+        if (query.response) {
+          query.response.choices = query.response.choices.concat(chunk.choices);
+        } else {
+          query.response = chunk;
         }
-      });
-    });
-    // post the data
-    post_req.write(post_data);
-    post_req.end();
-    post_req.on('error', error => {
-      console.error('connection on error.', error);
-      const q = this.queries.get(query.id);
-      q.error = `Something went wrong while trying to request model.${String(error)}`;
-      this.sendQueriesState();
-    });
+        this.sendQueriesState();
+      }
+    })().catch((err: unknown) => console.warn(`Error while reading streamed response for model ${modelInfo.id}`, err));
 
-    this.queries.set(query.id, query);
-    this.sendQueriesState();
     return query.id;
   }
 

packages/frontend/src/pages/ModelPlayground.svelte

@@ -1,6 +1,5 @@
 <script lang="ts">
   import type { ModelInfo } from '@shared/src/models/IModelInfo';
-  import type { ModelResponseChoice } from '@shared/src/models/IModelResponse';
   import Button from '../lib/button/Button.svelte';
   import { onMount } from 'svelte';
   import { studioClient } from '../utils/client';
@@ -10,13 +9,12 @@
   import type { PlaygroundState } from '@shared/src/models/IPlaygroundState';
   import Card from '/@/lib/Card.svelte';
   export let model: ModelInfo | undefined;
-  import Fa from 'svelte-fa';
   import { faPlay, faStop, faInfo, faWarning } from '@fortawesome/free-solid-svg-icons';
   import ErrorMessage from '/@/lib/ErrorMessage.svelte';
 
   let prompt = '';
   let queryId: number;
-  let result: ModelResponseChoice | undefined = undefined;
+  $: result = '';
   let inProgress = false;
   let error: string | undefined = undefined;
   let playgroundState: PlaygroundState | undefined = undefined;
@@ -41,7 +39,6 @@
     });
 
     const unsubscribeStates = playgroundStates.subscribe((states: PlaygroundState[]) => {
-      console.log('playgroundStates update', states);
       playgroundState = states.find((state) => state.modelId === model.id);
       if(playgroundState === undefined) {
         playgroundState = { modelId: model.id,  status: 'none' };
@@ -62,10 +59,10 @@
     }
 
     if (query.response) {
-      inProgress = false;
+      inProgress = !query.response.choices.some(choice => choice.finish_reason);
       prompt = query.prompt;
       if (query.response?.choices.length) {
-        result = query.response?.choices[0];
+        result = query.response?.choices.map(choice => choice.text).join('');
       }
     } else {
       inProgress = true;
@@ -79,7 +76,7 @@
       return;
     }
     inProgress = true;
-    result = undefined;
+    result = '';
     error = undefined;
 
     // do not display anything before we get a response from askPlayground
@@ -194,7 +191,7 @@
       readonly
       disabled
       rows="20"
-      bind:value={result.text}
+      bind:value={result}
       class="w-full p-2 outline-none text-sm bg-charcoal-800 rounded-sm text-gray-700 placeholder-gray-700"></textarea>
   {/if}
 </div>

packages/shared/src/models/IModelResponse.ts

@@ -4,7 +4,7 @@ export interface ModelResponse {
   id: string;
   model: string;
   choices: ModelResponseChoice[];
-  usage: ModelResponseUsage;
+  usage?: ModelResponseUsage;
 }
 
 export interface ModelResponseChoice {