Generate commit with Ollama

autocommit-with-ollama

#!/bin/bash

usage() {
  echo "Usage: $0 [OPTIONS]"
  echo ""
  echo "Options:"
  echo "  --help                  Show help"
  echo "  --model <model>         Specify mode to use (see ollama list)"
  echo "  --api-mode <api-mode>   Specify api mode to use : chat or generate"
  exit 0
}

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

TIMEOUT=20

# Traiter les options
while [[ $# -gt 0 ]]; do
  case "$1" in
    --help)
      usage
      ;;
    --model)
      MODEL="$2"
      shift 2
      ;;
    --api-mode)
      MODE="$2"
      shift 2
      ;;
    *)
      echo "Unknwon option: $1"
      usage
      ;;
  esac
done

# Check if Ollama exists
if ! type ollama &>/dev/null; then
  printf "\r${RED}You need to install Ollama (https://ollama.com/) !\n${NC}"
  exit
fi

# Check Ollama is running
if ! curl -s http://localhost:11434/ > /dev/null; then
  printf "\r${RED}Ollama is not running !\n${NC}"
  exit
fi

# Check staged files
if [[ -z $(git diff --cached) ]]; then
  printf "\r${RED}No staged files. Do 'git add' first !\n${NC}"
  exit 1
fi

MODES=(chat generate)
if [ ! -n "$MODE" ]; then
  printf "\r${YELLOW}Choose an Ollama API mode :\n${NC}"
  select MODE in ${MODES[@]}; do
    if [[ -n "$MODE" ]]; then
      break
    else
      echo "Invalid choice"
    fi
  done
elif ! echo "${MODES[@]}" | grep -w -q "$MODE"; then
  printf "\r${RED}API mode $MODE doesn't exists\n${NC}"
  exit
fi

echo ""


MODELS=$(ollama list | tail -n +2 | awk '{print $1}' | sed 's/:latest$//')
if [ ! -n "$MODEL" ]; then
  printf "\r${YELLOW}Choose an Ollama model :\n${NC}"
  select MODEL in $MODELS; do
    if [[ -n "$MODEL" ]]; then
      break
    else
      echo "Invalid choice"
    fi
  done
elif ! echo "${MODELS[@]}" | grep -w -q "$MODEL"; then
  printf "\r${RED}Model $MODEL doesn't exists\n${NC}"
  exit
fi

if [[ ! -n "$MODEL" ]]; then
  printf "\r${RED}You need to pull an Ollama model\n${NC}"
  exit
fi


# Get the diff
DIFF=$(git diff --cached --diff-algorithm=minimal)

# Define a prompt
PROMPT=$(cat <<-END
Generate a concise git commit message written in present tense for the following code diff with the given specifications below

Message language: en

Commit message must be a maximum of 50 characters

Exclude anything unnecessary such as translation. Your entire response will be passed directly into git commit.

Choose a type from the type-to-description JSON below that best describes the git diff:

docs: 'Documentation only changes',
style: 'Changes that do not affect the meaning of the code (white-space, formatting, missing semi-colons, etc)',
refactor: 'A code change that neither fixes a bug nor adds a feature',
perf: 'A code change that improves performance',
test: 'Adding missing tests or correcting existing tests',
build: 'Changes that affect the build system or external dependencies',
ci: 'Changes to our CI configuration files and scripts',
chore: "Other changes that don't modify src or test files",
revert: 'Reverts a previous commit',
feat: 'A new feature',
fix: 'A bug fix',

The output response must be in format : <type>(<optional scope>): <commit message>

Changeset:
$DIFF
END
)

spinner_frames=('⠋' '⠙' '⠹' '⠸' '⠼' '⠴' '⠦' '⠧' '⠇' '⠏')
spinner() {
  local pid=$1
  local delay=0.1
  local i=0
  local start_time=$(date +%s)

  echo ""
  while kill -0 $pid 2>/dev/null; do
    local now=$(date +%s)
    local elapsed=$((now - start_time))

    if (( elapsed >= TIMEOUT )); then
      kill -9 $pid 2>/dev/null
      printf "\r\n\n${RED}Timeout after ${TIMEOUT}s${NC}\n"
      return 1
    fi

    local frame="${spinner_frames[i++ % ${#spinner_frames[@]}]}"
    printf "\r${GREEN}${frame} Generating commit message with $MODEL model...${NC}"
    sleep $delay
  done

  return 0
}

send_curl_request() {
  local url=$1
  local body=$2

  RESPONSE_FILE=$(mktemp)
  HTTP_CODE_FILE=$(mktemp)

  {
    HTTP_CODE=$(curl -s -w "%{http_code}" -o "$RESPONSE_FILE" -X POST "$url" \
      -H "Content-Type: application/json" \
      -d "$body")
    echo "$HTTP_CODE" > "$HTTP_CODE_FILE"
  } &
  curl_pid=$!

  spinner $curl_pid >&2

  wait $curl_pid

  local RESPONSE=$(<"$RESPONSE_FILE")
  local HTTP_CODE=$(<"$HTTP_CODE_FILE")

  rm -f "$RESPONSE_FILE" "$HTTP_CODE_FILE"

  if [[ "$HTTP_CODE" == "" ]]; then
    exit 1
  fi

  if [[ "$HTTP_CODE" != "200" ]]; then
    echo -e "\n\n${RED}[$HTTP_CODE] An error occured during Ollama request :${NC}" >&2
    echo "$RESPONSE" >&2
    exit 1
  fi

  echo "$RESPONSE"
}

use_api_chat(){
  BODY=$(jq -n \
    --arg model "$MODEL" \
    --arg role "user" \
    --arg content "$PROMPT" \
    --argjson stream false \
    '{
      model: $model,
      stream: $stream,
      options: {
        temperature: 0.8
      },
      messages: [
        {
          role: $role,
          content: $content
        }
      ]
    }')

  RESPONSE=$(send_curl_request "http://localhost:11434/api/chat" "$BODY")

  if [[ ! -n "$RESPONSE" ]]; then
    exit
  fi

  COMMIT_MSG=$(echo "$RESPONSE" | jq -r '.message.content')
}

use_api_generate(){
  BODY=$(jq -n \
    --arg model "$MODEL" \
    --arg prompt "$PROMPT" \
    --argjson stream false \
    '{
      model: $model,
      prompt: $prompt,
      stream: $stream,
      options: {
        temperature: 0.8
      }
    }')

  RESPONSE=$(send_curl_request "http://localhost:11434/api/generate" "$BODY")

  if [[ ! -n "$RESPONSE" ]]; then
    exit
  fi

  COMMIT_MSG=$(echo "$RESPONSE" | jq -r 'select(.response) | .response' | tr -d '\n')
}

if [[ "$MODE" == "chat" ]]; then
    use_api_chat
elif [[ "$MODE" == "generate" ]]; then
    use_api_generate
fi

echo
echo ""
printf "\r${YELLOW}Suggested message :\n${NC}"
echo "----------------------------------------"
echo "$COMMIT_MSG"
echo "----------------------------------------"
echo ""

# Confirmation
printf "\r${YELLOW}Use this message ? (y/n) :\n${NC}"
read -p "" confirm
if [[ "$confirm" == "y" ]]; then
  git commit -m "$COMMIT_MSG"
else
  echo ""
  printf "\r${RED}Abort\n${NC}"
fi