commit-maker 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

commit_maker/main.py

@@ -1,6 +1,4 @@
-# CLI-утилита, которая будет создавать сообщение для коммита на основе ИИ.
-# noqa: F841
-
+# CLI utility that generates commit messages using AI.
 import argparse
 import importlib
 import os
@@ -8,105 +6,43 @@ import subprocess
 
 import requests
 import rich.console
-import rich_argparse
 
-# Константы
+from .colored import colored
+from .custom_int_prompt import CustomIntPrompt
+from .cut_think_part import cut_think
+from .mistral import MistralAI
+from .ollama import Ollama
+from .rich_custom_formatter import CustomFormatter
+
+# Constants
 mistral_api_key = os.environ.get("MISTRAL_API_KEY")
 console = rich.console.Console()
+prompt = CustomIntPrompt()
+available_langs = ["en", "ru"]
 
-# Настройка вывода --help (это не хардкод, такой способ указан в оф.
-# документации rich_argparse)
-rich_argparse.RichHelpFormatter.styles = {
-    "argparse.args": "cyan bold",
-    "argparse.groups": "green bold",
-    "argparse.metavar": "dark_cyan",
-    "argparse.prog": "dark_green bold",
-}
-
-
-# Функции для цветного вывода
-def bold(text: str) -> str:
-    """Возвращает жирный текст
-
-    Args:
-        text (str): Текст
-
-    Returns:
-        str: Жирный текст
-    """
-    bold_start = "\033[1m"
-    bold_end = "\033[0m"
-    return f"{bold_start}{text}{bold_end}"
-
-
-def colored(
-    string: str,
-    color: str,
-    text_bold: bool = True,
-) -> str:
-    """Функция для 'окраски' строк для красивого вывода
-
-    Args:
-        string (str): Строка, которую нужно покрасить
-        color (str): Цвет покраски ['red', 'yellow', 'green', 'magenta',\
-            'blue', 'cyan', 'reset']
-        text_bold (bool, optional): Жирный текст или нет. Defaults to True.
-
-    Returns:
-        str: Покрашенная строка
-
-    Example:
-        `print(colored(string='Success!', color='green'))` # Выводит 'Success!'
-        зеленого цвета
-    """
-    COLOR_RED = "\033[31m"
-    COLOR_GREEN = "\033[32m"
-    COLOR_YELLOW = "\033[33m"
-    COLOR_BLUE = "\033[94m"
-    COLOR_MAGENTA = "\033[95m"
-    COLOR_CYAN = "\033[96m"
-    COLOR_RESET = "\033[0m"
-    COLORS_DICT = {
-        "red": COLOR_RED,
-        "green": COLOR_GREEN,
-        "yellow": COLOR_YELLOW,
-        "blue": COLOR_BLUE,
-        "magenta": COLOR_MAGENTA,
-        "cyan": COLOR_CYAN,
-        "reset": COLOR_RESET,
-    }
-    return (
-        bold(f"{COLORS_DICT[color]}{string}{COLORS_DICT['reset']}")
-        if text_bold
-        else f"{COLORS_DICT[color]}{string}{COLORS_DICT['reset']}"
-    )
-
-
-# Парсер параметров
+# Argument parser
 parser = argparse.ArgumentParser(
     prog="commit_maker",
-    description="CLI-утилита, которая будет создавать сообщение "
-    "для коммита на основе ИИ. Можно использовать локальные модели/Mistral AI "
-    "API. Локальные модели используют ollama.",
-    formatter_class=rich_argparse.RichHelpFormatter,
+    description="CLI utility that generates commit messages using AI. "
+    "Supports local models/Mistral AI API. Local models use ollama.",
+    formatter_class=CustomFormatter,
 )
 
