main.py

import argparse
import time
from openai import OpenAI

from datamodel import CommitMessageSuggestion, DiffData
from diff import get_git_diffs
from utils import fallback_commit_message, is_git_repository
from rate_limiter import rate_limit


SYSTEM_PROMPT = """
You are a concise git commit message generator. You will be given labeled git diffs.
Rules:
- Only use information present in the diffs; do NOT invent changed behavior.
- Output JUST the commit message (no explanations).
- Primary format:
  <type>(<scope>): <short summary>
  [blank line]
  Optional body lines, each max 72 characters.

- Short summary: imperative mood, <= 50 characters.
- Body: optional, explain why the change was made or notable details.
- Choose one type from: feat, fix, docs, chore, refactor, test, perf, ci, style.
- If scope is unclear, omit the parentheses: e.g. "fix: ..."
- If diffs are insufficient to produce a reliable message, return the single token:
  NO_SUGGESTION

Examples:

Input:
[staged] file: a/src/app.py
@@
+def add(a, b):
+    return a + b
- def add_numbers(x, y):
-     return x + y

Output:
refactor(src): rename add_numbers to add and simplify signature

Input:
[staged] file: a/README.md
@@
+Add instructions for setup
-Remove old setup notes

Output:
docs: update README setup instructions

Now generate a commit message for the diffs that follow. Output only the message.
"""

MODEL_NAME = "gemma3:4b"
API_BASE_URL = "http://localhost:11434/v1/"

MAX_RETRIES = 3  # Number of retries for the LLM API call
BASE_DELAY = 1  # Base delay in seconds between retries

client = OpenAI(
    api_key="ollama",
    base_url=API_BASE_URL,
)


def generate_commit_message(diffs: list[DiffData]) -> str | None:
    diffs_str_list = []
    for diff in diffs:
        diff_section = f"[staged] file: a/{diff.filename}\n{diff.staged_changes}\n\n[unstaged] file: a/{diff.filename}\n{diff.unstaged_changes}\n"
        diffs_str_list.append(diff_section)

    diff_prompt = "\n".join(diffs_str_list)

    for attempt in range(MAX_RETRIES):
        commit_message = call_llm_api(diff_prompt)
        if commit_message is not None:
            return commit_message
        print(f"Attempt {attempt + 1} failed. Retrying...")
        delay = BASE_DELAY * (2**attempt)
        time.sleep(delay)
    return None


@rate_limit(max_calls=3, period=60)
def call_llm_api(prompt: str) -> str | None:
    print(f"Calling LLM: model={MODEL_NAME} prompt_len={len(prompt)} chars")

    try:
        completion = client.beta.chat.completions.parse(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            max_tokens=150,
            temperature=0.2,
            timeout=60,
            response_format=CommitMessageSuggestion,
        )

        response = completion.choices[0].message

        print(f"LLM response: {response}")

        if response.parsed:
            content = response.parsed.model_dump_json()
            if content is not None:
                suggestion = CommitMessageSuggestion.model_validate_json(content)
                if suggestion.title == "NO_SUGGESTION":
                    print("LLM indicated no suggestion could be made.")
                    return None
                full_message = suggestion.title
                if suggestion.body:
                    full_message += f"\n\n{suggestion.body}"
                return full_message
        return None
    except Exception as e:
        print(f"LLM API call failed: {str(e)}")
        return None


def main():
    parser = argparse.ArgumentParser(
        description="Generate commit message suggestions using an LLM."
    )
    parser.add_argument(
        "repo_path", nargs="?", help="Path to the git repository", default="."
    )
    args = parser.parse_args()

    if not is_git_repository(args.repo_path):
        print("Error: The specified path is not a valid git repository.")
        return

    diffs = get_git_diffs(args.repo_path)
    if not diffs:
        print("No changes detected in the repository.")
        return

    try:
        commit_message = generate_commit_message(diffs)
        if commit_message is None:
            print("No commit message generated by the LLM. Using fallback message.")
            commit_message = fallback_commit_message(diffs)

        print("Suggested commit message: ")
        print(commit_message)
    except Exception as e:
        print(f"An error occurred while generating the commit message: {str(e)}")
        print("Using fallback commit message.")
        print(fallback_commit_message(diffs))


if __name__ == "__main__":
    main()