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
+7
View File
@@ -0,0 +1,7 @@
import { ChatCompletionRequestMessage } from 'openai';
export interface AiEngine {
generateCommitMessage(
messages: Array<ChatCompletionRequestMessage>
): Promise<string | undefined>;
}
+36
View File
@@ -0,0 +1,36 @@
import axios, { AxiosError } from 'axios';
import { ChatCompletionRequestMessage } from 'openai';
import { AiEngine } from './Engine';
export class OllamaAi implements AiEngine {
async generateCommitMessage(
messages: Array<ChatCompletionRequestMessage>
): Promise<string | undefined> {
const model = 'mistral'; // todo: allow other models
let prompt = messages.map((x) => x.content).join('\n');
//hoftix: local models are not so clever so im changing the prompt a bit...
prompt += 'Summarize above git diff in 10 words or less';
const url = 'http://localhost:11434/api/generate';
const p = {
model,
prompt,
stream: false
};
try {
const response = await axios.post(url, p, {
headers: {
'Content-Type': 'application/json'
}
});
const answer = response.data?.response;
return answer;
} catch (err: any) {
const message = err.response?.data?.error ?? err.message;
throw new Error('local model issues. details: ' + message);
}
}
}
export const ollamaAi = new OllamaAi();
+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();