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>

src/commands/config.ts

@@ -22,7 +22,8 @@ export enum CONFIG_KEYS {
   OCO_MODEL = 'OCO_MODEL',
   OCO_LANGUAGE = 'OCO_LANGUAGE',
   OCO_MESSAGE_TEMPLATE_PLACEHOLDER = 'OCO_MESSAGE_TEMPLATE_PLACEHOLDER',
-  OCO_PROMPT_MODULE = 'OCO_PROMPT_MODULE'
+  OCO_PROMPT_MODULE = 'OCO_PROMPT_MODULE',
+  OCO_AI_PROVIDER = 'OCO_AI_PROVIDER',
 }
 
 export const DEFAULT_MODEL_TOKEN_LIMIT = 4096;
@@ -48,7 +49,8 @@ const validateConfig = (
 
 export const configValidators = {
   [CONFIG_KEYS.OCO_OPENAI_API_KEY](value: any, config: any = {}) {
-    validateConfig(CONFIG_KEYS.OCO_OPENAI_API_KEY, value, 'Cannot be empty');
+    //need api key unless running locally with ollama
+    validateConfig('API_KEY', value || config.OCO_AI_PROVIDER == 'ollama', 'You need to provide an API key');
     validateConfig(
       CONFIG_KEYS.OCO_OPENAI_API_KEY,
       value.startsWith('sk-'),
@@ -150,7 +152,20 @@ export const configValidators = {
     );
 
     return value;
-  }
+  },
+
+  [CONFIG_KEYS.OCO_AI_PROVIDER](value: any) {
+    validateConfig(
+      CONFIG_KEYS.OCO_AI_PROVIDER,
+      [
+        '',
+        'openai',
+        'ollama'
+      ].includes(value),
+      `${value} is not supported yet, use 'ollama' or 'openai' (default)`
+    );
+    return value;
+  },
 };
 
 export type ConfigType = {
@@ -172,7 +187,8 @@ export const getConfig = (): ConfigType | null => {
     OCO_LANGUAGE: process.env.OCO_LANGUAGE || 'en',
     OCO_MESSAGE_TEMPLATE_PLACEHOLDER:
       process.env.OCO_MESSAGE_TEMPLATE_PLACEHOLDER || '$msg',
-    OCO_PROMPT_MODULE: process.env.OCO_PROMPT_MODULE || 'conventional-commit'
+    OCO_PROMPT_MODULE: process.env.OCO_PROMPT_MODULE || 'conventional-commit',
+    OCO_AI_PROVIDER: process.env.OCO_AI_PROVIDER || 'openai'
   };
 
   const configExists = existsSync(configPath);

src/engine/Engine.ts

@@ -0,0 +1,7 @@
+import { ChatCompletionRequestMessage } from 'openai';
+
+export interface AiEngine {
+  generateCommitMessage(
+    messages: Array<ChatCompletionRequestMessage>
+  ): Promise<string | undefined>;
+}

src/engine/ollama.ts

@@ -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();

src/engine/openAi.ts

@@ -13,23 +13,27 @@ import {
   CONFIG_MODES,
   DEFAULT_MODEL_TOKEN_LIMIT,
   getConfig
-} from './commands/config';
-import { GenerateCommitMessageErrorEnum } from './generateCommitMessageFromGitDiff';
-import { tokenCount } from './utils/tokenCount';
+} 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;
+let apiKey = config?.OCO_OPENAI_API_KEY
 
 const [command, mode] = process.argv.slice(2);
 
-if (!apiKey && command !== 'config' && mode !== CONFIG_MODES.set) {
+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>. Make sure you add payment details, so API works.`'
+    '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'
@@ -40,7 +44,7 @@ if (!apiKey && command !== 'config' && mode !== CONFIG_MODES.set) {
 
 const MODEL = config?.OCO_MODEL || 'gpt-3.5-turbo';
 
-class OpenAi {
+class OpenAi implements AiEngine {
   private openAiApiConfiguration = new OpenAiApiConfiguration({
     apiKey: apiKey
   });
@@ -100,16 +104,6 @@ class OpenAi {
   };
 }
 
-export const getOpenCommitLatestVersion = async (): Promise<
-  string | undefined
-> => {
-  try {
-    const { stdout } = await execa('npm', ['view', 'opencommit', 'version']);
-    return stdout;
-  } catch (_) {
-    outro('Error while getting the latest version of opencommit');
-    return undefined;
-  }
-};
+
 
 export const api = new OpenAi();

src/generateCommitMessageFromGitDiff.ts

@@ -3,11 +3,11 @@ import {
   ChatCompletionRequestMessageRoleEnum
 } from 'openai';
 
-import { api } from './api';
 import { DEFAULT_MODEL_TOKEN_LIMIT, getConfig } from './commands/config';
 import { getMainCommitPrompt } from './prompts';
 import { mergeDiffs } from './utils/mergeDiffs';
 import { tokenCount } from './utils/tokenCount';
+import { getEngine } from './utils/engine';
 
 const config = getConfig();
 
@@ -67,7 +67,8 @@ export const generateCommitMessageByDiff = async (
 
     const messages = await generateCommitMessageChatCompletionPrompt(diff);
 
-    const commitMessage = await api.generateCommitMessage(messages);
+    const engine = getEngine()
+    const commitMessage = await engine.generateCommitMessage(messages);
 
     if (!commitMessage)
       throw new Error(GenerateCommitMessageErrorEnum.emptyMessage);
@@ -104,13 +105,14 @@ function getMessagesPromisesByChangesInFile(
     }
   }
 
+  const engine = getEngine()
   const commitMsgsFromFileLineDiffs = lineDiffsWithHeader.map(
     async (lineDiff) => {
       const messages = await generateCommitMessageChatCompletionPrompt(
         separator + lineDiff
       );
 
-      return api.generateCommitMessage(messages);
+      return engine.generateCommitMessage(messages);
     }
   );
 
@@ -177,7 +179,8 @@ export const getCommitMsgsPromisesFromFileDiffs = async (
         separator + fileDiff
       );
 
-      commitMessagePromises.push(api.generateCommitMessage(messages));
+      const engine = getEngine()
+      commitMessagePromises.push(engine.generateCommitMessage(messages));
     }
   }
 

src/modules/commitlint/config.ts

@@ -1,6 +1,5 @@
 import { spinner } from '@clack/prompts';
 
-import { api } from '../../api';
 import { getConfig } from '../../commands/config';
 import { i18n, I18nLocals } from '../../i18n';
 import { COMMITLINT_LLM_CONFIG_PATH } from './constants';
@@ -9,6 +8,7 @@ import { commitlintPrompts, inferPromptsFromCommitlintConfig } from './prompts';
 import { getCommitLintPWDConfig } from './pwd-commitlint';
 import { CommitlintLLMConfig } from './types';
 import * as utils from './utils';
+import { getEngine } from '../../utils/engine';
 
 const config = getConfig();
 const translation = i18n[(config?.OCO_LANGUAGE as I18nLocals) || 'en'];
@@ -55,8 +55,9 @@ export const configureCommitlintIntegration = async (force = false) => {
   //   consistencyPrompts.map((p) => p.content)
   // );
 
+  const engine = getEngine()
   let consistency =
-    (await api.generateCommitMessage(consistencyPrompts)) || '{}';
+    (await engine.generateCommitMessage(consistencyPrompts)) || '{}';
 
   // Cleanup the consistency answer. Sometimes 'gpt-3.5-turbo' sends rule's back.
   prompts.forEach((prompt) => (consistency = consistency.replace(prompt, '')));

src/utils/checkIsLatestVersion.ts

@@ -3,7 +3,7 @@ import chalk from 'chalk';
 import { outro } from '@clack/prompts';
 
 import currentPackage from '../../package.json';
-import { getOpenCommitLatestVersion } from '../api';
+import { getOpenCommitLatestVersion } from '../version';
 
 export const checkIsLatestVersion = async () => {
   const latestVersion = await getOpenCommitLatestVersion();

src/utils/engine.ts

@@ -0,0 +1,13 @@
+import { AiEngine } from '../engine/Engine';
+import { api } from '../engine/openAi';
+import { getConfig } from '../commands/config';
+import { ollamaAi } from '../engine/ollama';
+
+export function getEngine(): AiEngine {
+  const config = getConfig();
+  if (config?.OCO_AI_PROVIDER == 'ollama') {
+    return ollamaAi;
+  }
+  //open ai gpt by default
+  return api;
+}

src/version.ts

@@ -0,0 +1,14 @@
+import { outro } from "@clack/prompts";
+import { execa } from "execa";
+
+export const getOpenCommitLatestVersion = async (): Promise<
+  string | undefined
+> => {
+  try {
+    const { stdout } = await execa('npm', ['view', 'opencommit', 'version']);
+    return stdout;
+  } catch (_) {
+    outro('Error while getting the latest version of opencommit');
+    return undefined;
+  }
+};