-# Общие параметры
-general_params = parser.add_argument_group("Общие параметры")
+# General parameters
+general_params = parser.add_argument_group("General parameters")
 general_params.add_argument(
     "-l",
     "--local-models",
     action="store_true",
     default=False,
-    help="Запуск с использованием локальных моделей",
+    help="Use local models",
 )
 general_params.add_argument(
     "-d",
     "--dry-run",
     action="store_true",
     default=False,
-    help="Запуск с выводом сообщения на основе зайстейдженных "
-    "изменений, без создания коммита",
+    help="Dry run: show commit message without creating commit",
 )
 general_params.add_argument(
     "-V",
@@ -114,173 +50,65 @@ general_params.add_argument(
     action="version",
     version=f"%(prog)s {importlib.metadata.version('commit-maker')}",
 )
+general_params.add_argument(
+    "-o",
+    "--timeout",
+    type=int,
+    default=None,
+    help="Change timeout for models. Default is None.",
+)
 
-# Параметры генерации
-generation_params = parser.add_argument_group("Параметры генерации")
+# Generation parameters
+generation_params = parser.add_argument_group("Generation parameters")
 generation_params.add_argument(
     "-t",
     "--temperature",
     default=1.0,
     type=float,
-    help="Температура модели при создании месседжа.\
-        Находится на отрезке [0.0, 1.5]. Defaults to 1.0",
+    help="Model temperature for message generation. "
+         "Range: [0.0, 1.5]. Default: 1.0",
 )
 generation_params.add_argument(
     "-m",
     "--max-symbols",
     type=int,
     default=200,
-    help="Длина сообщения коммита. Defaults to 200",
+    help="Maximum commit message length. Default: 200",
 )
 generation_params.add_argument(
     "-M",
     "--model",
     type=str,
-    help="Модель, которую ollama будет использовать.",
+    help="Model to be used by ollama",
 )
 generation_params.add_argument(
     "-e",
     "--exclude",
     nargs="+",
     default=[],
-    help="Файлы, которые нужно игнорировать при создании сообщения коммита",
+    help="Files to exclude when generating commit message",
+)
+generation_params.add_argument(
+    "-w",
+    "--wish",
+    default=None,
+    type=str,
+    help="Custom wishes/edits for the commit message",
+)
+generation_params.add_argument(
+    "-L",
+    "--language",
+    choices=available_langs,
+    default="ru",
+    help="Language of generated commit message (en/ru)",
 )
 
 
-# Класс для использования API Mistral AI
-class MistralAI:
-    """Класс для общения с MistralAI.
-    Написан с помощью requests."""
-
-    def __init__(
-        self,
-        api_key: str,
-        model: str = "mistral-small-latest",
-    ):
-        """Инициализация класса
-
-        Args:
-            api_key (str): Апи ключ MistralAI
-        """
-        self.url = "https://api.mistral.ai/v1/chat/completions"
-        self.api_key = api_key
-        self.headers = {
-            "Content-Type": "application/json",
-            "Accept": "application/json",
-            "Authorization": f"Bearer {api_key}",
-        }
-        self.model = model
-
-    def message(
-        self,
-        message: str,
-        role: str = "user",
-        temperature: float = 0.7,
-    ) -> str:
-        """Функция сообщения
-
-        Args:
-            message (str): Сообщение
-            role (str, optional): Роль. Defaults to "user".
-
-        Returns:
-            str: Json-ответ/Err
-        """
-        data = {
-            "model": self.model,
-            "messages": [
-                {
-                    "role": role,
-                    "content": message,
-                }
-            ],
-            "temperature": 0.7,
-        }
-        try:
-            response = requests.post(
-                url=self.url,
-                json=data,
-                headers=self.headers,
-                timeout=60,
-            )
-            response.raise_for_status()
-            return response.json()["choices"][0]["message"]["content"]
-
-        except requests.exceptions.RequestException as e:
-            print(colored(f"Ошибка при обращении к Mistral AI: {e}", "red"))
-        except KeyError:
-            print(colored("Ошибка парсинга ответа от Mistral AI", "red"))
-
-
-# Класс для использования API Ollama
-class Ollama:
-    """Класс для общения с локальными моделями Ollama.
-    Написан с помощью requests."""
-
-    def __init__(
-        self,
-        model: str,
-    ):
-        """Инициализация класса"""
-        self.model = model
-        self.url = "http://localhost:11434/api/chat"
-        self.headers = {
-            "Content-Type": "application/json",
-            "Accept": "application/json",
-        }
-
-    def message(
-        self,
-        message: str,
-        temperature: float = 0.7,
-        role: str = "user",
-    ) -> str:
-        """Функция сообщения
-
-        Args:
-            message (str): Сообщение
-            model (str): Модель, с которой будем общаться
-            temperature (float, optional): Температура общения. Defaults to 0.7
-            role (str, optional): Роль в сообщении.
-
-        Returns:
-            str: Json-ответ/Err
-        """
-        data = {
-            "model": self.model,
-            "messages": [
-                {
-                    "role": role,
-                    "content": message,
-                }
-            ],
-            "options": {
-                "temperature": temperature,
-            },
-            "stream": False,
-        }
-
-        try:
-            response = requests.post(
-                url=self.url,
-                json=data,
-                headers=self.headers,
-                timeout=60,
-            )
-            response.raise_for_status()  # выбросит ошибку при плохом статусе
-            return response.json()["choices"][0]["message"]["content"]
-
-        except requests.exceptions.RequestException as e:
-            print(colored(f"Ошибка при обращении к Ollama: {e}", "red"))
-        except KeyError:
-            print(colored("Ошибка парсинга ответа от Ollama", "red"))
-
-
-# main функция
+# Main function
 
 
 def main() -> None:
-    # Парсинг аргументов
+    # Parsing arguments
     parsed_args = parser.parse_args()
     use_local_models = parsed_args.local_models
     max_symbols = parsed_args.max_symbols
@@ -288,24 +116,31 @@ def main() -> None:
     dry_run = parsed_args.dry_run
     temperature = parsed_args.temperature
     excluded_files = parsed_args.exclude
-
-    # Промпт для ИИ
-    prompt_for_ai = f"""Привет! Ты составитель коммитов для git.
-    Сгенерируй коммит-месседж на РУССКОМ языке, который:
-    1. Точно отражает суть изменений
-    2. Не превышает {max_symbols} символов
-    Опирайся на данные от 'git status' и 'git diff'.
-    В ответ на это сообщение тебе нужно предоставить
-    ТОЛЬКО коммит. Пиши просто обычный текст, без markdown!"""
+    wish = parsed_args.wish
+    timeout = parsed_args.timeout
+    lang = parsed_args.language
+
+    # AI prompt
+    prompt_for_ai = f"""You are a git commit message generator.
+    Generate a single commit message in
+    {"Russian" if lang == "ru" else "English"} that:
+    Clearly summarizes the purpose of the changes.
+    Does not exceed {max_symbols} characters.
+    Uses information from git status and git diff.
+    Takes into account user preferences: {wish}.
+    Output only the commit message — plain text, no markdown, no
+    explanations, no formatting."""
 
     try:
         if not use_local_models and not mistral_api_key:
-            print(
-                colored("Не найден MISTRAL_API_KEY для работы с API!", "red")
+            console.print(
+                "MISTRAL_API_KEY not found for API usage!",
+                style="red",
+                highlight=False,
             )
             return
 
-        # Получаем версию git, если он есть
+        # Get git version if available
         git_version = subprocess.run(  # noqa
             ["git", "--version"],
             capture_output=True,
@@ -313,98 +148,12 @@ def main() -> None:
             encoding="utf-8",
         ).stdout
 
-        if use_local_models:
-            # Получаем список моделей из Ollama, если Ollama есть
-            ollama_list_of_models = (
-                subprocess.run(
-                    ["ollama", "ls"],
-                    capture_output=True,
-                    text=True,
-                    encoding="utf-8",
-                )
-                .stdout.strip()
-                .split("\n")
-            )
-            ollama_list_of_models = [
-                i.split()[0] for i in ollama_list_of_models[1:]
-            ]
-        else:
-            ollama_list_of_models = 0
-
-        # Обработка отсутствия ollama
-        if not ollama_list_of_models and use_local_models:
-            print(
-                colored(
-                    "Ollama не установлена или список моделей пуст!", "yellow"
-                )
-                + " Для установки перейдите по https://ollama.com/download"
-            )
-            return None
-        elif not use_local_models and model:
-            print(
-                f"Для использования {model} локально используйте флаг "
-                + colored("--local-models", "yellow")
-                + ". Если нужна помощь: "
-                + colored("--help", "yellow")
-            )
-            return None
-        elif ollama_list_of_models and use_local_models:
-            if not model:
-                if len(ollama_list_of_models) > 1:
-                    print(
-                        colored(
-                            "Для использования локальных моделей необходимо "
-                            "выбрать модель:",
-                            "yellow",
-                        )
-                        + "\n"
-                        + "\n".join(
-                            [
-                                f"{i + 1}. {colored(model, 'magenta', False,)}"
-                                for i, model in enumerate(
-                                    ollama_list_of_models
-                                )
-                            ]
-                        )
-                    )
-                    model_is_selected = False
-                    while not model_is_selected:
-                        model = input(
-                            colored(
-                                "Введите число от 1 до "
-                                f"{len(ollama_list_of_models)}: ",
-                                "yellow",
-                            )
-                        )
-                        model = int(model) if model.isdigit() else -1
-                        if model > len(ollama_list_of_models) or model == -1:
-                            continue
-                        model = ollama_list_of_models[model - 1]
-                        model_is_selected = True
-                        break
-                else:
-                    model = ollama_list_of_models[0]
-            else:
-                if model not in ollama_list_of_models:
-                    print(
-                        colored(
-                            f"{model} не является доступной моделью! ", "red"
-                        )
-                        + "Доступные модели: "
-                        + colored(
-                            f"{', '.join(ollama_list_of_models)}", "yellow"
-                        )
-                    )
-                    return None
-        if model:
-            print("Выбрана модель: " + colored(model, "yellow"))
-
-        # Проверяем, есть ли .git
+        # Check if .git exists
         dot_git = ".git" in os.listdir("./")
 
-        # Если есть
+        # If .git exists
         if dot_git:
-            # Получаем разницу в коммитах
+            # Get commit differences
             git_status = subprocess.run(
                 ["git", "status"],
                 capture_output=True,
@@ -421,9 +170,7 @@ def main() -> None:
 
             if excluded_files:
                 git_diff_command = ["git", "diff", "--staged", "--", "."]
-                git_diff_command.extend(
-                    [f":!{file}" for file in excluded_files]
-                )
+                git_diff_command.extend([f":!{file}" for file in excluded_files])  # noqa
             else:
                 git_diff_command = ["git", "diff", "--staged"]
 
@@ -444,15 +191,18 @@ def main() -> None:
                         encoding="utf-8",
                     ).stdout
                 )
-            ):  # Проверка на отсутствие каких-либо изменений
-                print(colored("Нет добавленных изменений!", "red"))
+            ):  # Check for no changes
+                console.print(
+                    "[red]No changes added![/red]",
+                    highlight=False,
+                )
                 return None
             if not git_diff.stdout:
                 if not dry_run:
                     if (
                         input(
-                            colored("Нет застейдженных изменений!", "red")
-                            + " Добавить всё автоматически с помощью "
+                            colored("No staged changes!", "red")
+                            + " Add all automatically using "
                             + colored("git add -A", "yellow")
                             + "? [y/N]: "
                         )
@@ -462,17 +212,18 @@ def main() -> None:
                             ["git", "add", "-A"],
                         )
                     else:
