forked from di-sukharev/opencommit
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathgemini.ts
More file actions
133 lines (108 loc) · 4.23 KB
/
Copy pathgemini.ts
File metadata and controls
133 lines (108 loc) · 4.23 KB
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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import { ChatCompletionRequestMessage } from 'openai';
import { AiEngine } from './Engine';
import { Content, GenerativeModel, GoogleGenerativeAI, HarmBlockThreshold, HarmCategory, Part } from '@google/generative-ai';
import { CONFIG_MODES, ConfigType, DEFAULT_TOKEN_LIMITS, getConfig, MODEL_LIST } from '../commands/config';
import { intro, outro } from '@clack/prompts';
import chalk from 'chalk';
import axios from 'axios';
export class Gemini implements AiEngine {
private readonly config: ConfigType;
private readonly googleGenerativeAi: GoogleGenerativeAI;
private ai: GenerativeModel;
// vars
private maxTokens = {
input: DEFAULT_TOKEN_LIMITS.DEFAULT_MAX_TOKENS_INPUT,
output: DEFAULT_TOKEN_LIMITS.DEFAULT_MAX_TOKENS_OUTPUT
};
private basePath: string;
private apiKey: string;
private model: string;
constructor() {
this.config = getConfig() as ConfigType;
this.googleGenerativeAi = new GoogleGenerativeAI(this.config.OCO_GEMINI_API_KEY);
this.warmup();
}
async generateCommitMessage(messages: ChatCompletionRequestMessage[]): Promise<string | undefined> {
const systemInstruction = messages.filter(m => m.role === 'system')
.map(m => m.content)
.join('\n');
this.ai = this.googleGenerativeAi.getGenerativeModel({
model: this.model,
systemInstruction,
});
const contents = messages.filter(m => m.role !== 'system')
.map(m => ({ parts: [{ text: m.content } as Part], role: m.role == 'user' ? m.role : 'model', } as Content));
try {
const result = await this.ai.generateContent({
contents,
safetySettings: [
{
category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
},
{
category: HarmCategory.HARM_CATEGORY_HARASSMENT,
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
},
{
category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
},
{
category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
},
],
generationConfig: {
maxOutputTokens: this.maxTokens.output,
temperature: 0,
topP: 0.1,
},
});
return result.response.text();
} catch (error) {
const err = error as Error;
outro(`${chalk.red('✖')} ${err?.message || err}`);
if (
axios.isAxiosError<{ error?: { message: string } }>(error) &&
error.response?.status === 401
) {
const geminiError = error.response.data.error;
if (geminiError?.message) outro(geminiError.message);
outro(
'For help look into README https://github.com/di-sukharev/opencommit#setup'
);
}
throw err;
}
}
private warmup(): void {
if (this.config.OCO_TOKENS_MAX_INPUT !== undefined) this.maxTokens.input = this.config.OCO_TOKENS_MAX_INPUT;
if (this.config.OCO_TOKENS_MAX_OUTPUT !== undefined) this.maxTokens.output = this.config.OCO_TOKENS_MAX_OUTPUT;
this.basePath = this.config.OCO_GEMINI_BASE_PATH;
this.apiKey = this.config.OCO_GEMINI_API_KEY;
const [command, mode] = process.argv.slice(2);
const provider = this.config.OCO_AI_PROVIDER;
if (provider === 'gemini' && !this.apiKey &&
command !== 'config' && mode !== 'set') {
intro('opencommit');
outro('OCO_GEMINI_API_KEY is not set, please run `oco config set OCO_GEMINI_API_KEY=<your token> . If you are using GPT, make sure you add payment details, so API works.');
outro(
'For help look into README https://github.com/di-sukharev/opencommit#setup'
);
process.exit(1);
}
this.model = this.config.OCO_MODEL || MODEL_LIST.gemini[0];
if (provider === 'gemini' &&
!MODEL_LIST.gemini.includes(this.model) &&
command !== 'config' &&
mode !== CONFIG_MODES.set) {
outro(
`${chalk.red('✖')} Unsupported model ${this.model} for Gemini. Supported models are: ${MODEL_LIST.gemini.join(
', '
)}`
);
process.exit(1);
}
}
}