Skip to content

Commit 669f664

Browse files
vinyas-bharadwajvinyas-bharadwaj
committed
added ollama support for running local LLM's, Two new environment variables OLLAMA_URL and OLLAMA_MODEL
1 parent e4248b2 commit 669f664

4 files changed

Lines changed: 183 additions & 90 deletions

File tree

‎cmd/commit-msg/main.go‎

Lines changed: 105 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,11 @@ import (
1111
"github.com/dfanso/commit-msg/internal/gemini"
1212
"github.com/dfanso/commit-msg/internal/git"
1313
"github.com/dfanso/commit-msg/internal/grok"
14+
"github.com/dfanso/commit-msg/internal/ollama"
1415
"github.com/dfanso/commit-msg/internal/stats"
1516
"github.com/dfanso/commit-msg/pkg/types"
1617
"github.com/pterm/pterm"
18+
"github.com/joho/godotenv"
1719
)
1820

1921
// main is the entry point of the commit message generator
@@ -43,107 +45,120 @@ func main() {
4345
if apiKey == "" {
4446
log.Fatalf("CLAUDE_API_KEY is not set")
4547
}
48+
case "ollama":
49+
// No API key needed for an llm running locally on ollama
50+
apiKey = ""
4651
default:
4752
log.Fatalf("Invalid COMMIT_LLM value: %s", commitLLM)
4853
}
4954