-                        print(
-                            colored(
-                                "Добавьте необходимые файлы вручную.", "yellow"
-                            )
+                        console.print(
+                            "Add required files manually.",
+                            style="yellow",
+                            highlight=False,
                         )
                         return None
                 else:
-                    print(
-                        colored("Нечего коммитить!", "red")
-                        + " Добавьте необходимые файлы с помощью "
-                        + colored("git add <filename>", "yellow")
+                    console.print(
+                        "[red]Nothing to commit![/red]"
+                        " Add required files using "
+                        "[yellow]git add <filename>[/yellow]",
+                        highlight=False,
                     )
                     return None
                 git_diff = subprocess.run(
@@ -486,18 +237,123 @@ def main() -> None:
                 capture_output=True,
                 encoding="utf-8",
             ).stdout:
-                print(
-                    colored(
-                        "Обратите внимание на то, что у Вас "
-                        "есть незастейдженные изменения!",
-                        "red",
+                console.print(
+                    "[red]Note: You have unstaged changes![/red]"
+                    " To add more files, press "
+                    "[yellow]Ctrl + C[/yellow] and run "
+                    "[yellow]git add <filename>[/yellow]",
+                    highlight=False,
+                )
+
+            if use_local_models:
+                # Check Ollama installation
+                try:
+                    subprocess.run(
+                        ["ollama", "--version"],
+                        text=True,
+                        capture_output=True,
+                    )
+                except FileNotFoundError:
+                    console.print(
+                        "Ollama is not installed!",
+                        style="red bold",
                     )
-                    + " Для добавления дополнительных файлов "
-                    + colored("Ctrl + C", "yellow")
-                    + " и выполните "
-                    + colored("git add <filename>", "yellow")
-                    + "."
+                    return None
+
+                # Check if Ollama is running
+                ollama_served = (
+                    requests.get("http://localhost:11434").status_code == 200
                 )
+
+                if ollama_served:
+                    # Get list of models from Ollama
+                    ollama_models_json = requests.get(
+                        "http://localhost:11434/api/tags"
+                    ).json()
+                    if ollama_models_json["models"]:
+                        ollama_list_of_models = [
+                            i["model"] for i in ollama_models_json["models"]
+                        ]
+                    else:
+                        console.print(
+                            "[yellow]Ollama model list is empty!"
+                            "[/yellow] To install models, visit "
+                            "https://ollama.com/models",
+                            highlight=False,
+                        )
+                        return None
+                else:
+                    console.print(
+                        "[yellow]Ollama server not running\n"
+                        "or not installed![/yellow]"
+                    )
+                    return None
+            else:
+                ollama_list_of_models = 0
+
+            # Handle missing Ollama
+            if not ollama_list_of_models and use_local_models:
+                console.print(
+                    "[yellow]Ollama is not installed or model list is empty!"
+                    "[/yellow] To install, visit "
+                    "https://ollama.com/download",
+                    highlight=False,
+                )
+                return None
+            elif not use_local_models and model:
+                console.print(
+                    f"To use {model} locally, use the flag "
+                    "[yellow]--local-models[/yellow]. For help: "
+                    "[yellow]--help[/yellow]",
+                    highlight=False,
+                )
+                return None
+            elif ollama_list_of_models and use_local_models:
+                if not model:
+                    if len(ollama_list_of_models) > 1:
+                        console.print(
+                            "[yellow]Select a local model:[/yellow]\n"
+                            + "\n".join(
+                                [
+                                    f"[magenta]{i + 1}. {model}[/magenta]"
+                                    for i, model in enumerate(ollama_list_of_models)  # noqa
+                                ]
+                            ),
+                            highlight=False,
+                        )
+                        model_is_selected = False
+                        while not model_is_selected:
+                            model = prompt.ask(
+                                "[yellow]Enter a number from 1 to "
+                                f"{len(ollama_list_of_models)}[/yellow]",
+                            )
+                            if not (1 <= model <= len(ollama_list_of_models)):
+                                console.print(
+                                    "[red]Enter a valid number![/red]",
+                                    highlight=False,
+                                )
+                                continue
+                            model = ollama_list_of_models[model - 1]
+                            model_is_selected = True
+                            break
+                    else:
+                        model = ollama_list_of_models[0]
+                else:
+                    if model not in ollama_list_of_models:
+                        console.print(
+                            f"[red]{model} is not an available model!"
+                            "[/red] "
+                            "Available models: [yellow]"
+                            f"{', '.join(ollama_list_of_models)}[/yellow]",
+                            highlight=False,
+                        )
+                        return None
+            if model:
+                console.print(
+                    f"Selected model: [yellow]{model}[/yellow]",
+                    highlight=False,
+                )
+            # Create AI client
             if use_local_models:
                 client = Ollama(model=model)
             else:
@@ -509,19 +365,30 @@ def main() -> None:
                 retry = True
                 while retry:
                     with console.status(
-                        "[magenta bold]Создание сообщения коммита...",
+                        "[magenta bold]Generating commit message...",
                         spinner_style="magenta",
                     ):
-                        commit_message = client.message(
-                            message=prompt_for_ai
-                            + "Git status: "
-                            + git_status.stdout
-                            + "Git diff: "
-                            + git_diff.stdout,
-                            temperature=temperature,
+                        commit_message = cut_think(
+                            client.message(
+                                messages=[
+                                    {
+                                        "role": "system",
+                                        "content": prompt_for_ai,
+                                    },
+                                    {
+                                        "role": "user",
+                                        "content": "Git status: "
+                                        + git_status.stdout
+                                        + "Git diff: "
+                                        + git_diff.stdout,
+                                    },
+                                ],
+                                temperature=temperature,
+                                timeout=timeout,
+                            )
                         )
                     commit_with_message_from_ai = input(
-                        "Закоммитить с сообщением "
+                        "Commit with message "
                         + colored(f"'{commit_message}'", "yellow")
                         + "? [y/N/r]: "
                     )
@@ -533,37 +400,45 @@ def main() -> None:
                         ["git", "commit", "-m", f"{commit_message}"],
                         encoding="utf-8",
                     )
-                    print(colored("Коммит успешно создан!", "green"))
-            else:
-                commit_message = client.message(
-                    message=prompt_for_ai
-                    + "Git status: "
-                    + git_status.stdout
-                    + "Git diff: "
-                    + git_diff.stdout,
-                    temperature=temperature,
-                )
-                print(
-                    colored(
-                        "Коммит-месседж успешно сгенерирован:", "green", False
+                    console.print(
+                        "Commit created successfully!",
+                        style="green bold",
+                        highlight=False,
                     )
-                )
-                print(
-                    colored(
-                        commit_message,
-                        "yellow",
-                        False,
+            else:
+                with console.status(
+                    "[magenta bold]Generating commit message...",
+                    spinner_style="magenta",
+                ):
+                    commit_message = cut_think(
+                        client.message(
+                            messages=[
+                                {
+                                    "role": "system",
+                                    "content": prompt_for_ai,
+                                },
+                                {
+                                    "role": "user",
+                                    "content": "Git status: "
+                                    + git_status.stdout
+                                    + "Git diff: "
+                                    + git_diff.stdout,
+                                },
+                            ],
+                            temperature=temperature,
+                            timeout=timeout,
+                        )
                     )
-                )
+                console.print(commit_message, style="yellow", highlight=False)
                 return None
 
-        # Если нет
+        # If .git does not exist
         else:
             init_git_repo = (
                 True
                 if input(
-                    colored("Не инициализирован git репозиторий!", "red")
-                    + " Выполнить "
+                    colored("Git repository not initialized!", "red")
+                    + " Run "
                     + colored("git init", "yellow")
                     + "? [y/N]: "
                 )
@@ -593,21 +468,17 @@ def main() -> None:
                         ),
                     )
                     if input(
-                        "Сделать первый коммит с сообщением "
+                        "Make first commit with message "
                         + colored("'Initial commit?'", "yellow")
                         + " [y/N]: "
                     )
                     == "y"
                     else None
                 )
-    except FileNotFoundError as e:
-        print(
-            colored("Ollama не установлена!", "red")
-            if "ollama" in str(e)
-            else colored(str(e), "red")
-        )
-    except Exception as e:
-        print(colored("Ошибка:", "red") + " " + str(e))
+    except KeyboardInterrupt:
+        return None
+    except Exception:
+        console.print_exception()
 
 
 if __name__ == "__main__":