复制项目

internal/rpc/bot/agent.go

@@ -0,0 +1,153 @@
+package bot
+
+import (
+	"context"
+	"crypto/rand"
+	"time"
+
+	"git.imall.cloud/openim/chat/pkg/common/constant"
+	"git.imall.cloud/openim/chat/pkg/common/convert"
+	"git.imall.cloud/openim/chat/pkg/common/mctx"
+	"git.imall.cloud/openim/chat/pkg/protocol/bot"
+	pbconstant "git.imall.cloud/openim/protocol/constant"
+	"git.imall.cloud/openim/protocol/sdkws"
+	"git.imall.cloud/openim/protocol/user"
+	"github.com/openimsdk/tools/errs"
+	"github.com/openimsdk/tools/utils/datautil"
+)
+
+func (b *botSvr) CreateAgent(ctx context.Context, req *bot.CreateAgentReq) (*bot.CreateAgentResp, error) {
+	if req.Agent == nil {
+		return nil, errs.ErrArgs.WrapMsg("req.Agent is nil")
+	}
+
+	now := time.Now()
+	imToken, err := b.imCaller.ImAdminTokenWithDefaultAdmin(ctx)
+	if err != nil {
+		return nil, err
+	}
+	ctx = mctx.WithApiToken(ctx, imToken)
+	if req.Agent.UserID != "" {
+		req.Agent.UserID = constant.AgentUserIDPrefix + req.Agent.UserID
+		users, err := b.imCaller.GetUsersInfo(ctx, []string{req.Agent.UserID})
+		if err != nil {
+			return nil, err
+		}
+		if len(users) > 0 {
+			return nil, errs.ErrDuplicateKey.WrapMsg("agent userID already exists")
+		}
+	} else {
+		randUserIDs := make([]string, 5)
+		for i := range randUserIDs {
+			randUserIDs[i] = constant.AgentUserIDPrefix + genID(10)
+		}
+		users, err := b.imCaller.GetUsersInfo(ctx, randUserIDs)
+		if err != nil {
+			return nil, err
+		}
+		if len(users) == len(randUserIDs) {
+			return nil, errs.ErrDuplicateKey.WrapMsg("gen agent userID already exists, please try again")
+		}
+		userIDs := datautil.Batch(func(u *sdkws.UserInfo) string { return u.UserID }, users)
+		for _, uid := range randUserIDs {
+			if datautil.Contain(uid, userIDs...) {
+				continue
+			}
+			req.Agent.UserID = uid
+			break
+		}
+	}
+
+	if err := b.imCaller.AddNotificationAccount(ctx, &user.AddNotificationAccountReq{
+		UserID:         req.Agent.UserID,
+		NickName:       req.Agent.Nickname,
+		FaceURL:        req.Agent.FaceURL,
+		AppMangerLevel: pbconstant.AppRobotAdmin,
+	}); err != nil {
+		return nil, err
+	}
+	dbagent := convert.PB2DBAgent(req.Agent)
+	dbagent.CreateTime = now
+	err = b.database.CreateAgent(ctx, dbagent)
+	if err != nil {
+		return nil, err
+	}
+	return &bot.CreateAgentResp{}, nil
+}
+
+func (b *botSvr) UpdateAgent(ctx context.Context, req *bot.UpdateAgentReq) (*bot.UpdateAgentResp, error) {
+	if _, err := b.database.TakeAgent(ctx, req.UserID); err != nil {
+		return nil, errs.ErrArgs.Wrap()
+	}
+
+	if req.FaceURL != nil || req.Nickname != nil {
+		imReq := &user.UpdateNotificationAccountInfoReq{
+			UserID: req.UserID,
+		}
+		if req.Nickname != nil {
+			imReq.NickName = *req.Nickname
+		}
+		if req.FaceURL != nil {
+			imReq.FaceURL = *req.FaceURL
+		}
+		imToken, err := b.imCaller.ImAdminTokenWithDefaultAdmin(ctx)
+		if err != nil {
+			return nil, err
+		}
+		ctx = mctx.WithApiToken(ctx, imToken)
+		err = b.imCaller.UpdateNotificationAccount(ctx, imReq)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	update := ToDBAgentUpdate(req)
+	err := b.database.UpdateAgent(ctx, req.UserID, update)
+	if err != nil {
+		return nil, err
+	}
+
+	return &bot.UpdateAgentResp{}, nil
+}
+
+func (b *botSvr) PageFindAgent(ctx context.Context, req *bot.PageFindAgentReq) (*bot.PageFindAgentResp, error) {
+	total, agents, err := b.database.PageAgents(ctx, req.UserIDs, req.Pagination)
+	if err != nil {
+		return nil, err
+	}
+	//_, userType, err := mctx.Check(ctx)
+	//if err != nil {
+	//	return nil, err
+	//}
+	//if userType != constant.AdminUser {
+	for i := range agents {
+		agents[i].Key = ""
+	}
+	//}
+	return &bot.PageFindAgentResp{
+		Total:  total,
+		Agents: convert.BatchDB2PBAgent(agents),
+	}, nil
+}
+
+func (b *botSvr) DeleteAgent(ctx context.Context, req *bot.DeleteAgentReq) (*bot.DeleteAgentResp, error) {
+	err := b.database.DeleteAgents(ctx, req.UserIDs)
+	if err != nil {
+		return nil, err
+	}
+	return &bot.DeleteAgentResp{}, nil
+}
+
+func genID(l int) string {
+	data := make([]byte, l)
+	_, _ = rand.Read(data)
+	chars := []byte("0123456789")
+	for i := 0; i < len(data); i++ {
+		if i == 0 {
+			data[i] = chars[1:][data[i]%9]
+		} else {
+			data[i] = chars[data[i]%10]
+		}
+	}
+	return string(data)
+}

internal/rpc/bot/send.go

@@ -0,0 +1,92 @@
+package bot
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	"git.imall.cloud/openim/chat/pkg/botstruct"
+	"git.imall.cloud/openim/chat/pkg/common/imapi"
+	"git.imall.cloud/openim/chat/pkg/common/mctx"
+	"git.imall.cloud/openim/chat/pkg/protocol/bot"
+	"git.imall.cloud/openim/protocol/constant"
+	"github.com/openimsdk/tools/errs"
+	"github.com/sashabaranov/go-openai"
+)
+
+func (b *botSvr) SendBotMessage(ctx context.Context, req *bot.SendBotMessageReq) (*bot.SendBotMessageResp, error) {
+	agent, err := b.database.TakeAgent(ctx, req.AgentID)
+	if err != nil {
+		return nil, errs.ErrArgs.WrapMsg("agent not found")
+	}
+	//convRespID, err := b.database.TakeConversationRespID(ctx, req.ConversationID, req.AgentID)
+	//if err != nil && !errors.Is(err, mongo.ErrNoDocuments) {
+	//	return nil, err
+	//}
+	//var respID string
+	//if convRespID != nil {
+	//	respID = convRespID.PreviousResponseID
+	//}
+
+	aiCfg := openai.DefaultConfig(agent.Key)
+	aiCfg.BaseURL = agent.Url
+	aiCfg.HTTPClient = b.httpClient
+	client := openai.NewClientWithConfig(aiCfg)
+	aiReq := openai.ChatCompletionRequest{
+		Model: agent.Model,
+		Messages: []openai.ChatCompletionMessage{
+			{
+				Role:    openai.ChatMessageRoleSystem,
+				Content: agent.Prompts,
+			},
+			{
+				Role:    openai.ChatMessageRoleUser,
+				Content: req.Content,
+			},
+		},
+	}
+	aiCtx, cancel := context.WithTimeout(ctx, time.Duration(b.timeout)*time.Second)
+	defer cancel()
+	completion, err := client.CreateChatCompletion(aiCtx, aiReq)
+	if err != nil {
+		return nil, errs.Wrap(err)
+	}
+
+	imToken, err := b.imCaller.ImAdminTokenWithDefaultAdmin(ctx)
+	if err != nil {
+		return nil, err
+	}
+	ctx = mctx.WithApiToken(ctx, imToken)
+
+	content := "no response"
+	if len(completion.Choices) > 0 {
+		content = completion.Choices[0].Message.Content
+	}
+	err = b.imCaller.SendSimpleMsg(ctx, &imapi.SendSingleMsgReq{
+		SendID:  agent.UserID,
+		Content: content,
+	}, req.Key)
+	if err != nil {
+		return nil, err
+	}
+
+	//err = b.database.UpdateConversationRespID(ctx, req.ConversationID, agent.UserID, ToDBConversationRespIDUpdate(completion.ID))
+	//if err != nil {
+	//	return nil, err
+	//}
+	return &bot.SendBotMessageResp{}, nil
+}
+
+func getContent(contentType int32, content string) (string, error) {
+	switch contentType {
+	case constant.Text:
+		var elem botstruct.TextElem
+		err := json.Unmarshal([]byte(content), &elem)
+		if err != nil {
+			return "", errs.ErrArgs.WrapMsg(err.Error())
+		}
+		return elem.Content, nil
+	default:
+		return "", errs.New("un support contentType").Wrap()
+	}
+}

