1
/*
2
 *  This file is part of CoCalc: Copyright © 2020 Sagemath, Inc.
3
 *  License: MS-RSL – see LICENSE.md for details
4
 */
5

6
import { delay } from "awaiting";
7
import { EventEmitter } from "events";
8
import { redux } from "@cocalc/frontend/app-framework";
9
import {
10
  LanguageModel,
11
  LanguageServiceCore,
12
  getSystemPrompt,
13
  isFreeModel,
14
  model2service,
15
} from "@cocalc/util/db-schema/llm-utils";
16
import type { WebappClient } from "./client";
17
import type { History } from "./types";
18
import {
19
  LOCALIZATIONS,
20
  OTHER_SETTINGS_LOCALE_KEY,
21
  OTHER_SETTINGS_REPLY_ENGLISH_KEY,
22
} from "@cocalc/util/i18n/const";
23
import { sanitizeLocale } from "@cocalc/frontend/i18n";
24

25
interface QueryLLMProps {
26
  input: string;
27
  model: LanguageModel;
28
  system?: string;
29
  history?: History;
30
  project_id?: string;
31
  path?: string;
32
  chatStream?: ChatStream; // if given, uses chat stream
33
  tag?: string;
34
  startStreamExplicitly?: boolean;
35
}
36

37
export class LLMClient {
38
  private client: WebappClient;
39

40
  constructor(client: WebappClient) {
41
    this.client = client;
42
  }
43

44
  public async query(opts: QueryLLMProps): Promise<string> {
45
    return await this.queryLanguageModel(opts);
46
  }
47

48
  // ATTN/TODO: startExplicitly seems to be broken
49
  public queryStream(opts, startExplicitly = false): ChatStream {
50
    const chatStream = new ChatStream();
51
    (async () => {
52
      try {
53
        await this.queryLanguageModel({ ...opts, chatStream });
54
        if (!startExplicitly) {
55
          chatStream.emit("start");
56
        }
57
      } catch (err) {
58
        chatStream.emit("error", err);
59
      }
60
    })();
61
    return chatStream;
62
  }
63

64
  private async queryLanguageModel({
65
    input,
66
    model,
67
    system, // if not set, a default system prompt is used – disable by setting to ""
68
    history,
69
    project_id,
70
    path,
71
    chatStream,
72
    tag = "",
73
  }: QueryLLMProps): Promise<string> {
74
    system ??= getSystemPrompt(model, path);
75

76
    // remove all date entries from all history objects
77
    if (history != null) {
78
      for (const h of history) {
79
        delete h.date;
80
      }
81
    }
82

83
    if (!redux.getStore("projects").hasLanguageModelEnabled(project_id, tag)) {
84
      throw new Error(
85
        `Language model support is not currently enabled ${
86
          project_id ? "in this project" : "on this server"
87
        }. [tag=${tag}]`,
88
      );
89
    }
90

91
    input = input.trim();
92
    if (chatStream == null) {
93
      if (!input || input == "test") {
94
        return "Great! What can I assist you with today?";
95
      }
96
      if (input == "ping") {
97
        await delay(1000);
98
        return "Pong";
99
      }
100
    }
101

102
    // append a sentence to request to translate the output to the user's language – unless disabled
103
    const other_settings = redux.getStore("account").get("other_settings");
104
    const alwaysEnglish = !!other_settings.get(
105
      OTHER_SETTINGS_REPLY_ENGLISH_KEY,
106
    );
107
    const locale = sanitizeLocale(
108
      other_settings.get(OTHER_SETTINGS_LOCALE_KEY),
109
    );
110
    if (!alwaysEnglish && locale != "en") {
111
      const lang = LOCALIZATIONS[locale].name; // name is always in english
112
      system = `${system}\n\nYour answer must be written in the language ${lang}.`;
113
    }
114

115
    const is_cocalc_com = redux.getStore("customize").get("is_cocalc_com");
116

117
    if (!isFreeModel(model, is_cocalc_com)) {
118
      // Ollama and others are treated as "free"
119
      const service = model2service(model) as LanguageServiceCore;
120
      // when client gets non-free openai model request, check if allowed.  If not, show quota modal.
121
      const { allowed, reason } =
122
        await this.client.purchases_client.isPurchaseAllowed(service);
123

124
      if (!allowed) {
125
        await this.client.purchases_client.quotaModal({
126
          service,
127
          reason,
128
          allowed,
129
        });
130
      }
131
      // Now check again after modal dismissed...
132
      const x = await this.client.purchases_client.isPurchaseAllowed(service);
133
      if (!x.allowed) {
134
        throw Error(reason);
135
      }
136
    }
137

138
    // do not import until needed -- it is HUGE!
139
    const {
140
      numTokensUpperBound,
141
      truncateHistory,
142
      truncateMessage,
143
      getMaxTokens,
144
    } = await import("@cocalc/frontend/misc/llm");
145

146
    // We always leave some room for output:
147
    const maxTokens = getMaxTokens(model) - 1000;
148
    input = truncateMessage(input, maxTokens);
149
    const n = numTokensUpperBound(input, getMaxTokens(model));
150
    if (n >= maxTokens) {
151
      history = undefined;
152
    } else if (history != null) {
153
      history = truncateHistory(history, maxTokens - n, model);
154
    }
155
    // console.log("chatgpt", { input, system, history, project_id, path });
156
    const options = {
157
      input,
158
      system,
159
      project_id,
160
      path,
161
      history,
162
      model,
163
      tag: `app:${tag}`,
164
    };
165

166
    if (chatStream == null) {
167
      // not streaming
168
      return await this.client.conat_client.llm(options);
169
    }
170

171
    chatStream.once("start", async () => {
172
      // streaming version
173
      try {
174
        await this.client.conat_client.llm({
175
          ...options,
176
          stream: chatStream.process,
177
        });
178
      } catch (err) {
179
        chatStream.error(err);
180
      }
181
    });
182

183
    return "see stream for output";
184
  }
185
}
186

187
class ChatStream extends EventEmitter {
188
  constructor() {
189
    super();
190
  }
191

192
  process = (text: string | null) => {
193
    // emits undefined text when done (or err below)
194
    this.emit("token", text);
195
  };
196

197
  error = (err) => {
198
    this.emit("error", err);
199
  };
200
}
201

202
export type { ChatStream };
203

204