AI inline autocompletion with Codestral

cmd/micro/micro.go

@@ -516,6 +516,8 @@ func DoEvent() {
 		}
 	case f := <-timerChan:
 		f()
+	case f := <-action.AIResultChan:
+		f()
 	case <-sighup:
 		exit(0)
 	case <-util.Sigterm:

internal/action/actions.go

@@ -907,6 +907,10 @@ func (h *BufPane) OutdentSelection() bool {
 func (h *BufPane) Autocomplete() bool {
 	b := h.Buf
 
+	if b.InlineCompletion != "" {
+		return false
+	}
+
 	if h.Cursor.HasSelection() {
 		return false
 	}
@@ -944,6 +948,9 @@ func (h *BufPane) CycleAutocompleteBack() bool {
 
 // InsertTab inserts a tab or spaces
 func (h *BufPane) InsertTab() bool {
+	if h.acceptAICompletion() {
+		return true
+	}
 	b := h.Buf
 	indent := b.IndentString(util.IntOpt(b.Settings["tabsize"]))
 	tabBytes := len(indent)
@@ -1874,6 +1881,10 @@ func (h *BufPane) ToggleOverwriteMode() bool {
 
 // Escape leaves current mode
 func (h *BufPane) Escape() bool {
+	if h.Buf.InlineCompletion != "" {
+		h.dismissAICompletion()
+		return true
+	}
 	return true
 }
 
@@ -2346,3 +2357,10 @@ func (h *BufPane) RemoveAllMultiCursors() bool {
 func (h *BufPane) None() bool {
 	return true
 }
+
+// ManualTrigger immediately triggers AI completion (no debounce).
+// Bound to Ctrl-Space by default.
+func (h *BufPane) ManualTrigger() bool {
+	h.triggerAICompletionNow()
+	return true
+}

internal/action/bufpane.go

@@ -1,11 +1,15 @@
 package action
 
 import (
+	"context"
+	"log"
 	"strings"
+	"sync"
 	"time"
 
 	luar "layeh.com/gopher-luar"
 
+	"github.com/micro-editor/micro/v2/internal/ai"
 	"github.com/micro-editor/micro/v2/internal/buffer"
 	"github.com/micro-editor/micro/v2/internal/config"
 	"github.com/micro-editor/micro/v2/internal/display"
@@ -16,6 +20,35 @@ import (
 	lua "github.com/yuin/gopher-lua"
 )
 
+// AIResultChan delivers AI completion results to the main event loop.
+// The main loop reads from this channel and applies the result on the main thread.
+var AIResultChan = make(chan func(), 16)
+
+var aiManager     *ai.Manager
+var aiManagerErr  error
+var aiManagerOnce sync.Once
+
+func getAIManager() *ai.Manager {
+	aiManagerOnce.Do(func() {
+		if !config.GetGlobalOption("aicomplete").(bool) {
+			return
+		}
+		provider := config.GetGlobalOption("aicompleteprovider").(string)
+		model := config.GetGlobalOption("aicompletemodel").(string)
+		baseURL := config.GetGlobalOption("aicompleteurl").(string)
+		debounce := config.GetGlobalOption("aicompletedebounce").(float64)
+		aiManager, aiManagerErr = ai.NewManager(provider, model, baseURL, debounce)
+		if aiManagerErr != nil {
+			log.Println("AI completion:", aiManagerErr)
+			InfoBar.Error("AI completion: ", aiManagerErr.Error())
+		}
+	})
+	if aiManagerErr != nil {
+		return nil
+	}
+	return aiManager
+}
+
 type BufAction any
 
 // BufKeyAction represents an action bound to a key.
@@ -254,6 +287,9 @@ type BufPane struct {
 	// since we may not know the window geometry yet. In such case we finish
 	// its initialization a bit later, after the initial resize.
 	initialized bool
+
+	// AI inline completion state
+	aiCancel context.CancelFunc
 }
 
 func newBufPane(buf *buffer.Buffer, win display.BWindow, tab *Tab) *BufPane {
@@ -345,6 +381,7 @@ func (h *BufPane) resetMouse() {
 
 // OpenBuffer opens the given buffer in this pane.
 func (h *BufPane) OpenBuffer(b *buffer.Buffer) {
+	h.dismissAICompletion()
 	h.Buf.Close()
 	h.Buf = b
 	h.BWindow.SetBuffer(b)
@@ -357,6 +394,100 @@ func (h *BufPane) OpenBuffer(b *buffer.Buffer) {
 	h.lastClickTime = time.Time{}
 }
 
+func (h *BufPane) dismissAICompletion() {
+	h.Buf.InlineCompletion = ""
+	if h.aiCancel != nil {
+		h.aiCancel()
+		h.aiCancel = nil
+	}
+}
+
+func (h *BufPane) triggerAICompletion() {
+	mgr := getAIManager()
+	if mgr == nil {
+		if aiManagerErr != nil {
+			InfoBar.Error("AI completion: ", aiManagerErr.Error())
+		}
+		return
+	}
+
+	if !config.GetGlobalOption("aicomplete").(bool) {
+		return
+	}
+
+	h.dismissAICompletion()
+	InfoBar.Message("")
+
+	before := h.Buf.TextBeforeCursor()
+	after := h.Buf.TextAfterCursor()
+	fileType := h.Buf.Settings["filetype"].(string)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	h.aiCancel = cancel
+
+	mgr.RequestDelayed(ai.Request{
+		BeforeCursor: before,
+		AfterCursor:  after,
+		FileType:     fileType,
+		FileName:     h.Buf.AbsPath,
+	}, func(resp *ai.Response, err error) {
+		if ctx.Err() != nil {
+			return
+		}
+		if err != nil {
+			AIResultChan <- func() {
+				InfoBar.Message("AI: ", err.Error())
+			}
+			return
+		}
+		AIResultChan <- func() {
+			h.Buf.InlineCompletion = resp.Text
+			screen.Redraw()
+		}
+	})
+}
+
+func (h *BufPane) triggerAICompletionNow() {
+	mgr := getAIManager()
+	if mgr == nil {
+		if aiManagerErr != nil {
+			InfoBar.Error("AI completion: ", aiManagerErr.Error())
+		}
+		return
+	}
+	if !config.GetGlobalOption("aicomplete").(bool) {
+		return
+	}
+	h.dismissAICompletion()
+	InfoBar.Message("")
+
+	resp, err := mgr.RequestNow(ai.Request{
+		BeforeCursor: h.Buf.TextBeforeCursor(),
+		AfterCursor:  h.Buf.TextAfterCursor(),
+		FileType:     h.Buf.Settings["filetype"].(string),
+		FileName:     h.Buf.AbsPath,
+	})
+	if err != nil {
+		InfoBar.Message("AI: ", err.Error())
+		return
+	}
+	if resp != nil {
+		h.Buf.InlineCompletion = resp.Text
+		screen.Redraw()
+	}
+}
+
+func (h *BufPane) acceptAICompletion() bool {
+	text := h.Buf.InlineCompletion
+	if text == "" {
+		return false
+	}
+	h.dismissAICompletion()
+	h.Buf.Insert(h.Cursor.Loc, text)
+	h.Relocate()
+	return true
+}
+
 // GotoLoc moves the cursor to a new location and adjusts the view accordingly.
 // Use GotoLoc when the new location may be far away from the current location.
 func (h *BufPane) GotoLoc(loc buffer.Loc) {
@@ -556,6 +687,12 @@ func (h *BufPane) execAction(action BufAction, name string, te *tcell.EventMouse
 	if name != "Autocomplete" && name != "CycleAutocompleteBack" {
 		h.Buf.HasSuggestions = false
 	}
+	if name != "Autocomplete" && name != "CycleAutocompleteBack" &&
+		name != "Escape" && name != "InsertTab" && name != "ManualTrigger" &&
+		name != "IndentSelection" && name != "OutdentSelection" &&
+		name != "OutdentLine" {
+		h.dismissAICompletion()
+	}
 
 	if !h.PluginCB("pre"+name, te) {
 		return false
@@ -645,6 +782,7 @@ func (h *BufPane) DoRuneInsert(r rune) {
 		h.Relocate()
 		h.PluginCB("onRune", string(r))
 	}
+	h.triggerAICompletion()
 }
 
 // VSplitIndex opens the given buffer in a vertical split on the given side.
@@ -851,6 +989,7 @@ var BufKeyActions = map[string]BufKeyAction{
 	"Deselect":                  (*BufPane).Deselect,
 	"ClearInfo":                 (*BufPane).ClearInfo,
 	"None":                      (*BufPane).None,
+	"ManualTrigger":             (*BufPane).ManualTrigger,
 
 	// This was changed to InsertNewline but I don't want to break backwards compatibility
 	"InsertEnter": (*BufPane).InsertNewline,

internal/action/defaults_darwin.go

@@ -76,6 +76,7 @@ var bufdefaults = map[string]string{
 	"Ctrl-u":         "ToggleMacro",
 	"Ctrl-j":         "PlayMacro",
 	"Insert":         "ToggleOverwriteMode",
+	"Ctrl-Space":     "ManualTrigger",
 
 	// Emacs-style keybindings
 	"Alt-f": "WordRight",

internal/action/defaults_other.go

@@ -79,6 +79,7 @@ var bufdefaults = map[string]string{
 	"Ctrl-u":         "ToggleMacro",
 	"Ctrl-j":         "PlayMacro",
 	"Insert":         "ToggleOverwriteMode",
+	"Ctrl-Space":     "ManualTrigger",
 
 	// Emacs-style keybindings
 	"Alt-f": "WordRight",

internal/ai/codestral.go

@@ -0,0 +1,116 @@
+package ai
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"time"
+)
+
+const (
+	codestralDefaultBaseURL = "https://codestral.mistral.ai/v1/fim/completions"
+	codestralDefaultModel   = "codestral-latest"
+	codestralEnvKey         = "MISTRAL_API_KEY"
+)
+
+type CodestralProvider struct {
+	apiKey  string
+	model   string
+	baseURL string
+	client  *http.Client
+}
+
+func NewCodestralProvider(model, baseURL string) (*CodestralProvider, error) {
+	apiKey := os.Getenv(codestralEnvKey)
+	if apiKey == "" {
+		return nil, fmt.Errorf("codestral: %s environment variable not set", codestralEnvKey)
+	}
+	if model == "" {
+		model = codestralDefaultModel
+	}
+	if baseURL == "" {
+		baseURL = codestralDefaultBaseURL
+	}
+	return &CodestralProvider{
+		apiKey:  apiKey,
+		model:   model,
+		baseURL: baseURL,
+		client: &http.Client{
+			Timeout: 10 * time.Second,
+		},
+	}, nil
+}
+
+func (p *CodestralProvider) Name() string {
+	return "codestral"
+}
+
+type codestralRequest struct {
+	Model       string   `json:"model"`
+	Prompt      string   `json:"prompt"`
+	Suffix      string   `json:"suffix"`
+	MaxTokens   int      `json:"max_tokens"`
+	Temperature float64  `json:"temperature"`
+	Stop        []string `json:"stop,omitempty"`
+}
+
+type codestralResponse struct {
+	Choices []struct {
+		Message struct {
+			Content string `json:"content"`
+		} `json:"message"`
+	} `json:"choices"`
+}
+
+func (p *CodestralProvider) Complete(ctx context.Context, req Request) (*Response, error) {
+	body := codestralRequest{
+		Model:       p.model,
+		Prompt:      req.BeforeCursor,
+		Suffix:      req.AfterCursor,
+		MaxTokens:   256,
+		Temperature: 0,
+		Stop:        []string{"\n\n"},
+	}
+
+	data, err := json.Marshal(body)
+	if err != nil {
+		return nil, fmt.Errorf("codestral: marshal request: %w", err)
+	}
+
+	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL, bytes.NewReader(data))
+	if err != nil {
+		return nil, fmt.Errorf("codestral: create request: %w", err)
+	}
+	httpReq.Header.Set("Content-Type", "application/json")
+	httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
+
+	resp, err := p.client.Do(httpReq)
+	if err != nil {
+		return nil, fmt.Errorf("codestral: http request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	respBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("codestral: read response: %w", err)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("codestral: status %d: %s", resp.StatusCode, string(respBody))
+	}
+
+	var cresp codestralResponse
+	if err := json.Unmarshal(respBody, &cresp); err != nil {
+		return nil, fmt.Errorf("codestral: parse response: %w", err)
+	}
+
+	if len(cresp.Choices) == 0 {
+		return nil, fmt.Errorf("codestral: no choices in response")
+	}
+
+	return &Response{Text: cresp.Choices[0].Message.Content}, nil
+}