All files / api generate.ts

46.66% Statements 21/45
0% Branches 0/4
33.33% Functions 2/6
43.9% Lines 18/41

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 1161x 1x 1x                         1x 1x   1x                   1x         1x 1x 1x   1x   1x           1x 1x                                                                                                           1x           1x   1x     1x        
import * as vscode from 'vscode';
import { getGenerateUrl } from './getUrl';
import { showError } from '../lib/showError';
 
interface IParam {
  baseUrl: string;
  prompt: string;
  model: string;
  rules: { id: string; name: string; content: string; selected: boolean }[];
  onLoading: () => void;
  onStartStreaming: () => void;
  onFinishStreaming: () => void;
  onUpdateOutput: (params: { prompt: string; text: string; streamId: number }) => void;
}
 
let controller: AbortController | null = null;
let isStreaming = false;
 
export const startStreaming = async ({
  baseUrl,
  prompt,
  model,
  rules,
  onLoading,
  onStartStreaming,
  onFinishStreaming,
  onUpdateOutput,
}: IParam) => {
  Iif (isStreaming) {
    vscode.window.showWarningMessage('Streaming is already in progress');
    return;
  }
 
  const streamId = new Date().getTime();
  controller = new AbortController();
  isStreaming = true;
 
  onLoading();
 
  onUpdateOutput({
    prompt,
    text: '',
    streamId,
  });
 
  try {
    const response = await fetch(getGenerateUrl(baseUrl), {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        model,
        prompt:
          rules
            .filter(r => r.selected)
            .map(r => r.content)
            .join('\n') +
          '\n' +
          prompt,
        stream: true,
      }),
      signal: controller.signal,
    });
 
    Iif (!response.body) {
      throw new Error('Response body is null');
    }
 
    const reader = response.body.getReader();
    const decoder = new TextDecoder();
 
    onStartStreaming();
 
    while (true) {
      const { done, value } = await reader.read();
 
      Iif (done) {
        break;
      }
 
      const chunk = decoder.decode(value);
      const lines = chunk.split('\n').filter(line => line.trim());
 
      for (const line of lines) {
        onUpdateOutput({
          prompt,
          text: JSON.parse(line).response,
          streamId,
        });
      }
    }
  } catch (error: unknown) {
    showError(error);
  } finally {
    isStreaming = false;
    controller = null;
 
    onFinishStreaming();
  }
};
 
export const stopStreaming = () => {
  Iif (controller) {
    controller.abort();
  }
};
 
export const generate = (param: IParam) => {
  // Запускаем стриминг
  startStreaming(param);
 
  // Возвращаем объект с методами управления
  return {
    stop: stopStreaming,
  };
};