From a7daa31273f01549f89085dca192c64f8f80d12f Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Tue, 1 Apr 2025 17:06:48 +0800
Subject: [PATCH 1/7] perf: Resurrection Kael
---
pkg/httpd/chat.go | 292 ++++++++++++++++++----------
pkg/httpd/message.go | 11 +-
pkg/httpd/webserver.go | 6 +-
pkg/jms-sdk-go/model/account.go | 12 ++
pkg/jms-sdk-go/service/jms_asset.go | 6 +
pkg/jms-sdk-go/service/url.go | 1 +
pkg/proxy/chat.go | 170 ++++++++++++++++
pkg/srvconn/conn_openai.go | 12 +-
8 files changed, 401 insertions(+), 109 deletions(-)
create mode 100644 pkg/proxy/chat.go
diff --git a/pkg/httpd/chat.go b/pkg/httpd/chat.go
index fc2e6b6d1..9836cf941 100644
--- a/pkg/httpd/chat.go
+++ b/pkg/httpd/chat.go
@@ -1,163 +1,253 @@
package httpd
import (
+ "context"
"encoding/json"
+ "fmt"
"github.com/jumpserver/koko/pkg/common"
+ "github.com/jumpserver/koko/pkg/i18n"
+ "github.com/jumpserver/koko/pkg/logger"
+ "github.com/jumpserver/koko/pkg/proxy"
+ "github.com/jumpserver/koko/pkg/session"
"github.com/sashabaranov/go-openai"
"sync"
"time"
"github.com/jumpserver/koko/pkg/jms-sdk-go/model"
- "github.com/jumpserver/koko/pkg/logger"
"github.com/jumpserver/koko/pkg/srvconn"
)
var _ Handler = (*chat)(nil)
type chat struct {
- ws *UserWebsocket
+ ws *UserWebsocket
+ term *model.TerminalConfig
- conversationMap sync.Map
-
- termConf *model.TerminalConfig
+ // conversationMap: map[conversationID]*AIConversation
+ conversations sync.Map
}
func (h *chat) Name() string {
return ChatName
}
-func (h *chat) CleanUp() {
- h.CleanConversationMap()
-}
+func (h *chat) CleanUp() { h.cleanupAll() }
func (h *chat) CheckValidation() error {
return nil
}
func (h *chat) HandleMessage(msg *Message) {
- conversationID := msg.Id
- conversation := &AIConversation{}
-
- if conversationID == "" {
- id := common.UUID()
- conversation = &AIConversation{
- Id: id,
- Prompt: msg.Prompt,
- HistoryRecords: make([]string, 0),
- InterruptCurrentChat: false,
- }
+ if msg.Interrupt {
+ h.interrupt(msg.Id)
+ return
+ }
- // T000 Currently a websocket connection only retains one conversation
- h.CleanConversationMap()
- h.conversationMap.Store(id, conversation)
- } else {
- c, ok := h.conversationMap.Load(conversationID)
- if !ok {
- logger.Errorf("Ws[%s] conversation %s not found", h.ws.Uuid, conversationID)
- h.sendErrorMessage(conversationID, "conversation not found")
- return
+ conv, err := h.getOrCreateConversation(msg)
+ if err != nil {
+ h.sendError(msg.Id, err.Error())
+ return
+ }
+ conv.Question = msg.Data
+ conv.NewDialogue = true
+
+ go h.runChat(conv)
+}
+
+func (h *chat) getOrCreateConversation(msg *Message) (*AIConversation, error) {
+ if msg.Id != "" {
+ if v, ok := h.conversations.Load(msg.Id); ok {
+ return v.(*AIConversation), nil
}
- conversation = c.(*AIConversation)
+ return nil, fmt.Errorf("conversation %s not found", msg.Id)
}
- if msg.Interrupt {
- conversation.InterruptCurrentChat = true
- return
+ jmsSrv, err := proxy.NewChatJMSServer(
+ h.ws.user.String(), h.ws.ClientIP(),
+ h.ws.user.ID, h.ws.langCode, h.ws.apiClient, h.term,
+ )
+ if err != nil {
+ return nil, fmt.Errorf("create JMS server: %w", err)
}
- openAIParam := &OpenAIParam{
- AuthToken: h.termConf.GptApiKey,
- BaseURL: h.termConf.GptBaseUrl,
- Proxy: h.termConf.GptProxy,
- Model: h.termConf.GptModel,
- Prompt: conversation.Prompt,
+ sess := session.NewSession(jmsSrv.Session, h.sessionCallback)
+ session.AddSession(sess)
+
+ conv := &AIConversation{
+ Id: jmsSrv.Session.ID,
+ Prompt: msg.Prompt,
+ Context: make([]QARecord, 0),
+ JMSServer: jmsSrv,
}
- conversation.HistoryRecords = append(conversation.HistoryRecords, msg.Data)
- go h.chat(openAIParam, conversation)
-}
-
-func (h *chat) chat(
- chatGPTParam *OpenAIParam, conversation *AIConversation,
-) string {
- doneCh := make(chan string)
- answerCh := make(chan string)
- defer close(doneCh)
- defer close(answerCh)
-
- c := srvconn.NewOpenAIClient(
- chatGPTParam.AuthToken,
- chatGPTParam.BaseURL,
- chatGPTParam.Proxy,
+ h.conversations.Store(jmsSrv.Session.ID, conv)
+ go h.Monitor(conv)
+ return conv, nil
+}
+
+func (h *chat) sessionCallback(task *model.TerminalTask) error {
+ if task.Name == model.TaskKillSession {
+ h.endConversation(task.Args, "close", "kill session")
+ return nil
+ }
+ return fmt.Errorf("unknown session task %s", task.Name)
+}
+
+func (h *chat) runChat(conv *AIConversation) {
+ ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
+ defer cancel()
+
+ client := srvconn.NewOpenAIClient(
+ h.term.GptApiKey, h.term.GptBaseUrl, h.term.GptProxy,
)
- startIndex := len(conversation.HistoryRecords) - 15
- if startIndex < 0 {
- startIndex = 0
+ // Keep the last 8 contexts
+ if len(conv.Context) > 8 {
+ conv.Context = conv.Context[len(conv.Context)-8:]
}
- contents := conversation.HistoryRecords[startIndex:]
-
- openAIConn := &srvconn.OpenAIConn{
- Id: conversation.Id,
- Client: c,
- Prompt: chatGPTParam.Prompt,
- Model: chatGPTParam.Model,
- Contents: contents,
+ messages := buildChatMessages(conv)
+
+ conn := &srvconn.OpenAIConn{
+ Id: conv.Id,
+ Client: client,
+ Prompt: conv.Prompt,
+ Model: h.term.GptModel,
+ Question: conv.Question,
+ Context: messages,
+ AnswerCh: make(chan string),
+ DoneCh: make(chan string),
IsReasoning: false,
- AnswerCh: answerCh,
- DoneCh: doneCh,
- Type: h.termConf.ChatAIType,
+ Type: h.term.ChatAIType,
}
- go openAIConn.Chat(&conversation.InterruptCurrentChat)
- return h.processChatMessages(openAIConn)
+ // 启动 streaming
+ go conn.Chat(&conv.InterruptCurrentChat)
+
+ conv.JMSServer.Replay.WriteInput(conv.Question)
+
+ h.streamResponses(ctx, conv, conn)
+}
+
+func buildChatMessages(conv *AIConversation) []openai.ChatCompletionMessage {
+ msgs := make([]openai.ChatCompletionMessage, 0, len(conv.Context)*2)
+ for _, r := range conv.Context {
+ msgs = append(msgs,
+ openai.ChatCompletionMessage{Role: openai.ChatMessageRoleUser, Content: r.Question},
+ openai.ChatCompletionMessage{Role: openai.ChatMessageRoleAssistant, Content: r.Answer},
+ )
+ }
+ return msgs
}
-func (h *chat) processChatMessages(
- openAIConn *srvconn.OpenAIConn,
-) string {
- messageID := common.UUID()
- id := openAIConn.Id
+func (h *chat) streamResponses(
+ ctx context.Context, conv *AIConversation, conn *srvconn.OpenAIConn,
+) {
+ msgID := common.UUID()
for {
select {
- case answer := <-openAIConn.AnswerCh:
- h.sendSessionMessage(id, answer, messageID, "message", openAIConn.IsReasoning)
- case answer := <-openAIConn.DoneCh:
- h.sendSessionMessage(id, answer, messageID, "finish", false)
- return answer
+ case <-ctx.Done():
+ h.sendError(conv.Id, "chat timeout")
+ return
+ case ans := <-conn.AnswerCh:
+ h.sendMessage(conv.Id, msgID, ans, "message", conn.IsReasoning)
+ case ans := <-conn.DoneCh:
+ h.sendMessage(conv.Id, msgID, ans, "finish", false)
+ h.finalizeConversation(conv, ans)
+ return
}
}
}
-func (h *chat) sendSessionMessage(id, answer, messageID, messageType string, isReasoning bool) {
- message := ChatGPTMessage{
- Content: answer,
- ID: messageID,
+func (h *chat) finalizeConversation(conv *AIConversation, fullAnswer string) {
+ runes := []rune(fullAnswer)
+ snippet := fullAnswer
+ if len(runes) > 100 {
+ snippet = string(runes[:100])
+ }
+ conv.Context = append(conv.Context, QARecord{Question: conv.Question, Answer: snippet})
+
+ cmd := conv.JMSServer.GenerateCommandItem(h.ws.user.String(), conv.Question, fullAnswer)
+ go conv.JMSServer.CmdR.Record(cmd)
+ go conv.JMSServer.Replay.WriteOutput(fullAnswer)
+}
+
+func (h *chat) sendMessage(
+ convID, msgID, content, typ string, reasoning bool,
+) {
+ msg := ChatGPTMessage{
+ Content: content,
+ ID: msgID,
CreateTime: time.Now(),
- Type: messageType,
+ Type: typ,
Role: openai.ChatMessageRoleAssistant,
- IsReasoning: isReasoning,
+ IsReasoning: reasoning,
}
- data, _ := json.Marshal(message)
- msg := Message{
- Id: id,
- Type: "message",
- Data: string(data),
+ data, _ := json.Marshal(msg)
+ h.ws.SendMessage(&Message{Id: convID, Type: "message", Data: string(data)})
+}
+
+func (h *chat) sendError(convID, errMsg string) {
+ h.endConversation(convID, "error", errMsg)
+}
+
+func (h *chat) endConversation(convID, typ, msg string) {
+
+ defer func() {
+ if r := recover(); r != nil {
+ logger.Errorf("panic while sending message to session %s: %v", convID, r)
+ }
+ }()
+
+ if v, ok := h.conversations.Load(convID); ok {
+ if conv, ok2 := v.(*AIConversation); ok2 && conv.JMSServer != nil {
+ conv.JMSServer.Close(msg)
+ }
}
- h.ws.SendMessage(&msg)
+ h.conversations.Delete(convID)
+ h.ws.SendMessage(&Message{Id: convID, Type: typ, Data: msg})
}
-func (h *chat) sendErrorMessage(id, message string) {
- msg := Message{
- Id: id,
- Type: "error",
- Data: message,
+func (h *chat) interrupt(convID string) {
+ if v, ok := h.conversations.Load(convID); ok {
+ v.(*AIConversation).InterruptCurrentChat = true
}
- h.ws.SendMessage(&msg)
}
-func (h *chat) CleanConversationMap() {
- h.conversationMap.Range(func(key, value interface{}) bool {
- h.conversationMap.Delete(key)
+func (h *chat) cleanupAll() {
+ h.conversations.Range(func(key, _ interface{}) bool {
+ h.endConversation(key.(string), "close", "")
return true
})
}
+
+func (h *chat) Monitor(conv *AIConversation) {
+ lang := i18n.NewLang(h.ws.langCode)
+
+ lastActiveTime := time.Now()
+ maxIdleTime := time.Duration(h.term.MaxIdleTime) * time.Minute
+ MaxSessionTime := time.Now().Add(time.Duration(h.term.MaxSessionTime) * time.Hour)
+
+ for {
+ now := time.Now()
+ if MaxSessionTime.Before(now) {
+ msg := lang.T("Session max time reached, disconnect")
+ logger.Infof("Session[%s] max session time reached, disconnect", conv.Id)
+ h.endConversation(conv.Id, "close", msg)
+ return
+ }
+
+ outTime := lastActiveTime.Add(maxIdleTime)
+ if now.After(outTime) {
+ msg := fmt.Sprintf(lang.T("Connect idle more than %d minutes, disconnect"), h.term.MaxIdleTime)
+ logger.Infof("Session[%s] idle more than %d minutes, disconnect", conv.Id, h.term.MaxIdleTime)
+ h.endConversation(conv.Id, "close", msg)
+ return
+ }
+
+ if conv.NewDialogue {
+ lastActiveTime = time.Now()
+ conv.NewDialogue = false
+ }
+
+ time.Sleep(10 * time.Second)
+ }
+}
diff --git a/pkg/httpd/message.go b/pkg/httpd/message.go
index 9ed65854f..191fa4f27 100644
--- a/pkg/httpd/message.go
+++ b/pkg/httpd/message.go
@@ -1,6 +1,7 @@
package httpd
import (
+ "github.com/jumpserver/koko/pkg/proxy"
"time"
"github.com/jumpserver/koko/pkg/exchange"
@@ -163,11 +164,19 @@ type OpenAIParam struct {
Type string
}
+type QARecord struct {
+ Question string
+ Answer string
+}
+
type AIConversation struct {
Id string
Prompt string
- HistoryRecords []string
+ Question string
+ Context []QARecord
+ JMSServer *proxy.ChatJMSServer
InterruptCurrentChat bool
+ NewDialogue bool
}
type ChatGPTMessage struct {
diff --git a/pkg/httpd/webserver.go b/pkg/httpd/webserver.go
index 86ff6e280..ca494d1b4 100644
--- a/pkg/httpd/webserver.go
+++ b/pkg/httpd/webserver.go
@@ -158,9 +158,9 @@ func (s *Server) ChatAIWebsocket(ctx *gin.Context) {
}
userConn.handler = &chat{
- ws: userConn,
- conversationMap: sync.Map{},
- termConf: &termConf,
+ ws: userConn,
+ conversations: sync.Map{},
+ term: &termConf,
}
s.broadCaster.EnterUserWebsocket(userConn)
defer s.broadCaster.LeaveUserWebsocket(userConn)
diff --git a/pkg/jms-sdk-go/model/account.go b/pkg/jms-sdk-go/model/account.go
index 0e9e893ed..fbedc194b 100644
--- a/pkg/jms-sdk-go/model/account.go
+++ b/pkg/jms-sdk-go/model/account.go
@@ -55,6 +55,18 @@ type AccountDetail struct {
Privileged bool `json:"privileged"`
}
+type AssetChat struct {
+ ID string `json:"id"`
+ Name string `json:"name"`
+}
+
+type AccountChatDetail struct {
+ ID string `json:"id"`
+ Name string `json:"name"`
+ Username string `json:"username"`
+ Asset AssetChat `json:"asset"`
+}
+
type PermAccount struct {
Name string `json:"name"`
Username string `json:"username"`
diff --git a/pkg/jms-sdk-go/service/jms_asset.go b/pkg/jms-sdk-go/service/jms_asset.go
index fc15ce8cf..39d57c430 100644
--- a/pkg/jms-sdk-go/service/jms_asset.go
+++ b/pkg/jms-sdk-go/service/jms_asset.go
@@ -23,3 +23,9 @@ func (s *JMService) GetAccountSecretById(accountId string) (res model.AccountDet
_, err = s.authClient.Get(url, &res)
return
}
+
+func (s *JMService) GetAccountChat() (res model.AccountChatDetail, err error) {
+ url := fmt.Sprintf(AccountChatURL)
+ _, err = s.authClient.Get(url, &res)
+ return
+}
diff --git a/pkg/jms-sdk-go/service/url.go b/pkg/jms-sdk-go/service/url.go
index 893e40d5b..39051e826 100644
--- a/pkg/jms-sdk-go/service/url.go
+++ b/pkg/jms-sdk-go/service/url.go
@@ -77,6 +77,7 @@ const (
UserPermsAssetAccountsURL = "/api/v1/perms/users/%s/assets/%s/"
AccountSecretURL = "/api/v1/assets/account-secrets/%s/"
+ AccountChatURL = "/api/v1/accounts/accounts/chat/"
UserPermsAssetsURL = "/api/v1/perms/users/%s/assets/"
AssetLoginConfirmURL = "/api/v1/acls/login-asset/check/"
diff --git a/pkg/proxy/chat.go b/pkg/proxy/chat.go
new file mode 100644
index 000000000..0e845f206
--- /dev/null
+++ b/pkg/proxy/chat.go
@@ -0,0 +1,170 @@
+package proxy
+
+import (
+ "fmt"
+ "github.com/jumpserver/koko/pkg/common"
+ modelCommon "github.com/jumpserver/koko/pkg/jms-sdk-go/common"
+ "github.com/jumpserver/koko/pkg/jms-sdk-go/model"
+ "github.com/jumpserver/koko/pkg/jms-sdk-go/service"
+ "github.com/jumpserver/koko/pkg/logger"
+ "github.com/jumpserver/koko/pkg/session"
+ "strings"
+ "time"
+)
+
+type ChatReplyRecorder struct {
+ *ReplyRecorder
+}
+
+func (rh *ChatReplyRecorder) WriteInput(inputStr string) {
+ currentTime := time.Now()
+ formattedTime := currentTime.Format("2006-01-02 15:04:05")
+ inputStr = fmt.Sprintf("[%s]#: %s", formattedTime, inputStr)
+ rh.Record([]byte(inputStr))
+}
+
+func (rh *ChatReplyRecorder) WriteOutput(outputStr string) {
+ wrappedText := rh.wrapText(outputStr)
+ outputStr = "\r\n" + wrappedText + "\r\n"
+ rh.Record([]byte(outputStr))
+
+}
+
+func (rh *ChatReplyRecorder) wrapText(text string) string {
+ var wrappedTextBuilder strings.Builder
+ words := strings.Fields(text)
+ currentLineLength := 0
+
+ for _, word := range words {
+ wordLength := len(word)
+
+ if currentLineLength+wordLength > rh.Writer.Width {
+ wrappedTextBuilder.WriteString("\r\n" + word + " ")
+ currentLineLength = wordLength + 1
+ } else {
+ wrappedTextBuilder.WriteString(word + " ")
+ currentLineLength += wordLength + 1
+ }
+ }
+
+ return wrappedTextBuilder.String()
+}
+
+func NewChatJMSServer(
+ user, ip, userID, langCode string,
+ jmsService *service.JMService, conf *model.TerminalConfig) (*ChatJMSServer, error) {
+ accountInfo, err := jmsService.GetAccountChat()
+ if err != nil {
+ logger.Errorf("Get account chat info error: %s", err)
+ return nil, err
+ }
+
+ id := common.UUID()
+
+ apiSession := &model.Session{
+ ID: id,
+ User: user,
+ LoginFrom: model.LoginFromWeb,
+ RemoteAddr: ip,
+ Protocol: model.ActionALL,
+ Asset: accountInfo.Asset.Name,
+ Account: accountInfo.Name,
+ AccountID: accountInfo.ID,
+ AssetID: accountInfo.Asset.ID,
+ UserID: userID,
+ OrgID: "00000000-0000-0000-0000-000000000004",
+ Type: model.NORMALType,
+ LangCode: langCode,
+ DateStart: modelCommon.NewNowUTCTime(),
+ }
+
+ _, err2 := jmsService.CreateSession(*apiSession)
+ if err2 != nil {
+ return nil, err2
+ }
+
+ chat := &ChatJMSServer{
+ JmsService: jmsService,
+ Session: apiSession,
+ Conf: conf,
+ }
+
+ chat.CmdR = chat.GetCommandRecorder()
+ chat.Replay = chat.GetReplayRecorder()
+
+ if err1 := jmsService.RecordSessionLifecycleLog(id, model.AssetConnectSuccess,
+ model.EmptyLifecycleLog); err1 != nil {
+ logger.Errorf("Record session activity log err: %s", err1)
+ }
+
+ return chat, nil
+}
+
+type ChatJMSServer struct {
+ JmsService *service.JMService
+ Session *model.Session
+ CmdR *CommandRecorder
+ Replay *ChatReplyRecorder
+ Conf *model.TerminalConfig
+}
+
+func (s *ChatJMSServer) GenerateCommandItem(user, input, output string) *model.Command {
+ createdDate := time.Now()
+ return &model.Command{
+ SessionID: s.Session.ID,
+ OrgID: s.Session.OrgID,
+ Input: input,
+ Output: output,
+ User: user,
+ Server: s.Session.Asset,
+ Account: s.Session.Account,
+ Timestamp: createdDate.Unix(),
+ RiskLevel: model.NormalLevel,
+ DateCreated: createdDate.UTC(),
+ }
+}
+
+func (s *ChatJMSServer) GetReplayRecorder() *ChatReplyRecorder {
+ info := &ReplyInfo{
+ Width: 200,
+ Height: 200,
+ TimeStamp: time.Now(),
+ }
+ recorder, err := NewReplayRecord(s.Session.ID, s.JmsService,
+ NewReplayStorage(s.JmsService, s.Conf),
+ info)
+ if err != nil {
+ logger.Error(err)
+ }
+
+ return &ChatReplyRecorder{recorder}
+}
+
+func (s *ChatJMSServer) GetCommandRecorder() *CommandRecorder {
+ cmdR := CommandRecorder{
+ sessionID: s.Session.ID,
+ storage: NewCommandStorage(s.JmsService, s.Conf),
+ queue: make(chan *model.Command, 10),
+ closed: make(chan struct{}),
+ jmsService: s.JmsService,
+ }
+ go cmdR.record()
+ return &cmdR
+}
+
+func (s *ChatJMSServer) Close(msg string) {
+ session.RemoveSessionById(s.Session.ID)
+ if err := s.JmsService.SessionFinished(s.Session.ID, modelCommon.NewNowUTCTime()); err != nil {
+ logger.Errorf("finish session %s: %v", s.Session.ID, err)
+ }
+
+ s.CmdR.End()
+ s.Replay.End()
+
+ logObj := model.SessionLifecycleLog{Reason: msg, User: s.Session.User}
+ err := s.JmsService.RecordSessionLifecycleLog(s.Session.ID, model.AssetConnectFinished, logObj)
+ if err != nil {
+ logger.Errorf("record session lifecycle log %s: %v", s.Session.ID, err)
+ return
+ }
+}
diff --git a/pkg/srvconn/conn_openai.go b/pkg/srvconn/conn_openai.go
index 9d639a6a9..50b1e5ad6 100644
--- a/pkg/srvconn/conn_openai.go
+++ b/pkg/srvconn/conn_openai.go
@@ -93,7 +93,8 @@ type OpenAIConn struct {
Client *openai.Client
Model string
Prompt string
- Contents []string
+ Question string
+ Context []openai.ChatCompletionMessage
IsReasoning bool
AnswerCh chan string
DoneCh chan string
@@ -102,11 +103,10 @@ type OpenAIConn struct {
func (conn *OpenAIConn) Chat(interruptCurrentChat *bool) {
ctx := context.Background()
- var messages []openai.ChatCompletionMessage
- messages = append(messages, openai.ChatCompletionMessage{
+ messages := append(conn.Context, openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleUser,
- Content: strings.Join(conn.Contents, "\n"),
+ Content: conn.Question,
})
systemPrompt := conn.Prompt
@@ -182,6 +182,10 @@ func (conn *OpenAIConn) Chat(interruptCurrentChat *bool) {
newContent = response.Choices[0].Delta.Content
}
+ if newContent == "" {
+ continue
+ }
+
content += newContent
conn.AnswerCh <- content
}
From 45bc9cf8c107cad20d368ffdb1bdbb01d4d12203 Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Wed, 30 Apr 2025 17:41:30 +0800
Subject: [PATCH 2/7] perf: add chat model
---
pkg/httpd/chat.go | 8 +++++++-
pkg/httpd/message.go | 2 ++
2 files changed, 9 insertions(+), 1 deletion(-)
diff --git a/pkg/httpd/chat.go b/pkg/httpd/chat.go
index 9836cf941..3d5c4b80a 100644
--- a/pkg/httpd/chat.go
+++ b/pkg/httpd/chat.go
@@ -76,6 +76,7 @@ func (h *chat) getOrCreateConversation(msg *Message) (*AIConversation, error) {
conv := &AIConversation{
Id: jmsSrv.Session.ID,
Prompt: msg.Prompt,
+ Model: msg.ChatModel,
Context: make([]QARecord, 0),
JMSServer: jmsSrv,
}
@@ -106,11 +107,16 @@ func (h *chat) runChat(conv *AIConversation) {
}
messages := buildChatMessages(conv)
+ chatModel := conv.Model
+ if conv.Model == "" {
+ chatModel = h.term.GptModel
+ }
+
conn := &srvconn.OpenAIConn{
Id: conv.Id,
Client: client,
Prompt: conv.Prompt,
- Model: h.term.GptModel,
+ Model: chatModel,
Question: conv.Question,
Context: messages,
AnswerCh: make(chan string),
diff --git a/pkg/httpd/message.go b/pkg/httpd/message.go
index 191fa4f27..3c98fb2d4 100644
--- a/pkg/httpd/message.go
+++ b/pkg/httpd/message.go
@@ -19,6 +19,7 @@ type Message struct {
//Chat AI
Prompt string `json:"prompt"`
Interrupt bool `json:"interrupt"`
+ ChatModel string `json:"chat_model"`
//K8s
KubernetesId string `json:"k8s_id"`
@@ -173,6 +174,7 @@ type AIConversation struct {
Id string
Prompt string
Question string
+ Model string
Context []QARecord
JMSServer *proxy.ChatJMSServer
InterruptCurrentChat bool
From 505035df2e0c9fa47c74c7e893003c1b448e5218 Mon Sep 17 00:00:00 2001
From: zhaojisen <1301338853@qq.com>
Date: Wed, 30 Apr 2025 17:47:46 +0800
Subject: [PATCH 3/7] Perf: change UI
---
ui/components.d.ts | 1 +
ui/public/icons/logo.svg | 1 +
ui/src/hooks/helper/index.ts | 3 +
ui/src/index.css | 7 ++
ui/src/overrides.ts | 2 +-
.../chat/components/Conversation/index.vue | 70 ++++++++++++++++
.../components/Conversation/optionRender.tsx | 72 +++++++++++++++++
ui/src/views/chat/components/Header/index.vue | 21 +++++
ui/src/views/chat/components/Sider/index.vue | 59 ++++++++++++++
.../views/chat/components/Welcome/index.vue | 14 ++++
ui/src/views/chat/index.vue | 79 +++++++------------
11 files changed, 279 insertions(+), 50 deletions(-)
create mode 100755 ui/public/icons/logo.svg
create mode 100644 ui/src/views/chat/components/Conversation/index.vue
create mode 100644 ui/src/views/chat/components/Conversation/optionRender.tsx
create mode 100644 ui/src/views/chat/components/Header/index.vue
create mode 100644 ui/src/views/chat/components/Sider/index.vue
create mode 100644 ui/src/views/chat/components/Welcome/index.vue
diff --git a/ui/components.d.ts b/ui/components.d.ts
index 7d80f5542..7b71b1018 100644
--- a/ui/components.d.ts
+++ b/ui/components.d.ts
@@ -59,6 +59,7 @@ declare module 'vue' {
NResult: typeof import('naive-ui')['NResult']
NScrollbar: typeof import('naive-ui')['NScrollbar']
NSelect: typeof import('naive-ui')['NSelect']
+ NSpace: typeof import('naive-ui')['NSpace']
NSpin: typeof import('naive-ui')['NSpin']
NSplit: typeof import('naive-ui')['NSplit']
NSwitch: typeof import('naive-ui')['NSwitch']
diff --git a/ui/public/icons/logo.svg b/ui/public/icons/logo.svg
new file mode 100755
index 000000000..eb25227d1
--- /dev/null
+++ b/ui/public/icons/logo.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/ui/src/hooks/helper/index.ts b/ui/src/hooks/helper/index.ts
index 6b99888a4..9b6ef1586 100644
--- a/ui/src/hooks/helper/index.ts
+++ b/ui/src/hooks/helper/index.ts
@@ -306,6 +306,9 @@ export const generateWsURL = () => {
connectURL = BASE_WS_URL + '/koko/ws/terminal/?' + requireParams.toString();
break;
}
+ case 'Chat':
+ connectURL = BASE_WS_URL + `/koko/ws/chat/system`
+ break;
default: {
connectURL = urlParams ? `${BASE_WS_URL}/koko/ws/terminal/?${urlParams.toString()}` : '';
}
diff --git a/ui/src/index.css b/ui/src/index.css
index 3d552a61f..db1950261 100644
--- a/ui/src/index.css
+++ b/ui/src/index.css
@@ -1,2 +1,9 @@
@import "tailwindcss";
+.icon-hover-primary {
+ @apply cursor-pointer hover:text-[#16987D] focus:outline-none;
+}
+
+.icon-hover-danger {
+ @apply cursor-pointer hover:text-[#ff0000] focus:outline-none;
+}
diff --git a/ui/src/overrides.ts b/ui/src/overrides.ts
index 707313c2a..30a5dedbc 100644
--- a/ui/src/overrides.ts
+++ b/ui/src/overrides.ts
@@ -126,7 +126,7 @@ export const themeOverrides: GlobalThemeOverrides = {
},
Layout: {
color: 'rgba(0, 0, 0, 1)',
- siderColor: 'rgba(0, 0, 0, 1)',
+ siderColor: 'rgba(255, 255, 255, 0.09)',
headerColor: 'rgba(0, 0, 0, 1)'
}
};
diff --git a/ui/src/views/chat/components/Conversation/index.vue b/ui/src/views/chat/components/Conversation/index.vue
new file mode 100644
index 000000000..7c0761a39
--- /dev/null
+++ b/ui/src/views/chat/components/Conversation/index.vue
@@ -0,0 +1,70 @@
+
+
+
+ 今天
+
+
+
+
+
+ 随便聊聊
+
+
+
+
+
+
+
+
+
+
+
+
+
+ 昨天
+
+
+
+
+
+ 随便聊聊
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/ui/src/views/chat/components/Conversation/optionRender.tsx b/ui/src/views/chat/components/Conversation/optionRender.tsx
new file mode 100644
index 000000000..df8aed2c8
--- /dev/null
+++ b/ui/src/views/chat/components/Conversation/optionRender.tsx
@@ -0,0 +1,72 @@
+import { useI18n } from 'vue-i18n';
+import { NSpace, NText } from 'naive-ui';
+import { SquareArrowOutUpRight, PencilLine, Trash2 } from 'lucide-vue-next';
+
+import type { SelectOption } from 'naive-ui';
+import type { FunctionalComponent } from 'vue';
+import type { LucideProps } from 'lucide-vue-next';
+
+interface OptionItem {
+ value: string;
+ label: string;
+ textColor?: string;
+ iconColor?: string;
+ click: (chatId: string) => void;
+ icon: FunctionalComponent;
+}
+
+type EmitsType = {
+ (e: 'chat-share', shareId: string): void;
+ (e: 'chat-rename', shareId: string): void;
+ (e: 'chat-delete', shareId: string): void;
+};
+
+export const OptionRender = (emits: EmitsType): SelectOption[] => {
+ const { t } = useI18n();
+
+ const optionItems: OptionItem[] = [
+ {
+ value: 'share',
+ icon: SquareArrowOutUpRight,
+ label: t('Share'),
+ iconColor: 'white',
+ click: (chatId: string) => {
+ emits('chat-share', chatId);
+ }
+ },
+ {
+ value: 'rename',
+ icon: PencilLine,
+ label: t('Rename'),
+ iconColor: 'white',
+ click: (chatId: string) => {
+ emits('chat-rename', chatId);
+ }
+ },
+ {
+ value: 'delete',
+ icon: Trash2,
+ label: t('Delete'),
+ iconColor: '#fb2c36',
+ textColor: '!text-red-500',
+ click: (chatId: string) => {
+ emits('chat-delete', chatId);
+ }
+ }
+ ];
+
+ const commonClass = 'px-4 py-2 w-28 hover:bg-[#ffffff1A] cursor-pointer transition-all duration-300';
+
+ return optionItems.map(item => ({
+ value: item.value,
+ render: () => (
+ item.click(item.value)}>
+ {item.icon && }
+
+
+ {item.label}
+
+
+ )
+ }));
+};
diff --git a/ui/src/views/chat/components/Header/index.vue b/ui/src/views/chat/components/Header/index.vue
new file mode 100644
index 000000000..a62d03f3a
--- /dev/null
+++ b/ui/src/views/chat/components/Header/index.vue
@@ -0,0 +1,21 @@
+
+
+
+
+
+ GPT-4
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/ui/src/views/chat/components/Sider/index.vue b/ui/src/views/chat/components/Sider/index.vue
new file mode 100644
index 000000000..fc8881802
--- /dev/null
+++ b/ui/src/views/chat/components/Sider/index.vue
@@ -0,0 +1,59 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/ui/src/views/chat/components/Welcome/index.vue b/ui/src/views/chat/components/Welcome/index.vue
new file mode 100644
index 000000000..fcd0ec290
--- /dev/null
+++ b/ui/src/views/chat/components/Welcome/index.vue
@@ -0,0 +1,14 @@
+
+
+
+ 散财消灾
+
+
+
+
+
diff --git a/ui/src/views/chat/index.vue b/ui/src/views/chat/index.vue
index e0446676e..3018d7e89 100644
--- a/ui/src/views/chat/index.vue
+++ b/ui/src/views/chat/index.vue
@@ -1,48 +1,15 @@
-
-
+
-
- 随便聊聊
-
- 随便聊聊
- 随便聊聊
- 随便聊聊
- 随便聊聊
-
-
-
-
- 随便聊聊
-
-
-
-
- GPT-4
-
-
-
+
+
+
-
-
- 散财消灾
-
-
+
@@ -51,17 +18,31 @@
-
+// ws.value?.onopen(() => {
+// console.log('open')
+// })
+
+// ws.value?.onmessage((message: MessageEvent) => {
+// console.log(message)
+// })
+
+
+onMounted(() => {
+
+})
+
From 73c055a0a75436fe5066d03e40b05d24c801bd87 Mon Sep 17 00:00:00 2001
From: zhaojisen <1301338853@qq.com>
Date: Wed, 30 Apr 2025 18:33:34 +0800
Subject: [PATCH 4/7] Perf: Socket Connect
---
ui/src/hooks/helper/index.ts | 3 ++-
ui/src/hooks/useChat.ts | 45 +++++++++++++++++++++++++++++++
ui/src/store/modules/chat.ts | 8 ++++++
ui/src/types/modules/chat.type.ts | 5 ++++
ui/src/views/chat/index.vue | 5 ++--
5 files changed, 63 insertions(+), 3 deletions(-)
create mode 100644 ui/src/hooks/useChat.ts
create mode 100644 ui/src/store/modules/chat.ts
create mode 100644 ui/src/types/modules/chat.type.ts
diff --git a/ui/src/hooks/helper/index.ts b/ui/src/hooks/helper/index.ts
index 9b6ef1586..328edb554 100644
--- a/ui/src/hooks/helper/index.ts
+++ b/ui/src/hooks/helper/index.ts
@@ -307,7 +307,8 @@ export const generateWsURL = () => {
break;
}
case 'Chat':
- connectURL = BASE_WS_URL + `/koko/ws/chat/system`
+ // connectURL = BASE_WS_URL + `/koko/ws/chat/system`
+ connectURL = 'ws://localhost:5050' + `/koko/ws/chat/system/`
break;
default: {
connectURL = urlParams ? `${BASE_WS_URL}/koko/ws/terminal/?${urlParams.toString()}` : '';
diff --git a/ui/src/hooks/useChat.ts b/ui/src/hooks/useChat.ts
new file mode 100644
index 000000000..10c6d7bb3
--- /dev/null
+++ b/ui/src/hooks/useChat.ts
@@ -0,0 +1,45 @@
+import { ref } from 'vue'
+import { MessageType } from '@/enum'
+import { useWebSocket } from '@vueuse/core'
+import { generateWsURL } from '@/hooks/helper';
+
+export const useChat = () => {
+
+ const socket = ref();
+
+ const socketOnMessage = (message: MessageEvent) => {
+ const messageData = JSON.parse(message.data);
+
+ console.log(messageData);
+
+ switch (messageData.type) {
+ case MessageType.CONNECT:
+ break;
+ }
+ }
+
+ const createChatSocket = () => {
+ const url = generateWsURL()
+
+ const { ws } = useWebSocket(url)
+
+ if (!ws.value) {
+
+ }
+
+ ws.value.onopen(() => {
+ // TODO 心跳
+ console.log('Connected to websocket');
+ })
+ ws.value.onmessage = socketOnMessage
+
+ return {
+ socket: socket.value,
+ }
+ }
+
+
+ return {
+ createChatSocket
+ }
+}
\ No newline at end of file
diff --git a/ui/src/store/modules/chat.ts b/ui/src/store/modules/chat.ts
new file mode 100644
index 000000000..c776c81f1
--- /dev/null
+++ b/ui/src/store/modules/chat.ts
@@ -0,0 +1,8 @@
+import { defineStore } from 'pinia'
+
+export const useChatStore = defineStore('chat-store', {
+ state: () => ({}),
+ actions: {
+
+ }
+})
\ No newline at end of file
diff --git a/ui/src/types/modules/chat.type.ts b/ui/src/types/modules/chat.type.ts
new file mode 100644
index 000000000..a26defdc2
--- /dev/null
+++ b/ui/src/types/modules/chat.type.ts
@@ -0,0 +1,5 @@
+export interface ChatState {
+ username: string;
+
+ messageId: string
+}
\ No newline at end of file
diff --git a/ui/src/views/chat/index.vue b/ui/src/views/chat/index.vue
index 3018d7e89..6906f7536 100644
--- a/ui/src/views/chat/index.vue
+++ b/ui/src/views/chat/index.vue
@@ -25,13 +25,14 @@ import InputArea from './components/InputArea/index.vue';
import { onMounted } from 'vue'
import { useMessage } from 'naive-ui'
+import { useChat } from '@/hooks/useChat.ts'
import { useWebSocket } from '@vueuse/core'
import { generateWsURL } from '@/hooks/helper'
const message = useMessage()
+const { createChatSocket } = useChat()
-const url = generateWsURL()
-const { ws } = useWebSocket(url)
+createChatSocket()
// ws.value?.onopen(() => {
// console.log('open')
From e047b4844e863d48c18c76102d5a70699d0001e9 Mon Sep 17 00:00:00 2001
From: zhaojisen <1301338853@qq.com>
Date: Tue, 6 May 2025 14:37:16 +0800
Subject: [PATCH 5/7] Perf: Basic layout
---
ui/package.json | 2 +-
ui/src/index.css | 4 +--
.../views/chat/components/InputArea/index.vue | 9 +++++-
ui/src/views/chat/index.vue | 31 +++++++++----------
ui/yarn.lock | 8 ++---
5 files changed, 29 insertions(+), 25 deletions(-)
diff --git a/ui/package.json b/ui/package.json
index ffb133840..0f5108024 100644
--- a/ui/package.json
+++ b/ui/package.json
@@ -23,7 +23,7 @@
"clipboard-polyfill": "^4.1.0",
"dayjs": "^1.11.13",
"loglevel": "^1.9.1",
- "lucide-vue-next": "^0.487.0",
+ "lucide-vue-next": "^0.507.0",
"mitt": "^3.0.1",
"naive-ui": "^2.39.0",
"nora-zmodemjs": "^1.1.1",
diff --git a/ui/src/index.css b/ui/src/index.css
index db1950261..2534f1b56 100644
--- a/ui/src/index.css
+++ b/ui/src/index.css
@@ -1,9 +1,9 @@
@import "tailwindcss";
.icon-hover-primary {
- @apply cursor-pointer hover:text-[#16987D] focus:outline-none;
+ @apply cursor-pointer hover:text-[#16987D] focus:outline-none transition-all duration-300;
}
.icon-hover-danger {
- @apply cursor-pointer hover:text-[#ff0000] focus:outline-none;
+ @apply cursor-pointer hover:text-[#ff0000] focus:outline-none transition-all duration-300;
}
diff --git a/ui/src/views/chat/components/InputArea/index.vue b/ui/src/views/chat/components/InputArea/index.vue
index 225e4cc47..0d30f1d79 100644
--- a/ui/src/views/chat/components/InputArea/index.vue
+++ b/ui/src/views/chat/components/InputArea/index.vue
@@ -1,7 +1,13 @@
-
+
+
+
+
+
+
+
发送
@@ -11,6 +17,7 @@
diff --git a/ui/yarn.lock b/ui/yarn.lock
index 60fa84f0c..9c73e53f9 100644
--- a/ui/yarn.lock
+++ b/ui/yarn.lock
@@ -2006,10 +2006,10 @@ lru-cache@^5.1.1:
dependencies:
yallist "^3.0.2"
-lucide-vue-next@^0.487.0:
- version "0.487.0"
- resolved "https://registry.yarnpkg.com/lucide-vue-next/-/lucide-vue-next-0.487.0.tgz#18d9e6b21c37c823d1c4879203b323dd97ada5f0"
- integrity sha512-ilVgu9EHkfId7WSjmoPkzp13cuzpSGO5J16AmltjRHFoj6MlCBGY8BzsBU/ISKVlDheUxM+MsIYjNo/1hSEa6w==
+lucide-vue-next@^0.507.0:
+ version "0.507.0"
+ resolved "https://registry.yarnpkg.com/lucide-vue-next/-/lucide-vue-next-0.507.0.tgz#bd2d9a5f85550875fb6f85082dd108c46a1870fd"
+ integrity sha512-n0AZRmez4xq5vu5ZOfnrhC5wKBpu5tpwDHA6Ue2+sKyJv3USAPxqIerXLbdYbkHRW1tVZ3U9GE1MACsu6X/Z7Q==
magic-string@*, magic-string@^0.30.10, magic-string@^0.30.11:
version "0.30.11"
From 89d109e05bbeb2790099b0458f8646e6753d62c3 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]"
Date: Tue, 6 May 2025 06:47:46 +0000
Subject: [PATCH 6/7] perf: Update Dockerfile with new base image tag
---
Dockerfile | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/Dockerfile b/Dockerfile
index 21a318e04..464a33855 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM jumpserver/koko-base:20250429_034346 AS stage-build
+FROM jumpserver/koko-base:20250506_063744 AS stage-build
WORKDIR /opt/koko
ARG TARGETARCH
From b2183d19cb7794e7a2c697b1adc2992c9a631195 Mon Sep 17 00:00:00 2001
From: zhaojisen <1301338853@qq.com>
Date: Thu, 8 May 2025 15:13:38 +0800
Subject: [PATCH 7/7] Fixed: socket error
---
ui/components.d.ts | 4 ---
ui/src/hooks/useChat.ts | 36 +++++++++----------
.../components/Conversation/optionRender.tsx | 2 +-
ui/src/views/chat/components/Sider/index.vue | 7 ----
4 files changed, 19 insertions(+), 30 deletions(-)
diff --git a/ui/components.d.ts b/ui/components.d.ts
index 7b71b1018..e9abf586b 100644
--- a/ui/components.d.ts
+++ b/ui/components.d.ts
@@ -40,13 +40,11 @@ declare module 'vue' {
NFormItemGi: typeof import('naive-ui')['NFormItemGi']
NGrid: typeof import('naive-ui')['NGrid']
NGridItem: typeof import('naive-ui')['NGridItem']
- NH2: typeof import('naive-ui')['NH2']
NIcon: typeof import('naive-ui')['NIcon']
NImage: typeof import('naive-ui')['NImage']
NInput: typeof import('naive-ui')['NInput']
NLayout: typeof import('naive-ui')['NLayout']
NLayoutContent: typeof import('naive-ui')['NLayoutContent']
- NLayoutFooter: typeof import('naive-ui')['NLayoutFooter']
NLayoutHeader: typeof import('naive-ui')['NLayoutHeader']
NLayoutSider: typeof import('naive-ui')['NLayoutSider']
NList: typeof import('naive-ui')['NList']
@@ -63,13 +61,11 @@ declare module 'vue' {
NSpin: typeof import('naive-ui')['NSpin']
NSplit: typeof import('naive-ui')['NSplit']
NSwitch: typeof import('naive-ui')['NSwitch']
- NTab: typeof import('naive-ui')['NTab']
NTable: typeof import('naive-ui')['NTable']
NTabPane: typeof import('naive-ui')['NTabPane']
NTabs: typeof import('naive-ui')['NTabs']
NTag: typeof import('naive-ui')['NTag']
NText: typeof import('naive-ui')['NText']
- NTextarea: typeof import('naive-ui')['NTextarea']
NThing: typeof import('naive-ui')['NThing']
NTooltip: typeof import('naive-ui')['NTooltip']
NTree: typeof import('naive-ui')['NTree']
diff --git a/ui/src/hooks/useChat.ts b/ui/src/hooks/useChat.ts
index 10c6d7bb3..4eeea5479 100644
--- a/ui/src/hooks/useChat.ts
+++ b/ui/src/hooks/useChat.ts
@@ -1,10 +1,9 @@
-import { ref } from 'vue'
-import { MessageType } from '@/enum'
-import { useWebSocket } from '@vueuse/core'
+import { ref } from 'vue';
+import { MessageType } from '@/enum';
+import { useWebSocket } from '@vueuse/core';
import { generateWsURL } from '@/hooks/helper';
export const useChat = () => {
-
const socket = ref();
const socketOnMessage = (message: MessageEvent) => {
@@ -16,30 +15,31 @@ export const useChat = () => {
case MessageType.CONNECT:
break;
}
- }
+ };
const createChatSocket = () => {
- const url = generateWsURL()
+ const url = generateWsURL();
- const { ws } = useWebSocket(url)
+ const { ws } = useWebSocket(url);
if (!ws.value) {
-
+ return;
}
- ws.value.onopen(() => {
- // TODO 心跳
+ ws.value.onopen = () => {
console.log('Connected to websocket');
- })
- ws.value.onmessage = socketOnMessage
+ };
+ ws.value.onmessage = socketOnMessage;
+ ws.value.onclose = () => {
+ console.log('Disconnected from websocket');
+ };
return {
- socket: socket.value,
- }
- }
-
+ socket: socket.value
+ };
+ };
return {
createChatSocket
- }
-}
\ No newline at end of file
+ };
+};
diff --git a/ui/src/views/chat/components/Conversation/optionRender.tsx b/ui/src/views/chat/components/Conversation/optionRender.tsx
index df8aed2c8..544ea6332 100644
--- a/ui/src/views/chat/components/Conversation/optionRender.tsx
+++ b/ui/src/views/chat/components/Conversation/optionRender.tsx
@@ -55,7 +55,7 @@ export const OptionRender = (emits: EmitsType): SelectOption[] => {
}
];
- const commonClass = 'px-4 py-2 w-28 hover:bg-[#ffffff1A] cursor-pointer transition-all duration-300';
+ const commonClass = 'px-4 py-2 w-30 hover:bg-[#ffffff1A] cursor-pointer transition-all duration-300';
return optionItems.map(item => ({
value: item.value,
diff --git a/ui/src/views/chat/components/Sider/index.vue b/ui/src/views/chat/components/Sider/index.vue
index fc8881802..6e597e062 100644
--- a/ui/src/views/chat/components/Sider/index.vue
+++ b/ui/src/views/chat/components/Sider/index.vue
@@ -13,12 +13,6 @@
-
-
-
-
-
import { reactive } from 'vue';
-import { Search, SquarePen } from 'lucide-vue-next';
import Conversation from '../Conversation/index.vue';
import type { CSSProperties } from 'vue';