internal/rpc/bot/start.go

@@ -0,0 +1,53 @@
+package bot
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	"git.imall.cloud/openim/chat/pkg/common/config"
+	"git.imall.cloud/openim/chat/pkg/common/db/database"
+	"git.imall.cloud/openim/chat/pkg/common/imapi"
+	"git.imall.cloud/openim/chat/pkg/protocol/bot"
+	"github.com/openimsdk/tools/db/mongoutil"
+	"github.com/openimsdk/tools/discovery"
+	"google.golang.org/grpc"
+)
+
+type Config struct {
+	RpcConfig     config.Bot
+	RedisConfig   config.Redis
+	MongodbConfig config.Mongo
+	Discovery     config.Discovery
+	Share         config.Share
+}
+
+func Start(ctx context.Context, config *Config, client discovery.SvcDiscoveryRegistry, server *grpc.Server) error {
+	mgocli, err := mongoutil.NewMongoDB(ctx, config.MongodbConfig.Build())
+	if err != nil {
+		return err
+	}
+	var srv botSvr
+
+	srv.database, err = database.NewBotDatabase(mgocli)
+	if err != nil {
+		return err
+	}
+	srv.timeout = config.RpcConfig.Timeout
+	srv.httpClient = &http.Client{
+		Timeout: time.Duration(config.RpcConfig.Timeout) * time.Second,
+	}
+	im := imapi.New(config.Share.OpenIM.ApiURL, config.Share.OpenIM.Secret, config.Share.OpenIM.AdminUserID)
+	srv.imCaller = im
+	bot.RegisterBotServer(server, &srv)
+	return nil
+}
+
+type botSvr struct {
+	bot.UnimplementedBotServer
+	database   database.BotDatabase
+	httpClient *http.Client
+	timeout    int
+	imCaller   imapi.CallerInterface
+	//Admin           *chatClient.AdminClient
+}

internal/rpc/bot/update.go

@@ -0,0 +1,37 @@
+package bot
+
+import "git.imall.cloud/openim/chat/pkg/protocol/bot"
+
+func ToDBAgentUpdate(req *bot.UpdateAgentReq) map[string]any {
+	update := make(map[string]any)
+	if req.Key != nil {
+		update["key"] = req.Key
+	}
+	if req.Prompts != nil {
+		update["prompts"] = req.Prompts
+	}
+	if req.Model != nil {
+		update["model"] = req.Model
+	}
+	if req.FaceURL != nil {
+		update["face_url"] = req.FaceURL
+	}
+	if req.Nickname != nil {
+		update["nick_name"] = req.Nickname
+	}
+	if req.Identity != nil {
+		update["identity"] = req.Identity
+	}
+	if req.Url != nil {
+		update["url"] = req.Url
+	}
+
+	return update
+}
+
+func ToDBConversationRespIDUpdate(respID string) map[string]any {
+	update := map[string]any{
+		"previous_response_id": respID,
+	}
+	return update
+}