feat(wip): code completion with llm

api/openai/code_completion.go

@@ -0,0 +1,82 @@
+package openai
+
+import (
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/0xJacky/Nginx-UI/api"
+	"github.com/0xJacky/Nginx-UI/internal/helper"
+	"github.com/0xJacky/Nginx-UI/internal/llm"
+	"github.com/0xJacky/Nginx-UI/settings"
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	"github.com/uozi-tech/cosy"
+	"github.com/uozi-tech/cosy/logger"
+)
+
+var mutex sync.Mutex
+
+// CodeCompletion handles code completion requests
+func CodeCompletion(c *gin.Context) {
+	if !settings.OpenAISettings.EnableCodeCompletion {
+		cosy.ErrHandler(c, llm.ErrCodeCompletionNotEnabled)
+		return
+	}
+
+	var upgrader = websocket.Upgrader{
+		CheckOrigin: func(r *http.Request) bool {
+			return true
+		},
+	}
+	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		cosy.ErrHandler(c, err)
+		return
+	}
+	defer ws.Close()
+
+	for {
+		var codeCompletionRequest llm.CodeCompletionRequest
+		err := ws.ReadJSON(&codeCompletionRequest)
+		if err != nil {
+			if helper.IsUnexpectedWebsocketError(err) {
+				logger.Errorf("Error reading JSON: %v", err)
+			}
+			return
+		}
+
+		codeCompletionRequest.UserID = api.CurrentUser(c).ID
+
+		go func() {
+			start := time.Now()
+			completedCode, err := codeCompletionRequest.Send()
+			if err != nil {
+				logger.Errorf("Error sending code completion request: %v", err)
+				return
+			}
+			elapsed := time.Since(start)
+
+			mutex.Lock()
+			defer mutex.Unlock()
+
+			err = ws.WriteJSON(gin.H{
+				"code":          completedCode,
+				"request_id":    codeCompletionRequest.RequestID,
+				"completion_ms": elapsed.Milliseconds(),
+			})
+			if err != nil {
+				if helper.IsUnexpectedWebsocketError(err) {
+					logger.Errorf("Error writing JSON: %v", err)
+				}
+				return
+			}
+		}()
+	}
+}
+
+func GetCodeCompletionEnabledStatus(c *gin.Context) {
+	c.JSON(http.StatusOK, gin.H{
+		"enabled": settings.OpenAISettings.EnableCodeCompletion,
+	})
+}

api/openai/openai.go

@@ -4,15 +4,16 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/0xJacky/Nginx-UI/internal/chatbot"
+	"io"
+	"strings"
+	"time"
+
+	"github.com/0xJacky/Nginx-UI/internal/llm"
 	"github.com/0xJacky/Nginx-UI/settings"
 	"github.com/gin-gonic/gin"
 	"github.com/sashabaranov/go-openai"
 	"github.com/uozi-tech/cosy"
 	"github.com/uozi-tech/cosy/logger"
-	"io"
-	"strings"
-	"time"
 )
 
 const ChatGPTInitPrompt = `You are a assistant who can help users write and optimise the configurations of Nginx,
@@ -41,7 +42,7 @@ func MakeChatCompletionRequest(c *gin.Context) {
 	messages = append(messages, json.Messages...)
 
 	if json.Filepath != "" {
-		messages = chatbot.ChatCompletionWithContext(json.Filepath, messages)
+		messages = llm.ChatCompletionWithContext(json.Filepath, messages)
 	}
 
 	// SSE server
@@ -50,7 +51,7 @@ func MakeChatCompletionRequest(c *gin.Context) {
 	c.Writer.Header().Set("Connection", "keep-alive")
 	c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
 
-	openaiClient, err := chatbot.GetClient()
+	openaiClient, err := llm.GetClient()
 	if err != nil {
 		c.Stream(func(w io.Writer) bool {
 			c.SSEvent("message", gin.H{

api/openai/router.go

@@ -6,4 +6,7 @@ func InitRouter(r *gin.RouterGroup) {
 	// ChatGPT
 	r.POST("chatgpt", MakeChatCompletionRequest)
 	r.POST("chatgpt_record", StoreChatGPTRecord)
+	// Code Completion
+	r.GET("code_completion", CodeCompletion)
+	r.GET("code_completion/enabled", GetCodeCompletionEnabledStatus)
 }