50-
// Get current directory
51-
currentDir, err := os.Getwd()
52-
if err != nil {
53-
log.Fatalf("Failed to get current directory: %v", err)
55+
// Get current directory
56+
currentDir, err := os.Getwd()
57+
if err != nil {
58+
log.Fatalf("Failed to get current directory: %v", err)
59+
}
60+
61+
// Check if current directory is a git repository
62+
if !git.IsRepository(currentDir) {
63+
log.Fatalf("Current directory is not a Git repository: %s", currentDir)
64+
}
65+
66+
// Create a minimal config for the API
67+
config := &types.Config{
68+
GrokAPI: "https://api.x.ai/v1/chat/completions",
69+
}
70+
71+
// Create a repo config for the current directory
72+
repoConfig := types.RepoConfig{
73+
Path: currentDir,
74+
}
75+
76+
// Get file statistics before fetching changes
77+
fileStats, err := stats.GetFileStatistics(&repoConfig)
78+
if err != nil {
79+
log.Fatalf("Failed to get file statistics: %v", err)
80+
}
81+
82+
// Display header
83+
pterm.DefaultHeader.WithFullWidth().
84+
WithBackgroundStyle(pterm.NewStyle(pterm.BgDarkGray)).
85+
WithTextStyle(pterm.NewStyle(pterm.FgLightWhite)).
86+
Println("🚀 Commit Message Generator")
87+
88+
pterm.Println()
89+
90+
// Display file statistics with icons
91+
display.ShowFileStatistics(fileStats)
92+
93+
if fileStats.TotalFiles == 0 {
94+
pterm.Warning.Println("No changes detected in the Git repository.")
95+
return
96+
}
97+
98+
// Get the changes
99+
changes, err := git.GetChanges(&repoConfig)
100+
if err != nil {
101+
log.Fatalf("Failed to get Git changes: %v", err)
102+
}
103+
104+
if len(changes) == 0 {
105+
pterm.Warning.Println("No changes detected in the Git repository.")
106+
return
107+
}
108+
109+
pterm.Println()
110+
111+
// Show generating spinner
112+
spinnerGenerating, err := pterm.DefaultSpinner.
113+
WithSequence("⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏").
114+
Start("🤖 Generating commit message...")
115+
if err != nil {
116+
log.Fatalf("Failed to start spinner: %v", err)
117+
}
118+
119+
var commitMsg string
120+
if os.Getenv("COMMIT_LLM") == "gemini" {
121+
commitMsg, err = gemini.GenerateCommitMessage(config, changes, apiKey)
122+
} else if os.Getenv("COMMIT_LLM") == "chatgpt" {
123+
commitMsg, err = chatgpt.GenerateCommitMessage(config, changes, apiKey)
124+
} else if os.Getenv("COMMIT_LLM") == "claude" {
125+
commitMsg, err = claude.GenerateCommitMessage(config, changes, apiKey)
126+
} else if os.Getenv("COMMIT_LLM") == "ollama" {
127+
url := os.Getenv("OLLAMA_URL")
128+
if url == "" {
129+
url = "http://localhost:11434/api/generate"
54130
}
55-
56-
// Check if current directory is a git repository
57-
if !git.IsRepository(currentDir) {
58-
log.Fatalf("Current directory is not a Git repository: %s", currentDir)
59-
}
60-
61-
// Create a minimal config for the API
62-
config := &types.Config{
63-
GrokAPI: "https://api.x.ai/v1/chat/completions",
64-
}
65-
66-
// Create a repo config for the current directory
67-
repoConfig := types.RepoConfig{
68-
Path: currentDir,
69-
}
70-
71-
// Get file statistics before fetching changes
72-
fileStats, err := stats.GetFileStatistics(&repoConfig)
73-
if err != nil {
74-
log.Fatalf("Failed to get file statistics: %v", err)
131+
model := os.Getenv("OLLAMA_MODEL")
132+
if model == "" {
133+
model = "llama3:latest"
75134
}
135+
commitMsg, err = ollama.GenerateCommitMessage(config, changes, url, model)
136+
} else {
137+
commitMsg, err = grok.GenerateCommitMessage(config, changes, apiKey)
138+
}
76139

77-
// Display header
78-
pterm.DefaultHeader.WithFullWidth().
79-
WithBackgroundStyle(pterm.NewStyle(pterm.BgDarkGray)).
80-
WithTextStyle(pterm.NewStyle(pterm.FgLightWhite)).
81-
Println("🚀 Commit Message Generator")
82-
83-
pterm.Println()
84-
85-
// Display file statistics with icons
86-
display.ShowFileStatistics(fileStats)
140+
if err != nil {
141+
spinnerGenerating.Fail("Failed to generate commit message")
142+
log.Fatalf("Error: %v", err)
143+
}
87144

88-
if fileStats.TotalFiles == 0 {
89-
pterm.Warning.Println("No changes detected in the Git repository.")
90-
return
91-
}
145+
spinnerGenerating.Success("✅ Commit message generated successfully!")
92146

93-
// Get the changes
94-
changes, err := git.GetChanges(&repoConfig)
95-
if err != nil {
96-
log.Fatalf("Failed to get Git changes: %v", err)
97-
}
147+
pterm.Println()
98148

99-
if len(changes) == 0 {
100-
pterm.Warning.Println("No changes detected in the Git repository.")
101-
return
102-
}
149+
// Display the commit message in a styled panel
150+
display.ShowCommitMessage(commitMsg)
103151

104-
pterm.Println()
105-
106-
// Show generating spinner
107-
spinnerGenerating, err := pterm.DefaultSpinner.
108-
WithSequence("⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏").
109-
Start("🤖 Generating commit message...")
110-
if err != nil {
111-
log.Fatalf("Failed to start spinner: %v", err)
112-
}
113-
114-
var commitMsg string
115-
if os.Getenv("COMMIT_LLM") == "gemini" {
116-
commitMsg, err = gemini.GenerateCommitMessage(config, changes, apiKey)
117-
} else if os.Getenv("COMMIT_LLM") == "chatgpt" {
118-
commitMsg, err = chatgpt.GenerateCommitMessage(config, changes, apiKey)
119-
} else if os.Getenv("COMMIT_LLM") == "claude" {
120-
commitMsg, err = claude.GenerateCommitMessage(config, changes, apiKey)
121-
} else {
122-
commitMsg, err = grok.GenerateCommitMessage(config, changes, apiKey)
123-
}
124-
125-
if err != nil {
126-
spinnerGenerating.Fail("Failed to generate commit message")
127-
log.Fatalf("Error: %v", err)
128-
}
129-
130-
spinnerGenerating.Success("✅ Commit message generated successfully!")
131-
132-
pterm.Println()
133-
134-
// Display the commit message in a styled panel
135-
display.ShowCommitMessage(commitMsg)
136-
137-
// Copy to clipboard
138-
err = clipboard.WriteAll(commitMsg)
139-
if err != nil {
140-
pterm.Warning.Printf("⚠️ Could not copy to clipboard: %v\n", err)
141-
} else {
142-
pterm.Success.Println("📋 Commit message copied to clipboard!")
143-
}
152+
// Copy to clipboard
153+
err = clipboard.WriteAll(commitMsg)
154+
if err != nil {
155+
pterm.Warning.Printf("⚠️ Could not copy to clipboard: %v\n", err)
156+
} else {
157+
pterm.Success.Println("📋 Commit message copied to clipboard!")
158+
}
144159

145-
pterm.Println()
160+
pterm.Println()
146161

147-
// Display changes preview
148-
display.ShowChangesPreview(fileStats)
162+
// Display changes preview
163+
display.ShowChangesPreview(fileStats)
149164
}

‎go.mod‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ require (
3131
github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
3232
github.com/googleapis/gax-go/v2 v2.14.1 // indirect
3333
github.com/gookit/color v1.5.4 // indirect
34+
github.com/joho/godotenv v1.5.1 // indirect
3435
github.com/lithammer/fuzzysearch v1.1.8 // indirect
3536
github.com/mattn/go-runewidth v0.0.16 // indirect
3637
github.com/rivo/uniseg v0.4.7 // indirect

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,8 @@ github.com/gookit/color v1.4.2/go.mod h1:fqRyamkC1W8uxl+lxCQxOT09l/vYfZ+QeiX3rKQ
6161
github.com/gookit/color v1.5.0/go.mod h1:43aQb+Zerm/BWh2GnrgOQm7ffz7tvQXEKV6BFMl7wAo=
6262
github.com/gookit/color v1.5.4 h1:FZmqs7XOyGgCAxmWyPslpiok1k05wmY3SJTytgvYFs0=
6363
github.com/gookit/color v1.5.4/go.mod h1:pZJOeOS8DM43rXbp4AZo1n9zCU2qjpcRko0b6/QJi9w=
64+
github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
65+
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
6466
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
6567
github.com/klauspost/cpuid/v2 v2.0.10/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c=
6668
github.com/klauspost/cpuid/v2 v2.0.12/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c=

‎internal/ollama/ollama.go‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
package ollama
2+
3+
import (
4+
"fmt"
5+
"net/http"
6+
"encoding/json"
7+
"bytes"
8+
"io"
9+
10+
"github.com/dfanso/commit-msg/pkg/types"
11+
12+
)
13+
14+
type OllamaRequest struct {
15+
Model string `json:"model"`
16+
Prompt string `json:"prompt"`
17+
}
18+
19+
type OllamaResponse struct {
20+
Response string `json:"response"`
21+
Done bool `json:"done"`
22+
}
23+
24+
func GenerateCommitMessage(config *types.Config, changes string, url string, model string) (string, error) {
25+
// Use llama3:latest as the default model
26+
if model == "" {
27+
model = "llama3:latest"
28+
}
29+
30+
// Preparing the prompt
31+
prompt := fmt.Sprintf("%s\n\n%s", types.CommitPrompt, changes)
32+
33+
// Generating the request body - add stream: false for non-streaming response
34+
reqBody := map[string]interface{}{
35+
"model": model,
36+
"prompt": prompt,
37+
"stream": false,
38+
}
39+
40+
// Generating the body
41+
body, err := json.Marshal(reqBody)
42+
if err != nil {
43+
return "", fmt.Errorf("failed to marshal request: %v", err)
44+
}
45+
46+
resp, err := http.Post(url, "application/json", bytes.NewBuffer(body))
47+
if err != nil {
48+
return "", fmt.Errorf("failed to send request to Ollama: %v", err)
49+
}
50+
defer resp.Body.Close()
51+
52+
// Read the full response body for better error handling
53+
responseBody, err := io.ReadAll(resp.Body)
54+
if err != nil {
55+
return "", fmt.Errorf("failed to read response body: %v", err)
56+
}
57+
58+
// Check HTTP status
59+
if resp.StatusCode != http.StatusOK {
60+
return "", fmt.Errorf("Ollama API returned status %d: %s", resp.StatusCode, string(responseBody))
61+
}
62+
63+
// Since we set stream: false, we get a single response object
64+
var response OllamaResponse
65+
if err := json.Unmarshal(responseBody, &response); err != nil {
66+
return "", fmt.Errorf("failed to decode response: %v", err)
67+
}
68+
69+
// Check if we got any response
70+
if response.Response == "" {
71+
return "", fmt.Errorf("received empty response from Ollama")
72+
}
73+
74+
return response.Response, nil
75+
}

0 commit comments

Comments
 (0)