Skip to content

Commit

Permalink
feat(completion): Add new language model "codegeex-4" for inline comp…
Browse files Browse the repository at this point in the history
…letion
  • Loading branch information
crazygo committed Jul 11, 2024
1 parent 3e8e2d2 commit 15a6c80
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 6 deletions.
1 change: 1 addition & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,7 @@
"description": "%configuration.completions.provider.description%",
"enum": [
"anthropic",
"codegeex-4",
"openai",
"qianfan",
"tongyi",
Expand Down
20 changes: 14 additions & 6 deletions src/base/common/language-models/providers/zhipuaiProvider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,18 @@ export class ZhipuAILanguageModelProvider implements ILanguageModelProvider {
progress?: Progress<IChatResponseFragment>,
token?: CancellationToken,
): Promise<string> {
// https://open.bigmodel.cn/dev/api#codegeex-4
// Using these 4 stops is enough
options.stop = ["<|endoftext|>", "<|user|>", "<|assistant|>", "<|observation|>"];
const { model, ...rest } = options;
// magic for completions prompt
if (messages.length === 1 && messages[0].content){
const partials = messages[0].content.match(/^<\|fim_prefix\|>(.*)<\|fim_suffix\|>(.*)<\|fim_middle\|>$/s);
if (partials) {
messages[0].content = `<|code_suffix|>${partials[2]}<|code_prefix|>${partials[1]}<|code_middle|>`;
}
}


const llm = this._newLLM(options);

Expand All @@ -33,7 +44,7 @@ export class ZhipuAILanguageModelProvider implements ILanguageModelProvider {

const completion = await llm.chat.completions.create(
{
...rest,
...rest ,
stream: true,
model: this._resolveChatModel(model),
messages: messages,
Expand Down Expand Up @@ -69,10 +80,6 @@ export class ZhipuAILanguageModelProvider implements ILanguageModelProvider {
progress?: Progress<IChatResponseFragment>,
token?: CancellationToken,
): Promise<string> {
if (this.configService.get('completions.enableLegacyMode')) {
return this._legacyCompletionResponse(prompt, options, progress, token);
}

return this.provideChatResponse([{ role: ChatMessageRole.User, content: prompt }], options, progress, token);
}

Expand Down Expand Up @@ -199,7 +206,8 @@ export class ZhipuAILanguageModelProvider implements ILanguageModelProvider {
}

private _resolveComletionModel(model?: string) {
return model || this._resolveChatModel(this.configService.get<string>('zhipuai.model'));
// codegeex-4 is the best choice for code completion :)
return 'codegeex-4';
}

private _resolveEmbeddingModel(model?: string) {
Expand Down

0 comments on commit 15a6c80

Please sign in to comment.