108 lines
3.0 KiB
Go
108 lines
3.0 KiB
Go
package server
|
||
|
||
import (
|
||
"ai_scheduler/internal/entitys"
|
||
"ai_scheduler/internal/services"
|
||
"context"
|
||
"fmt"
|
||
"sync"
|
||
|
||
"gitea.cdlsxd.cn/self-tools/l-dingtalk-stream-sdk-go/chatbot"
|
||
"gitea.cdlsxd.cn/self-tools/l-dingtalk-stream-sdk-go/client"
|
||
"github.com/go-kratos/kratos/v2/log"
|
||
)
|
||
|
||
type DingBotServiceInterface interface {
|
||
GetServiceCfg() ([]entitys.DingTalkBot, error)
|
||
OnChatBotMessageReceived(ctx context.Context, data *chatbot.BotCallbackDataModel) (content []byte, err error)
|
||
}
|
||
|
||
type DingTalkBotServer struct {
|
||
Clients map[string]*client.StreamClient
|
||
}
|
||
|
||
// NewDingTalkBotServer 批量注册钉钉客户端cli
|
||
// 这里支持两种方式，一种是完全独立service,一种是直接用现成的service
|
||
// 独立的service,在本页的ProvideAllDingBotServices方法进行注册
|
||
// 现成的service参考services->dtalk_bot.go
|
||
// 具体使用请根据实际业务需求
|
||
func NewDingTalkBotServer(
|
||
services []DingBotServiceInterface,
|
||
) *DingTalkBotServer {
|
||
clients := make(map[string]*client.StreamClient)
|
||
for _, service := range services {
|
||
serviceConfigs, err := service.GetServiceCfg()
|
||
for _, serviceConf := range serviceConfigs {
|
||
if serviceConf.ClientId == "" || serviceConf.ClientSecret == "" {
|
||
continue
|
||
}
|
||
cli := DingBotServerInit(serviceConf.ClientId, serviceConf.ClientSecret, service)
|
||
if cli == nil {
|
||
log.Info("%s客户端初始失败:%s", serviceConf.BotIndex, err.Error())
|
||
continue
|
||
}
|
||
clients[serviceConf.BotIndex] = cli
|
||
}
|
||
}
|
||
return &DingTalkBotServer{
|
||
Clients: clients,
|
||
}
|
||
}
|
||
|
||
func ProvideAllDingBotServices(
|
||
dingBotSvc *services.DingBotService,
|
||
) []DingBotServiceInterface {
|
||
return []DingBotServiceInterface{dingBotSvc}
|
||
}
|
||
|
||
func (d *DingTalkBotServer) Run(ctx context.Context, botIndex string) {
|
||
if botIndex == "" {
|
||
log.Info("未指定机器人索引，跳过启动")
|
||
return
|
||
}
|
||
|
||
var targets []string
|
||
switch {
|
||
case botIndex == "All":
|
||
targets = make([]string, 0, len(d.Clients))
|
||
for name := range d.Clients {
|
||
targets = append(targets, name)
|
||
}
|
||
default:
|
||
if _, exists := d.Clients[botIndex]; exists {
|
||
targets = []string{botIndex}
|
||
} else {
|
||
log.Infof("未找到索引为 %s 的机器人", botIndex)
|
||
return
|
||
}
|
||
}
|
||
|
||
var wg sync.WaitGroup
|
||
errors := make([]error, 0, len(targets))
|
||
|
||
for _, name := range targets {
|
||
wg.Add(1)
|
||
go func(name string) {
|
||
defer wg.Done()
|
||
err := d.Clients[name].Start(ctx)
|
||
if err != nil {
|
||
log.Errorf("%s 启动失败: %v", name, err)
|
||
errors = append(errors, fmt.Errorf("%s: %w", name, err))
|
||
} else {
|
||
log.Infof("%s 启动成功", name)
|
||
}
|
||
}(name)
|
||
}
|
||
|
||
wg.Wait()
|
||
if len(errors) > 0 {
|
||
log.Errorf("部分机器人启动失败，总数: %d, 成功: %d, 失败: %d",
|
||
len(targets), len(targets)-len(errors), len(errors))
|
||
}
|
||
}
|
||
func DingBotServerInit(clientId string, clientSecret string, service DingBotServiceInterface) (cli *client.StreamClient) {
|
||
cli = client.NewStreamClient(client.WithAppCredential(client.NewAppCredentialConfig(clientId, clientSecret)))
|
||
cli.RegisterChatBotCallbackRouter(service.OnChatBotMessageReceived)
|
||
return
|
||
}
|