add ollama support (#269)

* add ollama support

---------

Co-authored-by: di-sukharev <dim.sukharev@gmail.com>
Co-authored-by: GPT10 <57486732+di-sukharev@users.noreply.github.com>
Co-authored-by: Jaroslaw Weber <jaroslaw.weber@adriel.com>
This commit is contained in:
Jaroslaw Weber
2024-02-27 23:01:13 +09:00
committed by GitHub
parent 57d9cc59b5
commit 1d6980faf3
14 changed files with 3028 additions and 3185 deletions
+109
View File
@@ -0,0 +1,109 @@
import axios from 'axios';
import chalk from 'chalk';
import { execa } from 'execa';
import {
ChatCompletionRequestMessage,
Configuration as OpenAiApiConfiguration,
OpenAIApi
} from 'openai';
import { intro, outro } from '@clack/prompts';
import {
CONFIG_MODES,
DEFAULT_MODEL_TOKEN_LIMIT,
getConfig
} from '../commands/config';
import { GenerateCommitMessageErrorEnum } from '../generateCommitMessageFromGitDiff';
import { tokenCount } from '../utils/tokenCount';
import { AiEngine } from './Engine';
const config = getConfig();
let maxTokens = config?.OCO_OPENAI_MAX_TOKENS;
let basePath = config?.OCO_OPENAI_BASE_PATH;
let apiKey = config?.OCO_OPENAI_API_KEY
const [command, mode] = process.argv.slice(2);
const isLocalModel = config?.OCO_AI_PROVIDER == 'ollama'
if (!apiKey && command !== 'config' && mode !== CONFIG_MODES.set && !isLocalModel) {
intro('opencommit');
outro(
'OCO_OPENAI_API_KEY is not set, please run `oco config set OCO_OPENAI_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);
}
const MODEL = config?.OCO_MODEL || 'gpt-3.5-turbo';
class OpenAi implements AiEngine {
private openAiApiConfiguration = new OpenAiApiConfiguration({
apiKey: apiKey
});
private openAI!: OpenAIApi;
constructor() {
if (basePath) {
this.openAiApiConfiguration.basePath = basePath;
}
this.openAI = new OpenAIApi(this.openAiApiConfiguration);
}
public generateCommitMessage = async (
messages: Array<ChatCompletionRequestMessage>
): Promise<string | undefined> => {
const params = {
model: MODEL,
messages,
temperature: 0,
top_p: 0.1,
max_tokens: maxTokens || 500
};
try {
const REQUEST_TOKENS = messages
.map((msg) => tokenCount(msg.content) + 4)
.reduce((a, b) => a + b, 0);
if (REQUEST_TOKENS > DEFAULT_MODEL_TOKEN_LIMIT - maxTokens) {
throw new Error(GenerateCommitMessageErrorEnum.tooMuchTokens);
}
const { data } = await this.openAI.createChatCompletion(params);
const message = data.choices[0].message;
return message?.content;
} catch (error) {
outro(`${chalk.red('✖')} ${JSON.stringify(params)}`);
const err = error as Error;
outro(`${chalk.red('✖')} ${err?.message || err}`);
if (
axios.isAxiosError<{ error?: { message: string } }>(error) &&
error.response?.status === 401
) {
const openAiError = error.response.data.error;
if (openAiError?.message) outro(openAiError.message);
outro(
'For help look into README https://github.com/di-sukharev/opencommit#setup'
);
}
throw err;
}
};
}
export const api = new OpenAi();