重构：把剩余 12 个顶层包全部迁到 internal/
继续之前的重构：把根目录残留的非数据/资源类子目录（agents、agenttool、ai、 autoreply、completion、conversation、llm、message、mqtpp、stream、 toolmanager、toolrouter）一并搬到 internal/ 下，并把全工程引用它们的 import 路径从 "meshtastic_mqtt_server/<x>" 重写为 "meshtastic_mqtt_server/internal/<x>"。 变更 - git mv 12 个顶层目录进 internal/，无函数体改动。 - 全工程 sed 把 import path 加上 internal/ 前缀，包括 main.go、 internal/bot/bot_service.go、internal/store/bot_store.go 中对 mqtpp 的引用，以及 ai 子系统内部 agents/agenttool/autoreply/conversation/ llm/toolmanager/toolrouter 之间的相互引用。 验证 - go build ./... 通过；go test ./... 全部包通过。 - AI 自动回复链路（main → ai.NewService → autoreply.Service → botSender） 保持不变，仅 import 路径调整。 - 根目录最终只剩 main.go / main_test.go / internal/ 加 go.mod / go.sum 与数据资源目录（dist、doc、firmware、py、win、meshmap_frontend）。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Claude
@@ -0,0 +1,207 @@
|
||||
package calculator
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"go/ast"
|
||||
"go/parser"
|
||||
"go/token"
|
||||
"math"
|
||||
"strconv"
|
||||
|
||||
"meshtastic_mqtt_server/internal/agenttool"
|
||||
|
||||
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
||||
)
|
||||
|
||||
// Tool is a calculator tool that can evaluate simple math expressions
|
||||
type Tool struct {
|
||||
enabled bool
|
||||
}
|
||||
|
||||
// Name returns the tool name
|
||||
func (t *Tool) Name() string {
|
||||
return "calculator"
|
||||
}
|
||||
|
||||
// Enabled returns whether the tool is enabled
|
||||
func (t *Tool) Enabled() bool {
|
||||
return t.enabled
|
||||
}
|
||||
|
||||
// ToolDefinition returns the OpenAI tool definition
|
||||
func (t *Tool) ToolDefinition(description string) *model.Tool {
|
||||
desc := "一个计算器工具，可以计算数学表达式。支持加减乘除、平方根、幂运算等。当用户需要计算数学表达式时使用此工具。"
|
||||
if description != "" {
|
||||
desc = description
|
||||
}
|
||||
return &model.Tool{
|
||||
Type: model.ToolTypeFunction,
|
||||
Function: &model.FunctionDefinition{
|
||||
Name: "calculator",
|
||||
Description: desc,
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"expression": map[string]any{
|
||||
"type": "string",
|
||||
"description": "要计算的数学表达式，例如 \"2 + 3 * 4\" 或 \"sqrt(16)\"",
|
||||
},
|
||||
},
|
||||
"required": []string{"expression"},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the calculator tool
|
||||
func (t *Tool) Execute(ctx context.Context, args string, runtime agenttool.Runtime) (string, error) {
|
||||
var params struct {
|
||||
Expression string `json:"expression"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(args), ¶ms); err != nil {
|
||||
return "", fmt.Errorf("failed to parse arguments: %w", err)
|
||||
}
|
||||
|
||||
if params.Expression == "" {
|
||||
return "", fmt.Errorf("expression is required")
|
||||
}
|
||||
|
||||
result, err := evaluateExpression(params.Expression)
|
||||
if err != nil {
|
||||
return fmt.Sprintf("计算错误: %v", err), nil
|
||||
}
|
||||
|
||||
return fmt.Sprintf("计算结果: %s = %g", params.Expression, result), nil
|
||||
}
|
||||
|
||||
// RawState returns the tool state
|
||||
func (t *Tool) RawState() any {
|
||||
return map[string]any{"enabled": t.enabled}
|
||||
}
|
||||
|
||||
// evaluateExpression evaluates a simple mathematical expression using Go AST
|
||||
func evaluateExpression(expr string) (float64, error) {
|
||||
// Parse the expression
|
||||
node, err := parser.ParseExpr(expr)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("无效的表达式: %w", err)
|
||||
}
|
||||
|
||||
return evalNode(node)
|
||||
}
|
||||
|
||||
// evalNode evaluates an AST node
|
||||
func evalNode(node ast.Expr) (float64, error) {
|
||||
switch n := node.(type) {
|
||||
case *ast.BasicLit:
|
||||
if n.Kind == token.INT || n.Kind == token.FLOAT {
|
||||
return strconv.ParseFloat(n.Value, 64)
|
||||
}
|
||||
return 0, fmt.Errorf("不支持的字面量类型: %v", n.Kind)
|
||||
|
||||
case *ast.BinaryExpr:
|
||||
left, err := evalNode(n.X)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
right, err := evalNode(n.Y)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
switch n.Op {
|
||||
case token.ADD:
|
||||
return left + right, nil
|
||||
case token.SUB:
|
||||
return left - right, nil
|
||||
case token.MUL:
|
||||
return left * right, nil
|
||||
case token.QUO:
|
||||
if right == 0 {
|
||||
return 0, fmt.Errorf("除数不能为零")
|
||||
}
|
||||
return left / right, nil
|
||||
case token.REM:
|
||||
return math.Mod(left, right), nil
|
||||
default:
|
||||
return 0, fmt.Errorf("不支持的运算符: %v", n.Op)
|
||||
}
|
||||
|
||||
case *ast.ParenExpr:
|
||||
return evalNode(n.X)
|
||||
|
||||
case *ast.UnaryExpr:
|
||||
val, err := evalNode(n.X)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
switch n.Op {
|
||||
case token.ADD:
|
||||
return val, nil
|
||||
case token.SUB:
|
||||
return -val, nil
|
||||
default:
|
||||
return 0, fmt.Errorf("不支持的一元运算符: %v", n.Op)
|
||||
}
|
||||
|
||||
case *ast.CallExpr:
|
||||
ident, ok := n.Fun.(*ast.Ident)
|
||||
if !ok {
|
||||
return 0, fmt.Errorf("不支持的函数调用形式")
|
||||
}
|
||||
if len(n.Args) != 1 {
|
||||
return 0, fmt.Errorf("函数只接受一个参数")
|
||||
}
|
||||
arg, err := evalNode(n.Args[0])
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
switch ident.Name {
|
||||
case "sqrt":
|
||||
if arg < 0 {
|
||||
return 0, fmt.Errorf("平方根的参数不能为负数")
|
||||
}
|
||||
return math.Sqrt(arg), nil
|
||||
case "abs":
|
||||
return math.Abs(arg), nil
|
||||
case "sin":
|
||||
return math.Sin(arg), nil
|
||||
case "cos":
|
||||
return math.Cos(arg), nil
|
||||
case "tan":
|
||||
return math.Tan(arg), nil
|
||||
case "log":
|
||||
if arg <= 0 {
|
||||
return 0, fmt.Errorf("对数的参数必须为正数")
|
||||
}
|
||||
return math.Log(arg), nil
|
||||
case "log10":
|
||||
if arg <= 0 {
|
||||
return 0, fmt.Errorf("对数的参数必须为正数")
|
||||
}
|
||||
return math.Log10(arg), nil
|
||||
case "exp":
|
||||
return math.Exp(arg), nil
|
||||
case "pow":
|
||||
// For pow, we need two arguments - this is simplified
|
||||
return 0, fmt.Errorf("pow 函数需要两个参数，请使用 ** 运算符代替")
|
||||
default:
|
||||
return 0, fmt.Errorf("不支持的函数: %s", ident.Name)
|
||||
}
|
||||
|
||||
default:
|
||||
return 0, fmt.Errorf("不支持的表达式节点类型: %T", node)
|
||||
}
|
||||
}
|
||||
|
||||
func init() {
|
||||
agenttool.Register(agenttool.Descriptor{
|
||||
Name: "calculator",
|
||||
Load: func(path string, options agenttool.LoadOptions) (agenttool.LoadedTool, error) {
|
||||
return &Tool{enabled: true}, nil
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package time
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"meshtastic_mqtt_server/internal/agenttool"
|
||||
|
||||
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
||||
)
|
||||
|
||||
// Tool is a time tool that can get current time and format dates
|
||||
type Tool struct {
|
||||
enabled bool
|
||||
}
|
||||
|
||||
// Name returns the tool name
|
||||
func (t *Tool) Name() string {
|
||||
return "time"
|
||||
}
|
||||
|
||||
// Enabled returns whether the tool is enabled
|
||||
func (t *Tool) Enabled() bool {
|
||||
return t.enabled
|
||||
}
|
||||
|
||||
// ToolDefinition returns the OpenAI tool definition
|
||||
func (t *Tool) ToolDefinition(description string) *model.Tool {
|
||||
desc := "一个时间工具，可以获取当前时间、格式化日期时间、计算时间差等。当用户询问时间或需要处理日期时间相关问题时使用此工具。"
|
||||
if description != "" {
|
||||
desc = description
|
||||
}
|
||||
return &model.Tool{
|
||||
Type: model.ToolTypeFunction,
|
||||
Function: &model.FunctionDefinition{
|
||||
Name: "time",
|
||||
Description: desc,
|
||||
Parameters: map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"action": map[string]any{
|
||||
"type": "string",
|
||||
"enum": []string{"now", "format", "parse"},
|
||||
"description": "执行的操作: now(获取当前时间), format(格式化时间), parse(解析时间)",
|
||||
},
|
||||
"format": map[string]any{
|
||||
"type": "string",
|
||||
"description": "时间格式字符串，例如 \"2006-01-02 15:04:05\"",
|
||||
},
|
||||
"time": map[string]any{
|
||||
"type": "string",
|
||||
"description": "要解析的时间字符串",
|
||||
},
|
||||
"timezone": map[string]any{
|
||||
"type": "string",
|
||||
"description": "时区，例如 \"Asia/Shanghai\" 或 \"Local\"",
|
||||
},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the time tool
|
||||
func (t *Tool) Execute(ctx context.Context, args string, runtime agenttool.Runtime) (string, error) {
|
||||
var params struct {
|
||||
Action string `json:"action"`
|
||||
Format string `json:"format"`
|
||||
Time string `json:"time"`
|
||||
Timezone string `json:"timezone"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(args), ¶ms); err != nil {
|
||||
return "", fmt.Errorf("failed to parse arguments: %w", err)
|
||||
}
|
||||
|
||||
// Get timezone
|
||||
loc := time.Local
|
||||
if params.Timezone != "" {
|
||||
var err error
|
||||
loc, err = time.LoadLocation(params.Timezone)
|
||||
if err != nil {
|
||||
loc = time.Local
|
||||
}
|
||||
}
|
||||
|
||||
now := time.Now().In(loc)
|
||||
|
||||
switch params.Action {
|
||||
case "now":
|
||||
formatStr := time.RFC3339
|
||||
if params.Format != "" {
|
||||
formatStr = params.Format
|
||||
}
|
||||
return fmt.Sprintf("当前时间: %s", now.Format(formatStr)), nil
|
||||
|
||||
case "format":
|
||||
formatStr := time.RFC3339
|
||||
if params.Format != "" {
|
||||
formatStr = params.Format
|
||||
}
|
||||
return fmt.Sprintf("格式化时间: %s", now.Format(formatStr)), nil
|
||||
|
||||
case "parse":
|
||||
if params.Time == "" {
|
||||
return "", fmt.Errorf("time 参数是必需的")
|
||||
}
|
||||
formatStr := time.RFC3339
|
||||
if params.Format != "" {
|
||||
formatStr = params.Format
|
||||
}
|
||||
parsed, err := time.ParseInLocation(formatStr, params.Time, loc)
|
||||
if err != nil {
|
||||
return fmt.Sprintf("解析时间失败: %v", err), nil
|
||||
}
|
||||
return fmt.Sprintf("解析结果: %s (Unix 时间戳: %d)", parsed.Format(time.RFC3339), parsed.Unix()), nil
|
||||
|
||||
default:
|
||||
return fmt.Sprintf("不支持的操作: %s", params.Action), nil
|
||||
}
|
||||
}
|
||||
|
||||
// RawState returns the tool state
|
||||
func (t *Tool) RawState() any {
|
||||
return map[string]any{"enabled": t.enabled}
|
||||
}
|
||||
|
||||
func init() {
|
||||
agenttool.Register(agenttool.Descriptor{
|
||||
Name: "time",
|
||||
Load: func(path string, options agenttool.LoadOptions) (agenttool.LoadedTool, error) {
|
||||
return &Tool{enabled: true}, nil
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package agenttool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
||||
)
|
||||
|
||||
// LoadOptions contains options for loading a tool
|
||||
type LoadOptions struct {
|
||||
Values map[string]any
|
||||
}
|
||||
|
||||
// Value returns a value from the load options
|
||||
func (o LoadOptions) Value(key string) any {
|
||||
if o.Values == nil {
|
||||
return nil
|
||||
}
|
||||
return o.Values[key]
|
||||
}
|
||||
|
||||
// Frame represents an event frame emitted by a tool
|
||||
type Frame struct {
|
||||
Type string
|
||||
Tool string
|
||||
Stage string
|
||||
Status string
|
||||
Message string
|
||||
Data map[string]any
|
||||
Error string
|
||||
Text string
|
||||
}
|
||||
|
||||
// EmitFunc is a function that emits frames
|
||||
type EmitFunc func(frame any)
|
||||
|
||||
// Runtime provides context for tool execution
|
||||
type Runtime struct {
|
||||
Profile any
|
||||
CompleteText func(context.Context, string, int) (string, error)
|
||||
Emit EmitFunc
|
||||
Now time.Time
|
||||
}
|
||||
|
||||
// LoadedTool is the interface that all tools must implement
|
||||
type LoadedTool interface {
|
||||
Name() string
|
||||
Enabled() bool
|
||||
ToolDefinition(description string) *model.Tool
|
||||
Execute(context.Context, string, Runtime) (string, error)
|
||||
RawState() any
|
||||
}
|
||||
|
||||
// Descriptor describes a tool that can be loaded
|
||||
type Descriptor struct {
|
||||
Name string
|
||||
Load func(path string, options LoadOptions) (LoadedTool, error)
|
||||
}
|
||||
|
||||
var (
|
||||
registryMu sync.RWMutex
|
||||
registry = map[string]Descriptor{}
|
||||
)
|
||||
|
||||
// Register registers a tool descriptor
|
||||
func Register(descriptor Descriptor) {
|
||||
name := strings.ToLower(strings.TrimSpace(descriptor.Name))
|
||||
if name == "" {
|
||||
panic("agenttool: tool name is empty")
|
||||
}
|
||||
if descriptor.Load == nil {
|
||||
panic(fmt.Sprintf("agenttool: %s load function is nil", name))
|
||||
}
|
||||
descriptor.Name = name
|
||||
|
||||
registryMu.Lock()
|
||||
defer registryMu.Unlock()
|
||||
if _, ok := registry[name]; ok {
|
||||
panic(fmt.Sprintf("agenttool: tool %s already registered", name))
|
||||
}
|
||||
registry[name] = descriptor
|
||||
}
|
||||
|
||||
// Lookup finds a tool descriptor by name
|
||||
func Lookup(name string) (Descriptor, bool) {
|
||||
registryMu.RLock()
|
||||
defer registryMu.RUnlock()
|
||||
descriptor, ok := registry[strings.ToLower(strings.TrimSpace(name))]
|
||||
return descriptor, ok
|
||||
}
|
||||
|
||||
// Names returns all registered tool names
|
||||
func Names() []string {
|
||||
registryMu.RLock()
|
||||
defer registryMu.RUnlock()
|
||||
|
||||
names := make([]string, 0, len(registry))
|
||||
for name := range registry {
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
return names
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"meshtastic_mqtt_server/internal/agenttool"
|
||||
_ "meshtastic_mqtt_server/internal/agents/calculator"
|
||||
_ "meshtastic_mqtt_server/internal/agents/time"
|
||||
"meshtastic_mqtt_server/internal/autoreply"
|
||||
"meshtastic_mqtt_server/internal/conversation"
|
||||
"meshtastic_mqtt_server/internal/llm"
|
||||
"meshtastic_mqtt_server/internal/toolmanager"
|
||||
"meshtastic_mqtt_server/internal/toolrouter"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ToolConfigStore is the interface for getting tool configuration
|
||||
type ToolConfigStore interface {
|
||||
GetLLMPrimaryConfigSystemPrompt() (string, error)
|
||||
GetLLMPrimaryConfigEnableTool() (bool, error)
|
||||
}
|
||||
|
||||
// Config holds the AI service configuration
|
||||
type Config struct {
|
||||
LLMProviders []llm.ProviderConfig
|
||||
DataDir string
|
||||
Enabled bool
|
||||
ToolConfigStore ToolConfigStore
|
||||
}
|
||||
|
||||
// Service manages all AI-related components
|
||||
type Service struct {
|
||||
LLMState *llm.State
|
||||
ToolRouter *toolrouter.State
|
||||
ToolMgr *toolmanager.Manager
|
||||
ConvStore *conversation.Store
|
||||
AutoReply *autoreply.Service
|
||||
MsgQueue *autoreply.DBMessageQueue
|
||||
|
||||
enabled bool
|
||||
}
|
||||
|
||||
// NewService creates a new AI service
|
||||
func NewService(cfg Config, db *gorm.DB, botSender autoreply.BotSender) (*Service, error) {
|
||||
if !cfg.Enabled {
|
||||
return &Service{enabled: false}, nil
|
||||
}
|
||||
|
||||
// Create data directories
|
||||
agentsDir := filepath.Join(cfg.DataDir, "agents")
|
||||
convDir := filepath.Join(cfg.DataDir, "conversations")
|
||||
if err := os.MkdirAll(agentsDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("failed to create agents directory: %w", err)
|
||||
}
|
||||
if err := os.MkdirAll(convDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("failed to create conversations directory: %w", err)
|
||||
}
|
||||
|
||||
// Initialize LLM state
|
||||
llmState, err := llm.NewState(cfg.LLMProviders)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to initialize LLM state: %w", err)
|
||||
}
|
||||
|
||||
// Initialize tool router
|
||||
toolRouterCfg := &toolrouter.Config{
|
||||
Enabled: true,
|
||||
Timeout: 30,
|
||||
MaxTokens: 512,
|
||||
SystemPrompt: "你是一个智能助手，可以调用工具来回答用户问题。\n用户正在通过 Mesh 网络与你对话，请保持回答简洁明了。\n工具结果优先于模型内置知识。",
|
||||
}
|
||||
toolRouter, err := toolrouter.NewState(toolRouterCfg, llmState)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to initialize tool router: %w", err)
|
||||
}
|
||||
|
||||
// Load tools
|
||||
toolMgr, err := toolmanager.Load(agentsDir, agenttool.LoadOptions{})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to load tools: %w", err)
|
||||
}
|
||||
|
||||
// Initialize conversation store
|
||||
convStore := conversation.NewStore(convDir)
|
||||
|
||||
// Initialize message queue
|
||||
msgQueue := autoreply.NewDBMessageQueue(db)
|
||||
|
||||
// Initialize auto-reply service
|
||||
autoReply := autoreply.NewService(
|
||||
llmState,
|
||||
toolRouter,
|
||||
toolMgr,
|
||||
convStore,
|
||||
msgQueue,
|
||||
botSender,
|
||||
cfg.ToolConfigStore,
|
||||
)
|
||||
|
||||
return &Service{
|
||||
LLMState: llmState,
|
||||
ToolRouter: toolRouter,
|
||||
ToolMgr: toolMgr,
|
||||
ConvStore: convStore,
|
||||
AutoReply: autoReply,
|
||||
MsgQueue: msgQueue,
|
||||
enabled: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Start starts the AI service
|
||||
func (s *Service) Start(ctx context.Context) error {
|
||||
if !s.enabled {
|
||||
return nil
|
||||
}
|
||||
return s.AutoReply.Start(ctx)
|
||||
}
|
||||
|
||||
// Stop stops the AI service
|
||||
func (s *Service) Stop() {
|
||||
if !s.enabled {
|
||||
return
|
||||
}
|
||||
s.AutoReply.Stop()
|
||||
s.ToolMgr.Close()
|
||||
}
|
||||
|
||||
// Enabled returns whether the AI service is enabled
|
||||
func (s *Service) Enabled() bool {
|
||||
return s.enabled
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package autoreply
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// BotServiceAdapter adapts the bot service to the BotSender interface
|
||||
type BotServiceAdapter struct {
|
||||
sendDirectTextFn func(ctx context.Context, botID uint64, toNodeNum int64, text string) error
|
||||
sendChannelTextFn func(ctx context.Context, botID uint64, channelID string, text string) error
|
||||
}
|
||||
|
||||
// NewBotServiceAdapter creates a new bot service adapter
|
||||
func NewBotServiceAdapter(
|
||||
sendDirectTextFn func(ctx context.Context, botID uint64, toNodeNum int64, text string) error,
|
||||
sendChannelTextFn func(ctx context.Context, botID uint64, channelID string, text string) error,
|
||||
) *BotServiceAdapter {
|
||||
return &BotServiceAdapter{
|
||||
sendDirectTextFn: sendDirectTextFn,
|
||||
sendChannelTextFn: sendChannelTextFn,
|
||||
}
|
||||
}
|
||||
|
||||
// SendDirectText sends a direct/private message to a specific node
|
||||
func (a *BotServiceAdapter) SendDirectText(ctx context.Context, botID uint64, toNodeNum int64, text string) error {
|
||||
if a.sendDirectTextFn == nil {
|
||||
return fmt.Errorf("send direct text function is nil")
|
||||
}
|
||||
return a.sendDirectTextFn(ctx, botID, toNodeNum, text)
|
||||
}
|
||||
|
||||
// SendChannelText sends a channel message to a specific channel
|
||||
func (a *BotServiceAdapter) SendChannelText(ctx context.Context, botID uint64, channelID string, text string) error {
|
||||
if a.sendChannelTextFn == nil {
|
||||
return fmt.Errorf("send channel text function is nil")
|
||||
}
|
||||
return a.sendChannelTextFn(ctx, botID, channelID, text)
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package autoreply
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const (
|
||||
statusPending = "pending"
|
||||
statusProcessing = "processing"
|
||||
statusProcessed = "processed"
|
||||
statusFailed = "error"
|
||||
)
|
||||
|
||||
// DBMessageQueue implements MessageQueue using GORM
|
||||
type DBMessageQueue struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewDBMessageQueue creates a new database-backed message queue
|
||||
func NewDBMessageQueue(db *gorm.DB) *DBMessageQueue {
|
||||
return &DBMessageQueue{db: db}
|
||||
}
|
||||
|
||||
// llmMessageQueueRecord is the database record for LLM messages
|
||||
type llmMessageQueueRecord struct {
|
||||
ID uint64 `gorm:"column:id;primaryKey;autoIncrement"`
|
||||
BotID uint64 `gorm:"column:bot_id;not null;index"`
|
||||
BotNodeID string `gorm:"column:bot_node_id;not null"`
|
||||
BotNodeNum int64 `gorm:"column:bot_node_num;not null"`
|
||||
FromNodeID string `gorm:"column:from_node_id;not null"`
|
||||
FromNodeNum int64 `gorm:"column:from_node_num;not null"`
|
||||
LongName *string `gorm:"column:long_name"`
|
||||
ShortName *string `gorm:"column:short_name"`
|
||||
Text string `gorm:"column:text;type:text;not null"`
|
||||
PacketID int64 `gorm:"column:packet_id;not null"`
|
||||
ChannelID *string `gorm:"column:channel_id"`
|
||||
Topic string `gorm:"column:topic;not null"`
|
||||
MessageType string `gorm:"column:message_type;not null;default:'direct'"` // "channel" 或 "direct"
|
||||
Status string `gorm:"column:status;not null;index"`
|
||||
Error string `gorm:"column:error;type:text"`
|
||||
Reply string `gorm:"column:reply;type:text"`
|
||||
ReceivedAt time.Time `gorm:"column:received_at;not null"`
|
||||
ProcessedAt *time.Time `gorm:"column:processed_at;index"`
|
||||
CreatedAt time.Time `gorm:"column:created_at;autoCreateTime;index"`
|
||||
}
|
||||
|
||||
func (llmMessageQueueRecord) TableName() string {
|
||||
return "llm_message_queue"
|
||||
}
|
||||
|
||||
// GetPendingMessages returns pending messages from the queue
|
||||
func (q *DBMessageQueue) GetPendingMessages(botID uint64, limit int) ([]QueuedMessage, error) {
|
||||
var records []llmMessageQueueRecord
|
||||
query := q.db.Where("status = ?", statusPending).Order("created_at ASC")
|
||||
if botID > 0 {
|
||||
query = query.Where("bot_id = ?", botID)
|
||||
}
|
||||
if limit > 0 {
|
||||
query = query.Limit(limit)
|
||||
}
|
||||
if err := query.Find(&records).Error; err != nil {
|
||||
return nil, fmt.Errorf("failed to query pending messages: %w", err)
|
||||
}
|
||||
|
||||
messages := make([]QueuedMessage, 0, len(records))
|
||||
for _, r := range records {
|
||||
messages = append(messages, QueuedMessage{
|
||||
ID: r.ID,
|
||||
BotID: r.BotID,
|
||||
BotNodeID: r.BotNodeID,
|
||||
BotNodeNum: r.BotNodeNum,
|
||||
FromNodeID: r.FromNodeID,
|
||||
FromNodeNum: r.FromNodeNum,
|
||||
LongName: r.LongName,
|
||||
ShortName: r.ShortName,
|
||||
Text: r.Text,
|
||||
PacketID: r.PacketID,
|
||||
ChannelID: r.ChannelID,
|
||||
Topic: r.Topic,
|
||||
MessageType: r.MessageType,
|
||||
ReceivedAt: r.ReceivedAt,
|
||||
})
|
||||
}
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
// MarkAsProcessing marks a message as being processed
|
||||
func (q *DBMessageQueue) MarkAsProcessing(id uint64) error {
|
||||
return q.db.Model(&llmMessageQueueRecord{}).Where("id = ?", id).Update("status", statusProcessing).Error
|
||||
}
|
||||
|
||||
// MarkAsProcessed marks a message as successfully processed
|
||||
func (q *DBMessageQueue) MarkAsProcessed(id uint64, reply string) error {
|
||||
now := time.Now()
|
||||
return q.db.Model(&llmMessageQueueRecord{}).Where("id = ?", id).Updates(map[string]any{
|
||||
"status": statusProcessed,
|
||||
"reply": reply,
|
||||
"processed_at": &now,
|
||||
}).Error
|
||||
}
|
||||
|
||||
// MarkAsFailed marks a message as failed
|
||||
func (q *DBMessageQueue) MarkAsFailed(id uint64, error string) error {
|
||||
return q.db.Model(&llmMessageQueueRecord{}).Where("id = ?", id).Updates(map[string]any{
|
||||
"status": statusFailed,
|
||||
"error": error,
|
||||
}).Error
|
||||
}
|
||||
@@ -0,0 +1,480 @@
|
||||
package autoreply
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"meshtastic_mqtt_server/internal/completion"
|
||||
"meshtastic_mqtt_server/internal/conversation"
|
||||
"meshtastic_mqtt_server/internal/llm"
|
||||
"meshtastic_mqtt_server/internal/message"
|
||||
"meshtastic_mqtt_server/internal/toolmanager"
|
||||
"meshtastic_mqtt_server/internal/toolrouter"
|
||||
|
||||
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
||||
)
|
||||
|
||||
const (
|
||||
// MaxReplyLength is the maximum length for Meshtastic messages
|
||||
MaxReplyLength = 200
|
||||
// PollInterval is how often to check the queue for new messages
|
||||
PollInterval = 5 * time.Second
|
||||
// MaxProcessingTime is the maximum time to spend processing a single message
|
||||
MaxProcessingTime = 120 * time.Second
|
||||
)
|
||||
|
||||
// MessageQueue is the interface for accessing the LLM message queue
|
||||
type MessageQueue interface {
|
||||
// GetPendingMessages returns pending messages for a bot
|
||||
GetPendingMessages(botID uint64, limit int) ([]QueuedMessage, error)
|
||||
// MarkAsProcessing marks a message as being processed
|
||||
MarkAsProcessing(id uint64) error
|
||||
// MarkAsProcessed marks a message as successfully processed
|
||||
MarkAsProcessed(id uint64, reply string) error
|
||||
// MarkAsFailed marks a message as failed
|
||||
MarkAsFailed(id uint64, error string) error
|
||||
}
|
||||
|
||||
// QueuedMessage represents a message in the LLM queue
|
||||
type QueuedMessage struct {
|
||||
ID uint64
|
||||
BotID uint64
|
||||
BotNodeID string
|
||||
BotNodeNum int64
|
||||
FromNodeID string
|
||||
FromNodeNum int64
|
||||
LongName *string
|
||||
ShortName *string
|
||||
Text string
|
||||
PacketID int64
|
||||
ChannelID *string
|
||||
Topic string
|
||||
MessageType string // "channel" 或 "direct"
|
||||
ReceivedAt time.Time
|
||||
}
|
||||
|
||||
// BotSender is the interface for sending bot messages
|
||||
type BotSender interface {
|
||||
// SendDirectText 发送私聊消息给指定节点
|
||||
SendDirectText(ctx context.Context, botID uint64, toNodeNum int64, text string) error
|
||||
// SendChannelText 发送频道消息到指定频道
|
||||
SendChannelText(ctx context.Context, botID uint64, channelID string, text string) error
|
||||
}
|
||||
|
||||
// ToolConfigStore is the interface for getting tool configuration
|
||||
type ToolConfigStore interface {
|
||||
GetLLMPrimaryConfigSystemPrompt() (string, error)
|
||||
GetLLMPrimaryConfigEnableTool() (bool, error)
|
||||
}
|
||||
|
||||
// Service manages automatic AI replies for bots
|
||||
type Service struct {
|
||||
llmState *llm.State
|
||||
toolRouter *toolrouter.State
|
||||
toolMgr *toolmanager.Manager
|
||||
convStore *conversation.Store
|
||||
msgQueue MessageQueue
|
||||
botSender BotSender
|
||||
toolConfigStore ToolConfigStore
|
||||
|
||||
running bool
|
||||
mu sync.Mutex
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewService creates a new auto-reply service
|
||||
func NewService(
|
||||
llmState *llm.State,
|
||||
toolRouter *toolrouter.State,
|
||||
toolMgr *toolmanager.Manager,
|
||||
convStore *conversation.Store,
|
||||
msgQueue MessageQueue,
|
||||
botSender BotSender,
|
||||
toolConfigStore ToolConfigStore,
|
||||
) *Service {
|
||||
return &Service{
|
||||
llmState: llmState,
|
||||
toolRouter: toolRouter,
|
||||
toolMgr: toolMgr,
|
||||
convStore: convStore,
|
||||
msgQueue: msgQueue,
|
||||
botSender: botSender,
|
||||
toolConfigStore: toolConfigStore,
|
||||
}
|
||||
}
|
||||
|
||||
// Start starts the auto-reply service
|
||||
func (s *Service) Start(ctx context.Context) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.running {
|
||||
return fmt.Errorf("auto-reply service is already running")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
s.cancel = cancel
|
||||
s.running = true
|
||||
|
||||
s.wg.Add(1)
|
||||
go s.run(ctx)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop stops the auto-reply service
|
||||
func (s *Service) Stop() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if !s.running {
|
||||
return
|
||||
}
|
||||
|
||||
if s.cancel != nil {
|
||||
s.cancel()
|
||||
}
|
||||
s.wg.Wait()
|
||||
s.running = false
|
||||
}
|
||||
|
||||
// run is the main processing loop
|
||||
func (s *Service) run(ctx context.Context) {
|
||||
defer s.wg.Done()
|
||||
|
||||
ticker := time.NewTicker(PollInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
s.processQueue(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// processQueue processes pending messages in the queue
|
||||
func (s *Service) processQueue(ctx context.Context) {
|
||||
// Get all bots (we'd typically get this from bot store, but for now
|
||||
// we'll rely on the queue to provide messages per bot)
|
||||
// For now, process up to 10 messages at a time
|
||||
messages, err := s.msgQueue.GetPendingMessages(0, 10)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
for _, msg := range messages {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
s.processMessage(ctx, msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// printJSON outputs a structured log message (imported from main package pattern)
|
||||
func printJSON(v any) {
|
||||
fmt.Printf("%+v\n", v)
|
||||
}
|
||||
|
||||
// processMessage processes a single queued message
|
||||
func (s *Service) processMessage(ctx context.Context, msg QueuedMessage) {
|
||||
// Mark message as processing
|
||||
if err := s.msgQueue.MarkAsProcessing(msg.ID); err != nil {
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_process_failed",
|
||||
"msg_id": msg.ID,
|
||||
"step": "mark_as_processing",
|
||||
"error": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_process_start",
|
||||
"msg_id": msg.ID,
|
||||
"bot_id": msg.BotID,
|
||||
"from_node_id": msg.FromNodeID,
|
||||
"text": msg.Text,
|
||||
})
|
||||
|
||||
// Create processing context with timeout
|
||||
procCtx, cancel := context.WithTimeout(ctx, MaxProcessingTime)
|
||||
defer cancel()
|
||||
|
||||
// Get or create conversation for this bot
|
||||
conv, err := s.convStore.GetOrCreateForBot(msg.BotID, msg.BotNodeID, msg.FromNodeID)
|
||||
if err != nil {
|
||||
errMsg := fmt.Sprintf("failed to get conversation: %v", err)
|
||||
printJSON(map[string]any{"event": "llm_process_failed", "msg_id": msg.ID, "step": "get_conversation", "error": errMsg})
|
||||
_ = s.msgQueue.MarkAsFailed(msg.ID, errMsg)
|
||||
return
|
||||
}
|
||||
|
||||
// Add the user message to the conversation
|
||||
userMsg := message.ChatMessage{
|
||||
Role: "user",
|
||||
Content: s.formatUserMessage(msg),
|
||||
}
|
||||
if err := s.convStore.AddMessage(conv.ID, userMsg); err != nil {
|
||||
errMsg := fmt.Sprintf("failed to add message: %v", err)
|
||||
printJSON(map[string]any{"event": "llm_process_failed", "msg_id": msg.ID, "step": "add_message", "error": errMsg})
|
||||
_ = s.msgQueue.MarkAsFailed(msg.ID, errMsg)
|
||||
return
|
||||
}
|
||||
|
||||
// Get the LLM profile
|
||||
profile := s.llmState.ActiveProfile()
|
||||
if profile == nil {
|
||||
errMsg := "no active LLM profile - check if LLM providers are configured"
|
||||
printJSON(map[string]any{"event": "llm_process_failed", "msg_id": msg.ID, "step": "get_profile", "error": errMsg})
|
||||
_ = s.msgQueue.MarkAsFailed(msg.ID, errMsg)
|
||||
return
|
||||
}
|
||||
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_process_profile",
|
||||
"msg_id": msg.ID,
|
||||
"model": profile.Config.Model,
|
||||
"base_url": profile.Config.BaseURL,
|
||||
})
|
||||
|
||||
// Reload conversation to get updated messages
|
||||
conv, err = s.convStore.Get(conv.ID)
|
||||
if err != nil {
|
||||
errMsg := fmt.Sprintf("failed to reload conversation: %v", err)
|
||||
printJSON(map[string]any{"event": "llm_process_failed", "msg_id": msg.ID, "step": "reload_conversation", "error": errMsg})
|
||||
_ = s.msgQueue.MarkAsFailed(msg.ID, errMsg)
|
||||
return
|
||||
}
|
||||
|
||||
// Get system prompt and tool enable flag from primary config
|
||||
var systemPrompt string
|
||||
enableTool := false
|
||||
if s.toolConfigStore != nil {
|
||||
systemPrompt, err = s.toolConfigStore.GetLLMPrimaryConfigSystemPrompt()
|
||||
if err != nil {
|
||||
printJSON(map[string]any{"event": "llm_system_prompt_warning", "msg_id": msg.ID, "error": err.Error()})
|
||||
}
|
||||
enableTool, err = s.toolConfigStore.GetLLMPrimaryConfigEnableTool()
|
||||
if err != nil {
|
||||
printJSON(map[string]any{"event": "llm_enable_tool_warning", "msg_id": msg.ID, "error": err.Error()})
|
||||
}
|
||||
}
|
||||
|
||||
// Print tool manager status for debugging
|
||||
toolCount := 0
|
||||
if s.toolMgr != nil {
|
||||
tools := s.toolMgr.Tools()
|
||||
toolCount = len(tools)
|
||||
toolNames := make([]string, 0, toolCount)
|
||||
for _, t := range tools {
|
||||
toolNames = append(toolNames, t.Name())
|
||||
}
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_tool_manager_status",
|
||||
"msg_id": msg.ID,
|
||||
"tool_count": toolCount,
|
||||
"tool_names": toolNames,
|
||||
"enable_tool": enableTool,
|
||||
})
|
||||
}
|
||||
|
||||
// Run the tool loop to get augmented messages - pass system prompt to tool router
|
||||
// Tool loop will handle system prompt and tool calling
|
||||
var augmentedMessages []*model.ChatCompletionMessage
|
||||
if enableTool && toolCount > 0 {
|
||||
augmentedMessages, err = toolrouter.RunAgentToolLoop(procCtx, s.toolRouter, profile, systemPrompt, conv.Messages, s.toolMgr, nil)
|
||||
if err != nil {
|
||||
printJSON(map[string]any{"event": "llm_tool_loop_warning", "msg_id": msg.ID, "error": err.Error()})
|
||||
// Continue with original messages if tool loop fails
|
||||
}
|
||||
}
|
||||
|
||||
printJSON(map[string]any{"event": "llm_process_completion_start", "msg_id": msg.ID, "has_system_prompt": systemPrompt != "", "augmented_messages": len(augmentedMessages)})
|
||||
|
||||
// Use augmented messages from tool loop (already includes system prompt and tool results)
|
||||
// If augmented messages is empty or nil, fallback to original messages with system prompt
|
||||
var reply string
|
||||
if len(augmentedMessages) > 0 {
|
||||
// Use augmented messages from tool loop (already converted to model.ChatCompletionMessage)
|
||||
reply, err = completion.CompleteTextWithArkMessages(procCtx, profile, augmentedMessages, 512)
|
||||
} else {
|
||||
// Fallback to simple completion
|
||||
reply, err = completion.CompleteText(procCtx, profile, systemPrompt, conv.Messages, 512)
|
||||
}
|
||||
if err != nil {
|
||||
errMsg := fmt.Sprintf("LLM completion failed: %v", err)
|
||||
printJSON(map[string]any{"event": "llm_process_failed", "msg_id": msg.ID, "step": "llm_completion", "error": errMsg})
|
||||
_ = s.msgQueue.MarkAsFailed(msg.ID, errMsg)
|
||||
return
|
||||
}
|
||||
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_process_completion_success",
|
||||
"msg_id": msg.ID,
|
||||
"reply_len": len(reply),
|
||||
})
|
||||
|
||||
// Clean and validate reply text
|
||||
reply = cleanReplyText(reply)
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_process_text_cleaned",
|
||||
"msg_id": msg.ID,
|
||||
"cleaned_len": len(reply),
|
||||
})
|
||||
|
||||
// Truncate reply for Meshtastic (UTF-8 safe truncation)
|
||||
if len([]byte(reply)) > MaxReplyLength {
|
||||
reply = truncateUTF8(reply, MaxReplyLength-3) + "..."
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_process_text_truncated",
|
||||
"msg_id": msg.ID,
|
||||
"truncated_len": len(reply),
|
||||
})
|
||||
}
|
||||
|
||||
// Final UTF-8 validation before sending
|
||||
if !utf8.ValidString(reply) {
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_process_utf8_warning",
|
||||
"msg_id": msg.ID,
|
||||
"message": "final text still invalid, using fallback",
|
||||
})
|
||||
reply = "抱歉，我暂时无法回复。请稍后再试。"
|
||||
}
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_process_final_check",
|
||||
"msg_id": msg.ID,
|
||||
"valid_utf8": utf8.ValidString(reply),
|
||||
"final_len": len(reply),
|
||||
})
|
||||
|
||||
// Add assistant reply to conversation
|
||||
assistantMsg := message.ChatMessage{
|
||||
Role: "assistant",
|
||||
Content: reply,
|
||||
}
|
||||
if err := s.convStore.AddMessage(conv.ID, assistantMsg); err != nil {
|
||||
// Non-fatal, continue
|
||||
}
|
||||
|
||||
// Send the reply via the bot - 根据消息类型决定发送方式
|
||||
// 频道消息：回复到原频道；私聊消息：回复给发送节点
|
||||
var sendErr error
|
||||
if msg.MessageType == "channel" && msg.ChannelID != nil && *msg.ChannelID != "" {
|
||||
// 频道消息 - 回复到原频道
|
||||
printJSON(map[string]any{"event": "llm_process_send_start", "msg_id": msg.ID, "channel_id": *msg.ChannelID, "message_type": "channel"})
|
||||
sendErr = s.botSender.SendChannelText(procCtx, msg.BotID, *msg.ChannelID, reply)
|
||||
} else {
|
||||
// 私聊消息 - 回复给发送节点
|
||||
printJSON(map[string]any{"event": "llm_process_send_start", "msg_id": msg.ID, "to_node_num": msg.FromNodeNum, "message_type": "direct"})
|
||||
sendErr = s.botSender.SendDirectText(procCtx, msg.BotID, msg.FromNodeNum, reply)
|
||||
}
|
||||
if sendErr != nil {
|
||||
errMsg := fmt.Sprintf("failed to send reply: %v", sendErr)
|
||||
printJSON(map[string]any{"event": "llm_process_failed", "msg_id": msg.ID, "step": "send_reply", "error": errMsg})
|
||||
_ = s.msgQueue.MarkAsFailed(msg.ID, errMsg)
|
||||
return
|
||||
}
|
||||
|
||||
// Mark message as processed
|
||||
_ = s.msgQueue.MarkAsProcessed(msg.ID, reply)
|
||||
printJSON(map[string]any{
|
||||
"event": "llm_process_success",
|
||||
"msg_id": msg.ID,
|
||||
"reply": reply,
|
||||
})
|
||||
}
|
||||
|
||||
// formatUserMessage formats the incoming message for the LLM
|
||||
func (s *Service) formatUserMessage(msg QueuedMessage) string {
|
||||
var sb strings.Builder
|
||||
|
||||
if msg.LongName != nil && *msg.LongName != "" {
|
||||
sb.WriteString(fmt.Sprintf("[来自 %s (%s)] ", *msg.LongName, msg.FromNodeID))
|
||||
} else if msg.ShortName != nil && *msg.ShortName != "" {
|
||||
sb.WriteString(fmt.Sprintf("[来自 %s (%s)] ", *msg.ShortName, msg.FromNodeID))
|
||||
} else {
|
||||
sb.WriteString(fmt.Sprintf("[来自 %s] ", msg.FromNodeID))
|
||||
}
|
||||
|
||||
sb.WriteString(msg.Text)
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// cleanReplyText cleans LLM reply text to ensure it's valid for Meshtastic
|
||||
func cleanReplyText(text string) string {
|
||||
// Force convert to valid UTF-8 using rune-by-rune processing
|
||||
v := make([]rune, 0, len(text))
|
||||
for i, r := range text {
|
||||
if r == utf8.RuneError {
|
||||
_, size := utf8.DecodeRuneInString(text[i:])
|
||||
if size == 1 {
|
||||
continue // Skip invalid runes
|
||||
}
|
||||
}
|
||||
// Skip all problematic characters
|
||||
if r < 32 && r != '\n' && r != '\r' && r != '\t' {
|
||||
continue
|
||||
}
|
||||
if r == 65533 || r == 0xfffd { // Unicode replacement character
|
||||
continue
|
||||
}
|
||||
v = append(v, r)
|
||||
}
|
||||
text = string(v)
|
||||
|
||||
// Additional cleanup - use only explicitly allowed ASCII printable + CJK
|
||||
var sb strings.Builder
|
||||
for _, r := range text {
|
||||
switch {
|
||||
case r == '\n':
|
||||
sb.WriteRune(' ') // Replace newlines with space for Meshtastic
|
||||
case r == '\r' || r == '\t':
|
||||
continue
|
||||
case r >= 32 && r <= 126: // Printable ASCII
|
||||
sb.WriteRune(r)
|
||||
case r >= 0x4E00 && r <= 0x9FFF: // CJK Unified Ideographs
|
||||
sb.WriteRune(r)
|
||||
case r >= 0x3400 && r <= 0x4DBF: // CJK Extension A
|
||||
sb.WriteRune(r)
|
||||
case r >= 0x20000 && r <= 0x2A6DF: // CJK Extension B
|
||||
sb.WriteRune(r)
|
||||
case r >= 0xFF01 && r <= 0xFF5E: // Fullwidth ASCII variants
|
||||
sb.WriteRune(r)
|
||||
case r == 0x3002 || r == 0xFF1F || r == 0xFF01 || r == 0xFF0C || r == 0xFF1A: // Fullwidth punctuation
|
||||
sb.WriteRune(r)
|
||||
default:
|
||||
continue // Skip all other characters
|
||||
}
|
||||
}
|
||||
|
||||
result := strings.TrimSpace(sb.String())
|
||||
if result == "" {
|
||||
result = "抱歉，我无法回复此消息。"
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// truncateUTF8 safely truncates a UTF-8 string to max bytes without breaking in the middle of a rune
|
||||
func truncateUTF8(s string, maxBytes int) string {
|
||||
if len([]byte(s)) <= maxBytes {
|
||||
return s
|
||||
}
|
||||
bytes := []byte(s)
|
||||
for i := maxBytes; i > 0; i-- {
|
||||
if utf8.RuneStart(bytes[i]) {
|
||||
return string(bytes[:i])
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -16,7 +16,7 @@ import (
|
||||
"gorm.io/gorm"
|
||||
|
||||
storepkg "meshtastic_mqtt_server/internal/store"
|
||||
"meshtastic_mqtt_server/mqtpp"
|
||||
"meshtastic_mqtt_server/internal/mqtpp"
|
||||
)
|
||||
|
||||
const botMaxTextBytes = 200
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
package completion
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"meshtastic_mqtt_server/internal/llm"
|
||||
"meshtastic_mqtt_server/internal/message"
|
||||
|
||||
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
||||
)
|
||||
|
||||
// ChatCompleter is a function that completes a chat
|
||||
type ChatCompleter func(ctx context.Context, profile *llm.Profile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error)
|
||||
|
||||
// CompleteChat completes a chat conversation
|
||||
func CompleteChat(ctx context.Context, profile *llm.Profile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
|
||||
if profile == nil || profile.Client == nil {
|
||||
return model.ChatCompletionResponse{}, fmt.Errorf("llm profile or client is nil")
|
||||
}
|
||||
|
||||
// Use context with timeout if provided
|
||||
if timeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
}
|
||||
|
||||
resp, err := profile.Client.CreateChatCompletion(ctx, req)
|
||||
if err != nil {
|
||||
return model.ChatCompletionResponse{}, fmt.Errorf("chat completion failed: %w", err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// CompleteText completes a text prompt using conversation messages
|
||||
// If systemPrompt is not empty, it will be added as the first message
|
||||
func CompleteText(ctx context.Context, profile *llm.Profile, systemPrompt string, messages []message.ChatMessage, maxTokens int) (string, error) {
|
||||
if profile == nil || profile.Client == nil {
|
||||
return "", fmt.Errorf("llm profile or client is nil")
|
||||
}
|
||||
|
||||
arkMessages := make([]*model.ChatCompletionMessage, 0, len(messages)+1)
|
||||
|
||||
// Add system prompt if provided
|
||||
if strings.TrimSpace(systemPrompt) != "" {
|
||||
content := &model.ChatCompletionMessageContent{
|
||||
StringValue: &systemPrompt,
|
||||
}
|
||||
arkMessages = append(arkMessages, &model.ChatCompletionMessage{
|
||||
Role: "system",
|
||||
Content: content,
|
||||
})
|
||||
}
|
||||
|
||||
for _, msg := range messages {
|
||||
content := &model.ChatCompletionMessageContent{
|
||||
StringValue: &msg.Content,
|
||||
}
|
||||
arkMessages = append(arkMessages, &model.ChatCompletionMessage{
|
||||
Role: msg.Role,
|
||||
Content: content,
|
||||
})
|
||||
}
|
||||
|
||||
req := model.CreateChatCompletionRequest{
|
||||
Model: profile.Config.Model,
|
||||
Messages: arkMessages,
|
||||
MaxTokens: &maxTokens,
|
||||
}
|
||||
|
||||
resp, err := profile.Client.CreateChatCompletion(ctx, req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("text completion failed: %w", err)
|
||||
}
|
||||
|
||||
if len(resp.Choices) == 0 {
|
||||
return "", fmt.Errorf("no completion choices returned")
|
||||
}
|
||||
|
||||
if resp.Choices[0].Message.Content != nil && resp.Choices[0].Message.Content.StringValue != nil {
|
||||
return *resp.Choices[0].Message.Content.StringValue, nil
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
|
||||
// CompleteTextWithArkMessages completes a text prompt using already converted Ark messages
|
||||
// This is used when messages have already been converted (e.g. after tool loop)
|
||||
func CompleteTextWithArkMessages(ctx context.Context, profile *llm.Profile, arkMessages []*model.ChatCompletionMessage, maxTokens int) (string, error) {
|
||||
if profile == nil || profile.Client == nil {
|
||||
return "", fmt.Errorf("llm profile or client is nil")
|
||||
}
|
||||
|
||||
req := model.CreateChatCompletionRequest{
|
||||
Model: profile.Config.Model,
|
||||
Messages: arkMessages,
|
||||
MaxTokens: &maxTokens,
|
||||
}
|
||||
|
||||
resp, err := profile.Client.CreateChatCompletion(ctx, req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("text completion failed: %w", err)
|
||||
}
|
||||
|
||||
if len(resp.Choices) == 0 {
|
||||
return "", fmt.Errorf("no completion choices returned")
|
||||
}
|
||||
|
||||
if resp.Choices[0].Message.Content != nil && resp.Choices[0].Message.Content.StringValue != nil {
|
||||
return *resp.Choices[0].Message.Content.StringValue, nil
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
@@ -0,0 +1,219 @@
|
||||
package conversation
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"meshtastic_mqtt_server/internal/message"
|
||||
)
|
||||
|
||||
// Store manages conversations stored as JSON files
|
||||
type Store struct {
|
||||
dir string
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// NewStore creates a new conversation store
|
||||
func NewStore(dir string) *Store {
|
||||
os.MkdirAll(dir, 0755)
|
||||
return &Store{dir: dir}
|
||||
}
|
||||
|
||||
// path returns the file path for a conversation ID
|
||||
func (s *Store) path(id string) string {
|
||||
return filepath.Join(s.dir, id+".json")
|
||||
}
|
||||
|
||||
// botPath returns the directory path for a specific bot
|
||||
func (s *Store) botPath(botID uint64) string {
|
||||
return filepath.Join(s.dir, fmt.Sprintf("bot_%d", botID))
|
||||
}
|
||||
|
||||
// Create creates a new conversation for a bot
|
||||
func (s *Store) Create(botID uint64, botNodeID string) (*message.Conversation, error) {
|
||||
return s.CreateWithPreset(botID, botNodeID, "")
|
||||
}
|
||||
|
||||
// CreateWithPreset creates a new conversation with a preset prompt
|
||||
func (s *Store) CreateWithPreset(botID uint64, botNodeID string, preset string) (*message.Conversation, error) {
|
||||
conv := &message.Conversation{
|
||||
ID: generateConversationID(botID),
|
||||
BotID: botID,
|
||||
BotNodeID: botNodeID,
|
||||
Title: "新对话",
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
PresetPrompt: strings.TrimSpace(preset),
|
||||
}
|
||||
if err := s.Save(conv); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return conv, nil
|
||||
}
|
||||
|
||||
// Save saves a conversation to disk
|
||||
func (s *Store) Save(conv *message.Conversation) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
conv.UpdatedAt = time.Now()
|
||||
return atomicWriteJSON(s.path(conv.ID), conv)
|
||||
}
|
||||
|
||||
// Get retrieves a conversation by ID
|
||||
func (s *Store) Get(id string) (*message.Conversation, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
data, err := os.ReadFile(s.path(id))
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, errors.New("对话不存在")
|
||||
}
|
||||
return nil, fmt.Errorf("读取对话失败: %w", err)
|
||||
}
|
||||
var conv message.Conversation
|
||||
if err := json.Unmarshal(data, &conv); err != nil {
|
||||
return nil, fmt.Errorf("解析对话失败: %w", err)
|
||||
}
|
||||
return &conv, nil
|
||||
}
|
||||
|
||||
// GetOrCreateForBot gets or creates a conversation for a bot
|
||||
func (s *Store) GetOrCreateForBot(botID uint64, botNodeID string, peerNodeID string) (*message.Conversation, error) {
|
||||
// Try to find an existing conversation with this peer
|
||||
convs, err := s.ListForBot(botID)
|
||||
if err == nil && len(convs) > 0 {
|
||||
// Use the most recent conversation (List already sorts by UpdatedAt desc)
|
||||
// Note: List returns convs with Messages = nil, so we need to reload
|
||||
return s.Get(convs[0].ID)
|
||||
}
|
||||
// Create a new conversation
|
||||
return s.Create(botID, botNodeID)
|
||||
}
|
||||
|
||||
// List returns all conversations
|
||||
func (s *Store) List() ([]message.Conversation, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
entries, err := os.ReadDir(s.dir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取对话目录失败: %w", err)
|
||||
}
|
||||
|
||||
var list []message.Conversation
|
||||
for _, e := range entries {
|
||||
if e.IsDir() || filepath.Ext(e.Name()) != ".json" {
|
||||
continue
|
||||
}
|
||||
data, err := os.ReadFile(filepath.Join(s.dir, e.Name()))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
var conv message.Conversation
|
||||
if err := json.Unmarshal(data, &conv); err != nil {
|
||||
continue
|
||||
}
|
||||
conv.Messages = nil // Don't return messages in list
|
||||
list = append(list, conv)
|
||||
}
|
||||
|
||||
sort.Slice(list, func(i, j int) bool {
|
||||
return list[i].UpdatedAt.After(list[j].UpdatedAt)
|
||||
})
|
||||
return list, nil
|
||||
}
|
||||
|
||||
// ListForBot returns all conversations for a specific bot
|
||||
func (s *Store) ListForBot(botID uint64) ([]message.Conversation, error) {
|
||||
all, err := s.List()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var filtered []message.Conversation
|
||||
for _, conv := range all {
|
||||
if conv.BotID == botID {
|
||||
filtered = append(filtered, conv)
|
||||
}
|
||||
}
|
||||
return filtered, nil
|
||||
}
|
||||
|
||||
// Delete deletes a conversation by ID
|
||||
func (s *Store) Delete(id string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if err := os.Remove(s.path(id)); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("删除对话失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteForBot deletes all conversations for a bot
|
||||
func (s *Store) DeleteForBot(botID uint64) error {
|
||||
convs, err := s.ListForBot(botID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, conv := range convs {
|
||||
_ = s.Delete(conv.ID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AddMessage adds a message to a conversation
|
||||
func (s *Store) AddMessage(convID string, msg message.ChatMessage) error {
|
||||
conv, err := s.Get(convID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
conv.Messages = append(conv.Messages, msg)
|
||||
if conv.Title == "" || conv.Title == "新对话" {
|
||||
conv.Title = GenerateTitle(conv.Messages)
|
||||
}
|
||||
return s.Save(conv)
|
||||
}
|
||||
|
||||
// atomicWriteJSON writes JSON to a file atomically
|
||||
func atomicWriteJSON(path string, v any) error {
|
||||
tmp := path + ".tmp"
|
||||
data, err := json.MarshalIndent(v, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.WriteFile(tmp, data, 0644); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, path)
|
||||
}
|
||||
|
||||
// GenerateTitle generates a title from conversation messages
|
||||
func GenerateTitle(messages []message.ChatMessage) string {
|
||||
for _, m := range messages {
|
||||
if m.Hidden {
|
||||
continue
|
||||
}
|
||||
if m.Role == "user" && strings.TrimSpace(m.Content) != "" {
|
||||
title := strings.TrimSpace(m.Content)
|
||||
title = strings.ReplaceAll(title, "\r\n", " ")
|
||||
title = strings.ReplaceAll(title, "\n", " ")
|
||||
runes := []rune(title)
|
||||
if len(runes) > 30 {
|
||||
return string(runes[:30]) + "..."
|
||||
}
|
||||
return title
|
||||
}
|
||||
}
|
||||
return "新对话"
|
||||
}
|
||||
|
||||
// generateConversationID generates a unique conversation ID
|
||||
func generateConversationID(botID uint64) string {
|
||||
return fmt.Sprintf("bot_%d_%d", botID, time.Now().UnixNano())
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
ark "github.com/volcengine/volcengine-go-sdk/service/arkruntime"
|
||||
)
|
||||
|
||||
// Profile represents an LLM provider configuration with client
|
||||
type Profile struct {
|
||||
Config ProviderConfig
|
||||
Client *ark.Client
|
||||
}
|
||||
|
||||
// ProviderConfig holds the configuration for an LLM provider
|
||||
type ProviderConfig struct {
|
||||
Name string
|
||||
Active bool
|
||||
APIKey string
|
||||
BaseURL string
|
||||
Model string
|
||||
Timeout int
|
||||
ContextWindowTokens int
|
||||
}
|
||||
|
||||
// State manages LLM profiles
|
||||
type State struct {
|
||||
mu sync.RWMutex
|
||||
profiles map[string]*Profile
|
||||
order []string
|
||||
activeName string
|
||||
}
|
||||
|
||||
// NewState creates a new LLM state from provider configurations
|
||||
func NewState(configs []ProviderConfig) (*State, error) {
|
||||
state := &State{
|
||||
profiles: make(map[string]*Profile, len(configs)),
|
||||
order: make([]string, 0, len(configs)),
|
||||
}
|
||||
for _, item := range configs {
|
||||
name := strings.TrimSpace(item.Name)
|
||||
if name == "" {
|
||||
return nil, errors.New("llm provider name cannot be empty")
|
||||
}
|
||||
if strings.TrimSpace(item.APIKey) == "" {
|
||||
return nil, fmt.Errorf("llm provider %s api_key is required", name)
|
||||
}
|
||||
if strings.TrimSpace(item.Model) == "" {
|
||||
return nil, fmt.Errorf("llm provider %s model is required", name)
|
||||
}
|
||||
if strings.TrimSpace(item.BaseURL) == "" {
|
||||
return nil, fmt.Errorf("llm provider %s base_url is required", name)
|
||||
}
|
||||
if item.Timeout <= 0 {
|
||||
item.Timeout = 120
|
||||
}
|
||||
if _, ok := state.profiles[name]; ok {
|
||||
return nil, fmt.Errorf("duplicate llm provider name: %s", name)
|
||||
}
|
||||
state.profiles[name] = &Profile{
|
||||
Config: item,
|
||||
Client: ark.NewClientWithApiKey(
|
||||
item.APIKey,
|
||||
ark.WithBaseUrl(item.BaseURL),
|
||||
ark.WithTimeout(time.Duration(item.Timeout)*time.Second),
|
||||
),
|
||||
}
|
||||
state.order = append(state.order, name)
|
||||
if item.Active && state.activeName == "" {
|
||||
state.activeName = name
|
||||
}
|
||||
}
|
||||
if len(state.order) == 0 {
|
||||
return nil, errors.New("at least one llm provider is required")
|
||||
}
|
||||
if state.activeName == "" {
|
||||
state.activeName = state.order[0]
|
||||
}
|
||||
return state, nil
|
||||
}
|
||||
|
||||
// ActiveProfile returns the currently active LLM profile
|
||||
func (s *State) ActiveProfile() *Profile {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.profiles[s.activeName]
|
||||
}
|
||||
|
||||
// GetProfile returns a profile by name, or the active profile if name is empty
|
||||
func (s *State) GetProfile(name string) (*Profile, error) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
return s.profiles[s.activeName], nil
|
||||
}
|
||||
profile, ok := s.profiles[name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("llm provider not found: %s", name)
|
||||
}
|
||||
return profile, nil
|
||||
}
|
||||
|
||||
// SwitchActive changes the active LLM profile
|
||||
func (s *State) SwitchActive(name string) (*Profile, error) {
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
return nil, errors.New("llm provider name cannot be empty")
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
profile, ok := s.profiles[name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("llm provider not found: %s", name)
|
||||
}
|
||||
s.activeName = name
|
||||
return profile, nil
|
||||
}
|
||||
|
||||
// ListProfiles returns all LLM profiles
|
||||
func (s *State) ListProfiles() []ProviderConfig {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
profiles := make([]ProviderConfig, 0, len(s.order))
|
||||
for _, name := range s.order {
|
||||
profile := s.profiles[name]
|
||||
cfg := profile.Config
|
||||
cfg.APIKey = "" // Hide API key in listings
|
||||
cfg.Active = name == s.activeName
|
||||
profiles = append(profiles, cfg)
|
||||
}
|
||||
return profiles
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package message
|
||||
|
||||
import "time"
|
||||
|
||||
// ChatMessage represents a single message in a conversation
|
||||
type ChatMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
ImageURL string `json:"image_url,omitempty"`
|
||||
ImageURLAlias string `json:"image_url,omitempty"`
|
||||
Hidden bool `json:"hidden,omitempty"`
|
||||
}
|
||||
|
||||
// Conversation represents a chat conversation with a bot
|
||||
type Conversation struct {
|
||||
ID string `json:"id"`
|
||||
BotID uint64 `json:"bot_id"`
|
||||
BotNodeID string `json:"bot_node_id"`
|
||||
Title string `json:"title"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
PresetPrompt string `json:"preset_prompt,omitempty"`
|
||||
Messages []ChatMessage `json:"messages,omitempty"`
|
||||
}
|
||||
|
||||
// ConversationPreset represents a preset configuration for a bot's conversation
|
||||
type ConversationPreset struct {
|
||||
ID uint64 `json:"id"`
|
||||
BotID uint64 `json:"bot_id"`
|
||||
Name string `json:"name"`
|
||||
SystemPrompt string `json:"system_prompt"`
|
||||
Enabled bool `json:"enabled"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
package mqtpp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
)
|
||||
|
||||
const NodeNumBroadcast uint32 = 0xffffffff
|
||||
|
||||
type PacketBuildOptions struct {
|
||||
FromNodeNum uint32
|
||||
ToNodeNum uint32
|
||||
PacketID uint32
|
||||
ChannelID string
|
||||
GatewayID string
|
||||
PSK []byte
|
||||
Encrypt bool
|
||||
ViaMQTT bool
|
||||
}
|
||||
|
||||
type TextMessageBuildOptions struct {
|
||||
PacketBuildOptions
|
||||
Text string
|
||||
}
|
||||
|
||||
type NodeInfoBuildOptions struct {
|
||||
PacketBuildOptions
|
||||
NodeID string
|
||||
LongName string
|
||||
ShortName string
|
||||
HWModel uint32
|
||||
Role uint32
|
||||
IsLicensed bool
|
||||
PublicKey []byte
|
||||
}
|
||||
|
||||
// AckBuildOptions 描述构造一个 Routing-NONE ACK（PSK 频道路径）所需字段。
|
||||
// RequestID 是被 ACK 原始包的 packet_id（写入 Data.request_id, tag 6）。
|
||||
type AckBuildOptions struct {
|
||||
PacketBuildOptions
|
||||
RequestID uint32
|
||||
}
|
||||
|
||||
func BuildTextMessageServiceEnvelope(opts TextMessageBuildOptions) ([]byte, error) {
|
||||
if opts.FromNodeNum == 0 {
|
||||
return nil, fmt.Errorf("from node number is required")
|
||||
}
|
||||
if opts.PacketID == 0 {
|
||||
return nil, fmt.Errorf("packet id is required")
|
||||
}
|
||||
if opts.ChannelID == "" {
|
||||
return nil, fmt.Errorf("channel id is required")
|
||||
}
|
||||
if strings.TrimSpace(opts.GatewayID) == "" {
|
||||
opts.GatewayID = NodeNumToID(opts.FromNodeNum)
|
||||
}
|
||||
if opts.Text == "" {
|
||||
return nil, fmt.Errorf("text is required")
|
||||
}
|
||||
if !utf8.ValidString(opts.Text) {
|
||||
return nil, fmt.Errorf("text must be valid utf-8")
|
||||
}
|
||||
|
||||
data := buildDataPacket(textMessageApp, []byte(opts.Text))
|
||||
packet, err := buildMeshPacket(opts.PacketBuildOptions, data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return buildServiceEnvelope(packet, opts.ChannelID, opts.GatewayID), nil
|
||||
}
|
||||
|
||||
func BuildNodeInfoServiceEnvelope(opts NodeInfoBuildOptions) ([]byte, error) {
|
||||
if opts.NodeID == "" {
|
||||
opts.NodeID = NodeNumToID(opts.FromNodeNum)
|
||||
}
|
||||
opts.NodeID = truncateUTF8Bytes(opts.NodeID, 16)
|
||||
opts.LongName = truncateUTF8Bytes(strings.TrimSpace(opts.LongName), 40)
|
||||
opts.ShortName = truncateUTF8Bytes(strings.TrimSpace(opts.ShortName), 5)
|
||||
if opts.LongName == "" {
|
||||
return nil, fmt.Errorf("long name is required")
|
||||
}
|
||||
if opts.ShortName == "" {
|
||||
return nil, fmt.Errorf("short name is required")
|
||||
}
|
||||
if len(opts.PublicKey) > 32 {
|
||||
opts.PublicKey = opts.PublicKey[:32]
|
||||
}
|
||||
user := buildUserPacket(opts)
|
||||
data := buildDataPacket(nodeInfoApp, user)
|
||||
packet, err := buildMeshPacket(opts.PacketBuildOptions, data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return buildServiceEnvelope(packet, opts.ChannelID, opts.GatewayID), nil
|
||||
}
|
||||
|
||||
// BuildAckServiceEnvelope 构造一个 PSK 频道上的 Routing ACK（error_reason=NONE）。
|
||||
// 与固件 MeshModule::allocAckNak/Router::sendAckNak 行为对齐：
|
||||
// - portnum = ROUTING_APP(5)
|
||||
// - Data.request_id = 原始包 ID
|
||||
// - Routing.which_variant = error_reason，error_reason = NONE(0)
|
||||
//
|
||||
// 用 PacketBuildOptions 中 ChannelID + PSK 加密；调用方负责把 ToNodeNum 设为原 from。
|
||||
func BuildAckServiceEnvelope(opts AckBuildOptions) ([]byte, error) {
|
||||
if opts.RequestID == 0 {
|
||||
return nil, fmt.Errorf("ack request_id is required")
|
||||
}
|
||||
if opts.ChannelID == "" {
|
||||
return nil, fmt.Errorf("channel id is required")
|
||||
}
|
||||
data := buildAckDataPacket(opts.RequestID)
|
||||
packet, err := buildMeshPacket(opts.PacketBuildOptions, data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if strings.TrimSpace(opts.GatewayID) == "" {
|
||||
opts.GatewayID = NodeNumToID(opts.FromNodeNum)
|
||||
}
|
||||
return buildServiceEnvelope(packet, opts.ChannelID, opts.GatewayID), nil
|
||||
}
|
||||
|
||||
// buildAckDataPacket 构造 Data { portnum=ROUTING_APP, payload=Routing{error_reason=NONE}, request_id=req }。
|
||||
func buildAckDataPacket(requestID uint32) []byte {
|
||||
// Routing payload: oneof variant=error_reason(tag 3), value=NONE(0) → 0x18 0x00
|
||||
routing := protowire.AppendTag(nil, 3, protowire.VarintType)
|
||||
routing = protowire.AppendVarint(routing, 0)
|
||||
|
||||
var out []byte
|
||||
out = protowire.AppendTag(out, 1, protowire.VarintType)
|
||||
out = protowire.AppendVarint(out, uint64(routingApp))
|
||||
out = protowire.AppendTag(out, 2, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, routing)
|
||||
out = protowire.AppendTag(out, 6, protowire.Fixed32Type)
|
||||
out = protowire.AppendFixed32(out, requestID)
|
||||
return out
|
||||
}
|
||||
|
||||
func NodeNumToID(nodeNum uint32) string {
|
||||
return nodeNumToID(nodeNum)
|
||||
}
|
||||
|
||||
func truncateUTF8Bytes(value string, maxBytes int) string {
|
||||
if maxBytes <= 0 || len(value) <= maxBytes {
|
||||
return value
|
||||
}
|
||||
out := make([]byte, 0, maxBytes)
|
||||
for _, r := range value {
|
||||
part := string(r)
|
||||
if len(out)+len(part) > maxBytes {
|
||||
break
|
||||
}
|
||||
out = append(out, part...)
|
||||
}
|
||||
return string(out)
|
||||
}
|
||||
|
||||
func ParseNodeID(nodeID string) (uint32, error) {
|
||||
value := strings.TrimSpace(nodeID)
|
||||
if value == "" {
|
||||
return 0, fmt.Errorf("node id is required")
|
||||
}
|
||||
value = strings.TrimPrefix(value, "!")
|
||||
if len(value) != 8 {
|
||||
return 0, fmt.Errorf("node id must be !xxxxxxxx")
|
||||
}
|
||||
num, err := strconv.ParseUint(value, 16, 32)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("invalid node id: %w", err)
|
||||
}
|
||||
return uint32(num), nil
|
||||
}
|
||||
|
||||
func buildDataPacket(portnum uint32, payload []byte) []byte {
|
||||
var out []byte
|
||||
out = protowire.AppendTag(out, 1, protowire.VarintType)
|
||||
out = protowire.AppendVarint(out, uint64(portnum))
|
||||
out = protowire.AppendTag(out, 2, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, payload)
|
||||
return out
|
||||
}
|
||||
|
||||
func buildUserPacket(opts NodeInfoBuildOptions) []byte {
|
||||
var out []byte
|
||||
out = protowire.AppendTag(out, 1, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, []byte(opts.NodeID))
|
||||
out = protowire.AppendTag(out, 2, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, []byte(opts.LongName))
|
||||
out = protowire.AppendTag(out, 3, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, []byte(opts.ShortName))
|
||||
if opts.HWModel != 0 {
|
||||
out = protowire.AppendTag(out, 5, protowire.VarintType)
|
||||
out = protowire.AppendVarint(out, uint64(opts.HWModel))
|
||||
}
|
||||
out = protowire.AppendTag(out, 6, protowire.VarintType)
|
||||
if opts.IsLicensed {
|
||||
out = protowire.AppendVarint(out, 1)
|
||||
} else {
|
||||
out = protowire.AppendVarint(out, 0)
|
||||
}
|
||||
out = protowire.AppendTag(out, 7, protowire.VarintType)
|
||||
out = protowire.AppendVarint(out, uint64(opts.Role))
|
||||
if len(opts.PublicKey) > 0 {
|
||||
out = protowire.AppendTag(out, 8, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, opts.PublicKey)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func buildMeshPacket(opts PacketBuildOptions, data []byte) ([]byte, error) {
|
||||
if opts.FromNodeNum == 0 {
|
||||
return nil, fmt.Errorf("from node number is required")
|
||||
}
|
||||
if opts.PacketID == 0 {
|
||||
return nil, fmt.Errorf("packet id is required")
|
||||
}
|
||||
if opts.ChannelID == "" {
|
||||
return nil, fmt.Errorf("channel id is required")
|
||||
}
|
||||
if strings.TrimSpace(opts.GatewayID) == "" {
|
||||
opts.GatewayID = NodeNumToID(opts.FromNodeNum)
|
||||
}
|
||||
var out []byte
|
||||
out = protowire.AppendTag(out, 1, protowire.Fixed32Type)
|
||||
out = protowire.AppendFixed32(out, opts.FromNodeNum)
|
||||
out = protowire.AppendTag(out, 2, protowire.Fixed32Type)
|
||||
out = protowire.AppendFixed32(out, opts.ToNodeNum)
|
||||
|
||||
if opts.Encrypt {
|
||||
if len(opts.PSK) == 0 {
|
||||
return nil, fmt.Errorf("psk is required for encrypted packet")
|
||||
}
|
||||
ciphertext, err := cryptAESCTR(opts.PSK, opts.FromNodeNum, opts.PacketID, data)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = protowire.AppendTag(out, 3, protowire.VarintType)
|
||||
out = protowire.AppendVarint(out, uint64(channelHash(opts.ChannelID, opts.PSK)))
|
||||
out = protowire.AppendTag(out, 5, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, ciphertext)
|
||||
} else {
|
||||
out = protowire.AppendTag(out, 4, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, data)
|
||||
}
|
||||
|
||||
out = protowire.AppendTag(out, 6, protowire.Fixed32Type)
|
||||
out = protowire.AppendFixed32(out, opts.PacketID)
|
||||
if opts.ViaMQTT {
|
||||
out = protowire.AppendTag(out, 14, protowire.VarintType)
|
||||
out = protowire.AppendVarint(out, 1)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func buildServiceEnvelope(packet []byte, channelID string, gatewayID string) []byte {
|
||||
var out []byte
|
||||
out = protowire.AppendTag(out, 1, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, packet)
|
||||
out = protowire.AppendTag(out, 2, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, []byte(channelID))
|
||||
out = protowire.AppendTag(out, 3, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, []byte(gatewayID))
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,215 @@
|
||||
package mqtpp
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestBuildTextMessageServiceEnvelopeRoundTrip(t *testing.T) {
|
||||
key, err := ExpandPSK("AQ==")
|
||||
if err != nil {
|
||||
t.Fatalf("ExpandPSK() error = %v", err)
|
||||
}
|
||||
|
||||
raw, err := BuildTextMessageServiceEnvelope(TextMessageBuildOptions{
|
||||
PacketBuildOptions: PacketBuildOptions{
|
||||
FromNodeNum: 0x12345678,
|
||||
ToNodeNum: NodeNumBroadcast,
|
||||
PacketID: 0x87654321,
|
||||
ChannelID: "LongFast",
|
||||
GatewayID: "!12345678",
|
||||
PSK: key,
|
||||
Encrypt: true,
|
||||
ViaMQTT: true,
|
||||
},
|
||||
Text: "hello from bot",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildTextMessageServiceEnvelope() error = %v", err)
|
||||
}
|
||||
|
||||
valid, _, record := MQTTPP("msh/2/e/LongFast/!12345678", raw, key, Options{})
|
||||
if !valid {
|
||||
t.Fatalf("MQTTPP() valid = false, record = %#v", record)
|
||||
}
|
||||
if record["type"] != "text_message" {
|
||||
t.Fatalf("record type = %v", record["type"])
|
||||
}
|
||||
if record["text"] != "hello from bot" {
|
||||
t.Fatalf("text = %v", record["text"])
|
||||
}
|
||||
if record["from_num"] != uint32(0x12345678) {
|
||||
t.Fatalf("from_num = %v", record["from_num"])
|
||||
}
|
||||
if record["packet_to_num"] != uint32(NodeNumBroadcast) {
|
||||
t.Fatalf("packet_to_num = %v", record["packet_to_num"])
|
||||
}
|
||||
if record["decrypt_success"] != true {
|
||||
t.Fatalf("decrypt_success = %v", record["decrypt_success"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildTextMessageServiceEnvelopeDirectRoundTrip(t *testing.T) {
|
||||
key, err := ExpandPSK("AQ==")
|
||||
if err != nil {
|
||||
t.Fatalf("ExpandPSK() error = %v", err)
|
||||
}
|
||||
|
||||
raw, err := BuildTextMessageServiceEnvelope(TextMessageBuildOptions{
|
||||
PacketBuildOptions: PacketBuildOptions{
|
||||
FromNodeNum: 0x12345678,
|
||||
ToNodeNum: 0x10203040,
|
||||
PacketID: 0x11111111,
|
||||
ChannelID: "LongFast",
|
||||
GatewayID: "!12345678",
|
||||
PSK: key,
|
||||
Encrypt: true,
|
||||
ViaMQTT: true,
|
||||
},
|
||||
Text: "direct hello",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildTextMessageServiceEnvelope() error = %v", err)
|
||||
}
|
||||
|
||||
valid, _, record := MQTTPP("msh/2/e/LongFast/!12345678", raw, key, Options{})
|
||||
if !valid {
|
||||
t.Fatalf("MQTTPP() valid = false, record = %#v", record)
|
||||
}
|
||||
if record["text"] != "direct hello" {
|
||||
t.Fatalf("text = %v", record["text"])
|
||||
}
|
||||
if record["packet_to"] != "!10203040" {
|
||||
t.Fatalf("packet_to = %v", record["packet_to"])
|
||||
}
|
||||
if record["packet_to_num"] != uint32(0x10203040) {
|
||||
t.Fatalf("packet_to_num = %v", record["packet_to_num"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildNodeInfoServiceEnvelopeRoundTrip(t *testing.T) {
|
||||
key, err := ExpandPSK("AQ==")
|
||||
if err != nil {
|
||||
t.Fatalf("ExpandPSK() error = %v", err)
|
||||
}
|
||||
|
||||
raw, err := BuildNodeInfoServiceEnvelope(NodeInfoBuildOptions{
|
||||
PacketBuildOptions: PacketBuildOptions{
|
||||
FromNodeNum: 0x12345678,
|
||||
ToNodeNum: NodeNumBroadcast,
|
||||
PacketID: 0x22222222,
|
||||
ChannelID: "LongFast",
|
||||
GatewayID: "!12345678",
|
||||
PSK: key,
|
||||
Encrypt: true,
|
||||
ViaMQTT: true,
|
||||
},
|
||||
NodeID: "!12345678",
|
||||
LongName: "MQTT Bot",
|
||||
ShortName: "BT",
|
||||
HWModel: 255,
|
||||
Role: 0,
|
||||
IsLicensed: false,
|
||||
PublicKey: []byte{1, 2, 3},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildNodeInfoServiceEnvelope() error = %v", err)
|
||||
}
|
||||
|
||||
valid, _, record := MQTTPP("msh/2/e/LongFast/!12345678", raw, key, Options{})
|
||||
if !valid {
|
||||
t.Fatalf("MQTTPP() valid = false, record = %#v", record)
|
||||
}
|
||||
if record["type"] != "nodeinfo" {
|
||||
t.Fatalf("record type = %v", record["type"])
|
||||
}
|
||||
if record["long_name"] != "MQTT Bot" {
|
||||
t.Fatalf("long_name = %v", record["long_name"])
|
||||
}
|
||||
if record["short_name"] != "BT" {
|
||||
t.Fatalf("short_name = %v", record["short_name"])
|
||||
}
|
||||
if record["hw_model"] != "PRIVATE_HW" {
|
||||
t.Fatalf("hw_model = %v", record["hw_model"])
|
||||
}
|
||||
if record["role"] != "CLIENT" {
|
||||
t.Fatalf("role = %v", record["role"])
|
||||
}
|
||||
if record["is_licensed"] != false {
|
||||
t.Fatalf("is_licensed = %v", record["is_licensed"])
|
||||
}
|
||||
if record["public_key"] != "010203" {
|
||||
t.Fatalf("public_key = %v", record["public_key"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildNodeInfoTruncatesNanopbStrings(t *testing.T) {
|
||||
key, err := ExpandPSK("AQ==")
|
||||
if err != nil {
|
||||
t.Fatalf("ExpandPSK() error = %v", err)
|
||||
}
|
||||
|
||||
raw, err := BuildNodeInfoServiceEnvelope(NodeInfoBuildOptions{
|
||||
PacketBuildOptions: PacketBuildOptions{FromNodeNum: 0x12345678, ToNodeNum: NodeNumBroadcast, PacketID: 0x33333333, ChannelID: "LongFast", GatewayID: "!12345678", PSK: key, Encrypt: true, ViaMQTT: true},
|
||||
NodeID: "!12345678",
|
||||
LongName: "这是一个非常非常非常非常长的机器人节点名称",
|
||||
ShortName: "机器人",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildNodeInfoServiceEnvelope() error = %v", err)
|
||||
}
|
||||
valid, _, record := MQTTPP("msh/2/e/LongFast/!12345678", raw, key, Options{})
|
||||
if !valid {
|
||||
t.Fatalf("MQTTPP() valid = false, record = %#v", record)
|
||||
}
|
||||
if len([]byte(record["long_name"].(string))) > 40 {
|
||||
t.Fatalf("long_name byte length = %d", len([]byte(record["long_name"].(string))))
|
||||
}
|
||||
if len([]byte(record["short_name"].(string))) > 5 {
|
||||
t.Fatalf("short_name byte length = %d", len([]byte(record["short_name"].(string))))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAckServiceEnvelopeRoundTrip(t *testing.T) {
|
||||
key, err := ExpandPSK("AQ==")
|
||||
if err != nil {
|
||||
t.Fatalf("ExpandPSK: %v", err)
|
||||
}
|
||||
const requestID uint32 = 0xabcd1234
|
||||
raw, err := BuildAckServiceEnvelope(AckBuildOptions{
|
||||
PacketBuildOptions: PacketBuildOptions{
|
||||
FromNodeNum: 0x10101010,
|
||||
ToNodeNum: 0x20202020,
|
||||
PacketID: 0x30303030,
|
||||
ChannelID: "LongFast",
|
||||
GatewayID: "!10101010",
|
||||
PSK: key,
|
||||
Encrypt: true,
|
||||
ViaMQTT: true,
|
||||
},
|
||||
RequestID: requestID,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildAckServiceEnvelope: %v", err)
|
||||
}
|
||||
valid, _, record := MQTTPP("msh/2/e/LongFast/!10101010", raw, key, Options{})
|
||||
if !valid {
|
||||
t.Fatalf("MQTTPP not valid: %#v", record)
|
||||
}
|
||||
if record["portnum"] != "ROUTING_APP" {
|
||||
t.Fatalf("portnum = %v", record["portnum"])
|
||||
}
|
||||
if record["type"] != "routing" {
|
||||
t.Fatalf("type = %v", record["type"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseNodeID(t *testing.T) {
|
||||
num, err := ParseNodeID("!1234abcd")
|
||||
if err != nil {
|
||||
t.Fatalf("ParseNodeID() error = %v", err)
|
||||
}
|
||||
if num != 0x1234abcd {
|
||||
t.Fatalf("num = %#x", num)
|
||||
}
|
||||
if NodeNumToID(num) != "!1234abcd" {
|
||||
t.Fatalf("NodeNumToID() = %s", NodeNumToID(num))
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,48 @@
|
||||
package mqtpp
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
)
|
||||
|
||||
func TestMQTTPPEncryptedPacketDefaultRejected(t *testing.T) {
|
||||
raw := encryptedServiceEnvelopeTestPayload()
|
||||
valid, payload, record := MQTTPP("msh/test", raw, nil, Options{})
|
||||
if valid {
|
||||
t.Fatalf("valid = true, want false")
|
||||
}
|
||||
if payload != nil {
|
||||
t.Fatalf("payload = %v, want nil", payload)
|
||||
}
|
||||
if record["type"] != "encrypted_packet" {
|
||||
t.Fatalf("type = %v, want encrypted_packet", record["type"])
|
||||
}
|
||||
if record["error"] != "cannot be decrypted" {
|
||||
t.Fatalf("error = %v, want cannot be decrypted", record["error"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestMQTTPPEncryptedPacketAllowed(t *testing.T) {
|
||||
raw := encryptedServiceEnvelopeTestPayload()
|
||||
valid, payload, record := MQTTPP("msh/test", raw, nil, Options{AllowEncryptedForwarding: true})
|
||||
if !valid {
|
||||
t.Fatalf("valid = false, want true: %+v", record)
|
||||
}
|
||||
if string(payload) != string(raw) {
|
||||
t.Fatalf("payload = %v, want raw payload", payload)
|
||||
}
|
||||
if record["type"] != "encrypted_packet" {
|
||||
t.Fatalf("type = %v, want encrypted_packet", record["type"])
|
||||
}
|
||||
if record["error"] != nil {
|
||||
t.Fatalf("error = %v, want nil", record["error"])
|
||||
}
|
||||
}
|
||||
|
||||
func encryptedServiceEnvelopeTestPayload() []byte {
|
||||
packet := protowire.AppendTag(nil, 5, protowire.BytesType)
|
||||
packet = protowire.AppendBytes(packet, []byte{1, 2, 3, 4})
|
||||
envelope := protowire.AppendTag(nil, 1, protowire.BytesType)
|
||||
return protowire.AppendBytes(envelope, packet)
|
||||
}
|
||||
@@ -0,0 +1,363 @@
|
||||
package mqtpp
|
||||
|
||||
import (
|
||||
"crypto/aes"
|
||||
"crypto/ecdh"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"crypto/subtle"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
)
|
||||
|
||||
// PKIChannelID 是固件在 ServiceEnvelope/MQTT topic 中标识 PKI 加密包时使用的字面量
|
||||
// 见 firmware/src/mqtt/MQTT.cpp 中 `channelId = isPKIEncrypted ? "PKI" : channels.getGlobalId(chIndex);`
|
||||
const PKIChannelID = "PKI"
|
||||
|
||||
// pkcOverhead 与固件 MESHTASTIC_PKC_OVERHEAD 一致：8 字节 AES-CCM 认证标签 + 4 字节 extraNonce
|
||||
const pkcOverhead = 12
|
||||
|
||||
// PKITextMessageBuildOptions 描述构造一条 PKI 加密 DM 所需的全部上下文
|
||||
type PKITextMessageBuildOptions struct {
|
||||
FromNodeNum uint32
|
||||
ToNodeNum uint32
|
||||
PacketID uint32
|
||||
GatewayID string
|
||||
ViaMQTT bool
|
||||
SenderPrivate []byte // X25519 32 字节私钥
|
||||
RecipientPub []byte // X25519 32 字节公钥
|
||||
SenderPublic []byte // 可选；附在 MeshPacket.public_key (tag 16)
|
||||
Text string
|
||||
}
|
||||
|
||||
// BuildPKITextMessageServiceEnvelope 构造一条遵循固件实现的 PKI 私聊文本消息：
|
||||
// - data 包: portnum=TEXT_MESSAGE_APP, payload=text
|
||||
// - 共享密钥: SHA256(X25519(senderPriv, recipientPub))
|
||||
// - AES-CCM(M=8,L=2,AAD=0); nonce = packetId(8B LE) | fromNode(4B LE) | extraNonce(4B LE，覆盖 fromNode 后续 4 字节)
|
||||
// - encrypted bytes 末尾追加 8 字节 auth + 4 字节 extraNonce(LE)
|
||||
// - MeshPacket.channel = 0, pki_encrypted(tag17)=1
|
||||
// - ServiceEnvelope.channel_id 固定 "PKI"
|
||||
func BuildPKITextMessageServiceEnvelope(opts PKITextMessageBuildOptions) ([]byte, error) {
|
||||
if opts.FromNodeNum == 0 {
|
||||
return nil, fmt.Errorf("from node number is required")
|
||||
}
|
||||
if opts.ToNodeNum == 0 || opts.ToNodeNum == NodeNumBroadcast {
|
||||
return nil, fmt.Errorf("pki direct message requires a non-broadcast destination")
|
||||
}
|
||||
if opts.PacketID == 0 {
|
||||
return nil, fmt.Errorf("packet id is required")
|
||||
}
|
||||
if opts.Text == "" {
|
||||
return nil, fmt.Errorf("text is required")
|
||||
}
|
||||
if !utf8.ValidString(opts.Text) {
|
||||
return nil, fmt.Errorf("text must be valid utf-8")
|
||||
}
|
||||
if len(opts.SenderPrivate) != 32 {
|
||||
return nil, fmt.Errorf("sender private key must be 32 bytes")
|
||||
}
|
||||
if len(opts.RecipientPub) != 32 {
|
||||
return nil, fmt.Errorf("recipient public key must be 32 bytes")
|
||||
}
|
||||
if strings.TrimSpace(opts.GatewayID) == "" {
|
||||
opts.GatewayID = NodeNumToID(opts.FromNodeNum)
|
||||
}
|
||||
|
||||
plaintext := buildDataPacket(textMessageApp, []byte(opts.Text))
|
||||
|
||||
sharedKey, err := pkiSharedKey(opts.SenderPrivate, opts.RecipientPub)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var extraNonceBuf [4]byte
|
||||
if _, err := rand.Read(extraNonceBuf[:]); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
extraNonce := binary.LittleEndian.Uint32(extraNonceBuf[:])
|
||||
|
||||
ciphertext, auth, err := aesCCMEncrypt(sharedKey, pkiNonce(opts.PacketID, opts.FromNodeNum, extraNonce), plaintext)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
encrypted := make([]byte, 0, len(ciphertext)+pkcOverhead)
|
||||
encrypted = append(encrypted, ciphertext...)
|
||||
encrypted = append(encrypted, auth...)
|
||||
encrypted = append(encrypted, extraNonceBuf[:]...)
|
||||
|
||||
packet := buildPKIMeshPacket(opts.FromNodeNum, opts.ToNodeNum, opts.PacketID, opts.ViaMQTT, encrypted, opts.SenderPublic)
|
||||
return buildServiceEnvelope(packet, PKIChannelID, opts.GatewayID), nil
|
||||
}
|
||||
|
||||
// PKIAckBuildOptions 描述构造一个 PKI 加密的 Routing-NONE ACK 所需字段。
|
||||
type PKIAckBuildOptions struct {
|
||||
FromNodeNum uint32 // 我们（机器人）的节点号
|
||||
ToNodeNum uint32 // 原发送者
|
||||
PacketID uint32 // 新生成的 ACK 自身的 packet id
|
||||
RequestID uint32 // 被 ACK 的原始包 packet id
|
||||
GatewayID string
|
||||
ViaMQTT bool
|
||||
SenderPrivate []byte
|
||||
RecipientPub []byte
|
||||
SenderPublic []byte
|
||||
}
|
||||
|
||||
// BuildPKIAckServiceEnvelope 构造一条 PKI 加密的 Routing-NONE ACK，与固件
|
||||
// MeshModule::allocAckNak + Router::perhapsEncode (PKI 分支) 行为对齐。
|
||||
func BuildPKIAckServiceEnvelope(opts PKIAckBuildOptions) ([]byte, error) {
|
||||
if opts.FromNodeNum == 0 {
|
||||
return nil, fmt.Errorf("from node number is required")
|
||||
}
|
||||
if opts.ToNodeNum == 0 || opts.ToNodeNum == NodeNumBroadcast {
|
||||
return nil, fmt.Errorf("pki ack requires a non-broadcast destination")
|
||||
}
|
||||
if opts.PacketID == 0 {
|
||||
return nil, fmt.Errorf("packet id is required")
|
||||
}
|
||||
if opts.RequestID == 0 {
|
||||
return nil, fmt.Errorf("request id is required")
|
||||
}
|
||||
if len(opts.SenderPrivate) != 32 || len(opts.RecipientPub) != 32 {
|
||||
return nil, fmt.Errorf("pki keys must be 32 bytes each")
|
||||
}
|
||||
if strings.TrimSpace(opts.GatewayID) == "" {
|
||||
opts.GatewayID = NodeNumToID(opts.FromNodeNum)
|
||||
}
|
||||
|
||||
plaintext := buildAckDataPacket(opts.RequestID)
|
||||
|
||||
sharedKey, err := pkiSharedKey(opts.SenderPrivate, opts.RecipientPub)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var extraNonceBuf [4]byte
|
||||
if _, err := rand.Read(extraNonceBuf[:]); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
extraNonce := binary.LittleEndian.Uint32(extraNonceBuf[:])
|
||||
ciphertext, auth, err := aesCCMEncrypt(sharedKey, pkiNonce(opts.PacketID, opts.FromNodeNum, extraNonce), plaintext)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
encrypted := make([]byte, 0, len(ciphertext)+pkcOverhead)
|
||||
encrypted = append(encrypted, ciphertext...)
|
||||
encrypted = append(encrypted, auth...)
|
||||
encrypted = append(encrypted, extraNonceBuf[:]...)
|
||||
|
||||
packet := buildPKIMeshPacket(opts.FromNodeNum, opts.ToNodeNum, opts.PacketID, opts.ViaMQTT, encrypted, opts.SenderPublic)
|
||||
return buildServiceEnvelope(packet, PKIChannelID, opts.GatewayID), nil
|
||||
}
|
||||
|
||||
|
||||
func pkiSharedKey(privateKey, publicKey []byte) ([]byte, error) {
|
||||
curve := ecdh.X25519()
|
||||
priv, err := curve.NewPrivateKey(privateKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid sender private key: %w", err)
|
||||
}
|
||||
pub, err := curve.NewPublicKey(publicKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid recipient public key: %w", err)
|
||||
}
|
||||
shared, err := priv.ECDH(pub)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("x25519 ecdh failed: %w", err)
|
||||
}
|
||||
digest := sha256.Sum256(shared)
|
||||
return digest[:], nil
|
||||
}
|
||||
|
||||
// pkiNonce 完整复刻固件 CryptoEngine::initNonce(fromNode, packetId, extraNonce) 的字节布局。
|
||||
// 固件实现（mesh/CryptoEngine.cpp）：
|
||||
//
|
||||
// memcpy(nonce + 0, &packetId, 8); // packetId 是 uint64，写入 nonce[0..8)
|
||||
// memcpy(nonce + 8, &fromNode, 4); // fromNode 写入 nonce[8..12)
|
||||
// if (extraNonce)
|
||||
// memcpy(nonce + 4, &extraNonce, 4); // extraNonce 覆盖 nonce[4..8)
|
||||
//
|
||||
// 因此 13 字节 nonce 布局为：packetId_lo(4B LE) | extraNonce_or_packetId_hi(4B LE) | fromNode(4B LE) | 0x00
|
||||
func pkiNonce(packetID, fromNode, extraNonce uint32) []byte {
|
||||
nonce := make([]byte, 16)
|
||||
binary.LittleEndian.PutUint64(nonce[0:8], uint64(packetID)) // packetId 是 uint64，高 32 位为 0
|
||||
binary.LittleEndian.PutUint32(nonce[8:12], fromNode)
|
||||
if extraNonce != 0 {
|
||||
binary.LittleEndian.PutUint32(nonce[4:8], extraNonce)
|
||||
}
|
||||
// CCM L=2 → nonce 占 15-L=13 字节
|
||||
return nonce[:13]
|
||||
}
|
||||
|
||||
// aesCCMEncrypt 使用与固件相同的参数（AES-CCM, M=8 即 8 字节 tag, L=2, 无 AAD）。
|
||||
func aesCCMEncrypt(key, nonce, plaintext []byte) (ciphertext []byte, auth []byte, err error) {
|
||||
if len(nonce) != 13 {
|
||||
return nil, nil, fmt.Errorf("ccm nonce must be 13 bytes")
|
||||
}
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
const tagLen = 8
|
||||
if len(plaintext) > 0xffff {
|
||||
return nil, nil, fmt.Errorf("plaintext too large for L=2 ccm")
|
||||
}
|
||||
|
||||
// CBC-MAC 鉴权
|
||||
var x [aes.BlockSize]byte
|
||||
var b [aes.BlockSize]byte
|
||||
b[0] = byte((tagLen-2)/2) << 3 // M', AAD=0 时 Adata=0
|
||||
b[0] |= byte(2 - 1) // L'=L-1
|
||||
copy(b[1:], nonce[:13])
|
||||
binary.BigEndian.PutUint16(b[14:], uint16(len(plaintext)))
|
||||
block.Encrypt(x[:], b[:])
|
||||
|
||||
// 鉴权明文
|
||||
for offset := 0; offset < len(plaintext); offset += aes.BlockSize {
|
||||
end := offset + aes.BlockSize
|
||||
if end > len(plaintext) {
|
||||
end = len(plaintext)
|
||||
}
|
||||
var blk [aes.BlockSize]byte
|
||||
copy(blk[:], plaintext[offset:end])
|
||||
for i := range x {
|
||||
x[i] ^= blk[i]
|
||||
}
|
||||
block.Encrypt(x[:], x[:])
|
||||
}
|
||||
|
||||
// CTR 流：A_i = L' | nonce | counter_be16
|
||||
var a [aes.BlockSize]byte
|
||||
a[0] = byte(2 - 1)
|
||||
copy(a[1:], nonce[:13])
|
||||
encryptCounter := func(i uint16) [aes.BlockSize]byte {
|
||||
var ai [aes.BlockSize]byte
|
||||
copy(ai[:], a[:])
|
||||
binary.BigEndian.PutUint16(ai[14:], i)
|
||||
var s [aes.BlockSize]byte
|
||||
block.Encrypt(s[:], ai[:])
|
||||
return s
|
||||
}
|
||||
|
||||
ciphertext = make([]byte, len(plaintext))
|
||||
for i, offset := 1, 0; offset < len(plaintext); i, offset = i+1, offset+aes.BlockSize {
|
||||
s := encryptCounter(uint16(i))
|
||||
end := offset + aes.BlockSize
|
||||
if end > len(plaintext) {
|
||||
end = len(plaintext)
|
||||
}
|
||||
for j := offset; j < end; j++ {
|
||||
ciphertext[j] = plaintext[j] ^ s[j-offset]
|
||||
}
|
||||
}
|
||||
|
||||
// auth = T XOR S_0
|
||||
s0 := encryptCounter(0)
|
||||
auth = make([]byte, tagLen)
|
||||
for i := 0; i < tagLen; i++ {
|
||||
auth[i] = x[i] ^ s0[i]
|
||||
}
|
||||
return ciphertext, auth, nil
|
||||
}
|
||||
|
||||
// aesCCMDecrypt 与 encrypt 对称，验证标签后返回明文。仅用于测试与可能的回程解密。
|
||||
func aesCCMDecrypt(key, nonce, ciphertext, auth []byte) ([]byte, error) {
|
||||
if len(nonce) != 13 {
|
||||
return nil, fmt.Errorf("ccm nonce must be 13 bytes")
|
||||
}
|
||||
if len(auth) != 8 {
|
||||
return nil, fmt.Errorf("ccm auth tag must be 8 bytes")
|
||||
}
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 先 CTR 解密
|
||||
var a [aes.BlockSize]byte
|
||||
a[0] = byte(2 - 1)
|
||||
copy(a[1:], nonce[:13])
|
||||
encryptCounter := func(i uint16) [aes.BlockSize]byte {
|
||||
var ai [aes.BlockSize]byte
|
||||
copy(ai[:], a[:])
|
||||
binary.BigEndian.PutUint16(ai[14:], i)
|
||||
var s [aes.BlockSize]byte
|
||||
block.Encrypt(s[:], ai[:])
|
||||
return s
|
||||
}
|
||||
plain := make([]byte, len(ciphertext))
|
||||
for i, offset := 1, 0; offset < len(ciphertext); i, offset = i+1, offset+aes.BlockSize {
|
||||
s := encryptCounter(uint16(i))
|
||||
end := offset + aes.BlockSize
|
||||
if end > len(ciphertext) {
|
||||
end = len(ciphertext)
|
||||
}
|
||||
for j := offset; j < end; j++ {
|
||||
plain[j] = ciphertext[j] ^ s[j-offset]
|
||||
}
|
||||
}
|
||||
|
||||
// 再 CBC-MAC 校验
|
||||
var x [aes.BlockSize]byte
|
||||
var b [aes.BlockSize]byte
|
||||
b[0] = byte((8-2)/2) << 3
|
||||
b[0] |= byte(2 - 1)
|
||||
copy(b[1:], nonce[:13])
|
||||
binary.BigEndian.PutUint16(b[14:], uint16(len(plain)))
|
||||
block.Encrypt(x[:], b[:])
|
||||
for offset := 0; offset < len(plain); offset += aes.BlockSize {
|
||||
end := offset + aes.BlockSize
|
||||
if end > len(plain) {
|
||||
end = len(plain)
|
||||
}
|
||||
var blk [aes.BlockSize]byte
|
||||
copy(blk[:], plain[offset:end])
|
||||
for i := range x {
|
||||
x[i] ^= blk[i]
|
||||
}
|
||||
block.Encrypt(x[:], x[:])
|
||||
}
|
||||
s0 := encryptCounter(0)
|
||||
expected := make([]byte, 8)
|
||||
for i := 0; i < 8; i++ {
|
||||
expected[i] = x[i] ^ s0[i]
|
||||
}
|
||||
if subtle.ConstantTimeCompare(expected, auth) != 1 {
|
||||
return nil, fmt.Errorf("aes-ccm auth mismatch")
|
||||
}
|
||||
return plain, nil
|
||||
}
|
||||
|
||||
// buildPKIMeshPacket 构造一个 PKI 加密的 MeshPacket：
|
||||
// - tag 1/2: from/to (fixed32)
|
||||
// - tag 3 channel = 0 （省略，默认即为 0）
|
||||
// - tag 5 encrypted (含 ciphertext|auth|extraNonce)
|
||||
// - tag 6 packet_id
|
||||
// - tag 14 via_mqtt
|
||||
// - tag 16 public_key（可选，附带发送者公钥）
|
||||
// - tag 17 pki_encrypted = 1
|
||||
func buildPKIMeshPacket(from, to, packetID uint32, viaMQTT bool, encrypted []byte, senderPublic []byte) []byte {
|
||||
var out []byte
|
||||
out = protowire.AppendTag(out, 1, protowire.Fixed32Type)
|
||||
out = protowire.AppendFixed32(out, from)
|
||||
out = protowire.AppendTag(out, 2, protowire.Fixed32Type)
|
||||
out = protowire.AppendFixed32(out, to)
|
||||
out = protowire.AppendTag(out, 5, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, encrypted)
|
||||
out = protowire.AppendTag(out, 6, protowire.Fixed32Type)
|
||||
out = protowire.AppendFixed32(out, packetID)
|
||||
if viaMQTT {
|
||||
out = protowire.AppendTag(out, 14, protowire.VarintType)
|
||||
out = protowire.AppendVarint(out, 1)
|
||||
}
|
||||
if len(senderPublic) == 32 {
|
||||
out = protowire.AppendTag(out, 16, protowire.BytesType)
|
||||
out = protowire.AppendBytes(out, senderPublic)
|
||||
}
|
||||
out = protowire.AppendTag(out, 17, protowire.VarintType)
|
||||
out = protowire.AppendVarint(out, 1)
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
package mqtpp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/ecdh"
|
||||
"crypto/rand"
|
||||
"encoding/binary"
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/encoding/protowire"
|
||||
)
|
||||
|
||||
func TestBuildPKITextMessageRoundTrip(t *testing.T) {
|
||||
curve := ecdh.X25519()
|
||||
senderPriv, err := curve.GenerateKey(rand.Reader)
|
||||
if err != nil {
|
||||
t.Fatalf("generate sender key: %v", err)
|
||||
}
|
||||
recipientPriv, err := curve.GenerateKey(rand.Reader)
|
||||
if err != nil {
|
||||
t.Fatalf("generate recipient key: %v", err)
|
||||
}
|
||||
|
||||
const text = "hello over PKI 你好"
|
||||
const fromNum uint32 = 0x12345678
|
||||
const toNum uint32 = 0xa1b2c3d4
|
||||
const packetID uint32 = 0xdeadbeef
|
||||
|
||||
raw, err := BuildPKITextMessageServiceEnvelope(PKITextMessageBuildOptions{
|
||||
FromNodeNum: fromNum,
|
||||
ToNodeNum: toNum,
|
||||
PacketID: packetID,
|
||||
GatewayID: NodeNumToID(fromNum),
|
||||
ViaMQTT: true,
|
||||
SenderPrivate: senderPriv.Bytes(),
|
||||
RecipientPub: recipientPriv.PublicKey().Bytes(),
|
||||
SenderPublic: senderPriv.PublicKey().Bytes(),
|
||||
Text: text,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildPKITextMessageServiceEnvelope: %v", err)
|
||||
}
|
||||
|
||||
env, err := parseServiceEnvelope(raw)
|
||||
if err != nil {
|
||||
t.Fatalf("parseServiceEnvelope: %v", err)
|
||||
}
|
||||
if env.ChannelID != PKIChannelID {
|
||||
t.Fatalf("channel_id = %q want %q", env.ChannelID, PKIChannelID)
|
||||
}
|
||||
if env.GatewayID != NodeNumToID(fromNum) {
|
||||
t.Fatalf("gateway_id = %q", env.GatewayID)
|
||||
}
|
||||
pkt := env.Packet
|
||||
if pkt.From != fromNum || pkt.To != toNum || pkt.ID != packetID {
|
||||
t.Fatalf("packet header mismatch: %+v", pkt)
|
||||
}
|
||||
if !pkt.PKIEncrypted {
|
||||
t.Fatalf("pki_encrypted = false")
|
||||
}
|
||||
if !pkt.ViaMQTT {
|
||||
t.Fatalf("via_mqtt = false")
|
||||
}
|
||||
if pkt.Channel != 0 {
|
||||
t.Fatalf("channel = %d want 0", pkt.Channel)
|
||||
}
|
||||
if pkt.PayloadVariant != "encrypted" || len(pkt.Encrypted) <= pkcOverhead {
|
||||
t.Fatalf("encrypted payload missing: %+v", pkt)
|
||||
}
|
||||
|
||||
// 收件人用对端私钥 + 发件人公钥推导共享密钥并解密
|
||||
sharedKey, err := pkiSharedKey(recipientPriv.Bytes(), senderPriv.PublicKey().Bytes())
|
||||
if err != nil {
|
||||
t.Fatalf("pkiSharedKey: %v", err)
|
||||
}
|
||||
encryptedLen := len(pkt.Encrypted) - pkcOverhead
|
||||
ciphertext := pkt.Encrypted[:encryptedLen]
|
||||
auth := pkt.Encrypted[encryptedLen : encryptedLen+8]
|
||||
extraNonce := binary.LittleEndian.Uint32(pkt.Encrypted[encryptedLen+8:])
|
||||
plaintext, err := aesCCMDecrypt(sharedKey, pkiNonce(packetID, fromNum, extraNonce), ciphertext, auth)
|
||||
if err != nil {
|
||||
t.Fatalf("aesCCMDecrypt: %v", err)
|
||||
}
|
||||
data, err := parseDataPacket(plaintext)
|
||||
if err != nil {
|
||||
t.Fatalf("parseDataPacket: %v", err)
|
||||
}
|
||||
if data.Portnum != textMessageApp {
|
||||
t.Fatalf("portnum = %d", data.Portnum)
|
||||
}
|
||||
if string(data.Payload) != text {
|
||||
t.Fatalf("text = %q want %q", string(data.Payload), text)
|
||||
}
|
||||
|
||||
// 同样用 MQTTPP 解析路径：PKI 包对外应被识别为 encrypted_packet（无法解密），
|
||||
// 但用错的 PSK 不应误报“channel hash mismatch” 之外的奇怪错误。
|
||||
dummyPSK, _ := ExpandPSK("AQ==")
|
||||
_, _, record := MQTTPP("msh/2/e/PKI/!12345678", raw, dummyPSK, Options{AllowEncryptedForwarding: true})
|
||||
if record["channel_id"] != PKIChannelID {
|
||||
t.Fatalf("MQTTPP record channel_id = %v", record["channel_id"])
|
||||
}
|
||||
if record["pki_encrypted"] != true {
|
||||
t.Fatalf("pki_encrypted record = %v", record["pki_encrypted"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPKINonceLayoutMatchesFirmware(t *testing.T) {
|
||||
// 复刻 firmware initNonce(fromNode, packetId, extraNonce) 期望的字节布局：
|
||||
// nonce[0..8) = packetId(uint64 LE)
|
||||
// nonce[4..8) 被 extraNonce(uint32 LE) 覆盖（当 extraNonce != 0）
|
||||
// nonce[8..12) = fromNode(uint32 LE)
|
||||
// nonce[12] = 0
|
||||
got := pkiNonce(0xaabbccdd, 0x11223344, 0x55667788)
|
||||
want := []byte{
|
||||
0xdd, 0xcc, 0xbb, 0xaa, // packetId low 4 bytes，未被 extraNonce 覆盖前
|
||||
0x88, 0x77, 0x66, 0x55, // extraNonce 覆盖 nonce[4..8)
|
||||
0x44, 0x33, 0x22, 0x11, // fromNode
|
||||
0x00,
|
||||
}
|
||||
if !bytes.Equal(got, want) {
|
||||
t.Fatalf("pkiNonce = % x\nwant % x", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildPKITextMessageRejectsBroadcast(t *testing.T) {
|
||||
curve := ecdh.X25519()
|
||||
priv, _ := curve.GenerateKey(rand.Reader)
|
||||
pub, _ := curve.GenerateKey(rand.Reader)
|
||||
if _, err := BuildPKITextMessageServiceEnvelope(PKITextMessageBuildOptions{
|
||||
FromNodeNum: 0x1,
|
||||
ToNodeNum: NodeNumBroadcast,
|
||||
PacketID: 0x2,
|
||||
SenderPrivate: priv.Bytes(),
|
||||
RecipientPub: pub.PublicKey().Bytes(),
|
||||
Text: "hi",
|
||||
}); err == nil {
|
||||
t.Fatalf("expected error for broadcast destination")
|
||||
}
|
||||
}
|
||||
|
||||
// 确认 MeshPacket 中确实带上 pki_encrypted (tag 17) 与 public_key (tag 16)
|
||||
func TestBuildPKIMeshPacketTags(t *testing.T) {
|
||||
encrypted := []byte{0x01, 0x02, 0x03}
|
||||
pub := make([]byte, 32)
|
||||
for i := range pub {
|
||||
pub[i] = byte(i)
|
||||
}
|
||||
raw := buildPKIMeshPacket(0x11, 0x22, 0x33, true, encrypted, pub)
|
||||
tags := map[protowire.Number]bool{}
|
||||
if err := walkFields(raw, func(num protowire.Number, _ protowire.Type, _ any) error {
|
||||
tags[num] = true
|
||||
return nil
|
||||
}); err != nil {
|
||||
t.Fatalf("walkFields: %v", err)
|
||||
}
|
||||
for _, want := range []protowire.Number{1, 2, 5, 6, 14, 16, 17} {
|
||||
if !tags[want] {
|
||||
t.Fatalf("missing tag %d", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 端到端：发送方构造 PKI 包，接收方通过 PKIKeyResolver 解密并还原文本消息记录。
|
||||
func TestMQTTPPDecryptsPKIWithResolver(t *testing.T) {
|
||||
curve := ecdh.X25519()
|
||||
senderPriv, _ := curve.GenerateKey(rand.Reader)
|
||||
recipientPriv, _ := curve.GenerateKey(rand.Reader)
|
||||
|
||||
const text = "hello PKI inbound"
|
||||
const fromNum uint32 = 0xaaaa1111
|
||||
const toNum uint32 = 0xbbbb2222
|
||||
const packetID uint32 = 0x77777777
|
||||
|
||||
raw, err := BuildPKITextMessageServiceEnvelope(PKITextMessageBuildOptions{
|
||||
FromNodeNum: fromNum,
|
||||
ToNodeNum: toNum,
|
||||
PacketID: packetID,
|
||||
GatewayID: NodeNumToID(fromNum),
|
||||
ViaMQTT: true,
|
||||
SenderPrivate: senderPriv.Bytes(),
|
||||
RecipientPub: recipientPriv.PublicKey().Bytes(),
|
||||
SenderPublic: senderPriv.PublicKey().Bytes(),
|
||||
Text: text,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("build: %v", err)
|
||||
}
|
||||
|
||||
resolver := func(to, from uint32) ([]byte, []byte, bool) {
|
||||
if to != toNum || from != fromNum {
|
||||
return nil, nil, false
|
||||
}
|
||||
return recipientPriv.Bytes(), senderPriv.PublicKey().Bytes(), true
|
||||
}
|
||||
dummyPSK, _ := ExpandPSK("AQ==")
|
||||
valid, _, record := MQTTPP("msh/2/e/PKI/!aaaa1111", raw, dummyPSK, Options{PKIKeyResolver: resolver})
|
||||
if !valid {
|
||||
t.Fatalf("MQTTPP not valid: %#v", record)
|
||||
}
|
||||
if record["type"] != "text_message" {
|
||||
t.Fatalf("type = %v, want text_message", record["type"])
|
||||
}
|
||||
if record["text"] != text {
|
||||
t.Fatalf("text = %v", record["text"])
|
||||
}
|
||||
if record["pki_encrypted"] != true {
|
||||
t.Fatalf("pki_encrypted = %v", record["pki_encrypted"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildPKIAckRoundTrip(t *testing.T) {
|
||||
curve := ecdh.X25519()
|
||||
botPriv, _ := curve.GenerateKey(rand.Reader)
|
||||
devicePriv, _ := curve.GenerateKey(rand.Reader)
|
||||
|
||||
const fromNum uint32 = 0x0000beef // bot
|
||||
const toNum uint32 = 0xfeed0000 // 原 device
|
||||
const ackPacketID uint32 = 0xaaaa5555
|
||||
const requestID uint32 = 0xdeadbeef
|
||||
|
||||
raw, err := BuildPKIAckServiceEnvelope(PKIAckBuildOptions{
|
||||
FromNodeNum: fromNum,
|
||||
ToNodeNum: toNum,
|
||||
PacketID: ackPacketID,
|
||||
RequestID: requestID,
|
||||
GatewayID: NodeNumToID(fromNum),
|
||||
ViaMQTT: true,
|
||||
SenderPrivate: botPriv.Bytes(),
|
||||
RecipientPub: devicePriv.PublicKey().Bytes(),
|
||||
SenderPublic: botPriv.PublicKey().Bytes(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildPKIAckServiceEnvelope: %v", err)
|
||||
}
|
||||
|
||||
// 设备侧解密
|
||||
env, err := parseServiceEnvelope(raw)
|
||||
if err != nil {
|
||||
t.Fatalf("parse: %v", err)
|
||||
}
|
||||
if env.ChannelID != PKIChannelID {
|
||||
t.Fatalf("channel_id = %q", env.ChannelID)
|
||||
}
|
||||
pkt := env.Packet
|
||||
if !pkt.PKIEncrypted || pkt.From != fromNum || pkt.To != toNum || pkt.ID != ackPacketID {
|
||||
t.Fatalf("ack header mismatch: %+v", pkt)
|
||||
}
|
||||
encryptedLen := len(pkt.Encrypted) - pkcOverhead
|
||||
cipher := pkt.Encrypted[:encryptedLen]
|
||||
auth := pkt.Encrypted[encryptedLen : encryptedLen+8]
|
||||
extraNonce := binary.LittleEndian.Uint32(pkt.Encrypted[encryptedLen+8:])
|
||||
sharedKey, err := pkiSharedKey(devicePriv.Bytes(), botPriv.PublicKey().Bytes())
|
||||
if err != nil {
|
||||
t.Fatalf("shared: %v", err)
|
||||
}
|
||||
plain, err := aesCCMDecrypt(sharedKey, pkiNonce(ackPacketID, fromNum, extraNonce), cipher, auth)
|
||||
if err != nil {
|
||||
t.Fatalf("decrypt: %v", err)
|
||||
}
|
||||
data, err := parseDataPacket(plain)
|
||||
if err != nil {
|
||||
t.Fatalf("data: %v", err)
|
||||
}
|
||||
if data.Portnum != routingApp {
|
||||
t.Fatalf("portnum = %d, want ROUTING_APP(%d)", data.Portnum, routingApp)
|
||||
}
|
||||
|
||||
// Routing payload 解析: 期望 oneof error_reason=NONE(0)，即 wire 字节 0x18 0x00
|
||||
wantRouting := []byte{0x18, 0x00}
|
||||
if !bytes.Equal(data.Payload, wantRouting) {
|
||||
t.Fatalf("routing payload = % x, want % x", data.Payload, wantRouting)
|
||||
}
|
||||
}
|
||||
@@ -11,7 +11,7 @@ import (
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"meshtastic_mqtt_server/mqtpp"
|
||||
"meshtastic_mqtt_server/internal/mqtpp"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
package stream
|
||||
|
||||
import "context"
|
||||
|
||||
// Frame represents an event in the stream
|
||||
type Frame struct {
|
||||
Type string `json:"type,omitempty"`
|
||||
Tool string `json:"tool,omitempty"`
|
||||
Stage string `json:"stage,omitempty"`
|
||||
Status string `json:"status,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
Data map[string]any `json:"data,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
}
|
||||
|
||||
// EmitFunc is a function that emits frames
|
||||
type EmitFunc func(frame Frame)
|
||||
|
||||
// ContextKey type for context keys
|
||||
type contextKey string
|
||||
|
||||
const (
|
||||
// TrackerContextKey is the key for the stream tracker in context
|
||||
TrackerContextKey contextKey = "stream_tracker"
|
||||
)
|
||||
|
||||
// Tracker tracks token usage during streaming
|
||||
type Tracker struct {
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
ToolCalls int
|
||||
}
|
||||
|
||||
// NewTracker creates a new stream tracker
|
||||
func NewTracker() *Tracker {
|
||||
return &Tracker{}
|
||||
}
|
||||
|
||||
// AddTool adds tool call token usage
|
||||
func (t *Tracker) AddTool(promptTokens, completionTokens int) {
|
||||
t.PromptTokens += promptTokens
|
||||
t.CompletionTokens += completionTokens
|
||||
t.ToolCalls++
|
||||
}
|
||||
|
||||
// TrackerFromContext retrieves the tracker from context
|
||||
func TrackerFromContext(ctx context.Context) *Tracker {
|
||||
if tracker, ok := ctx.Value(TrackerContextKey).(*Tracker); ok {
|
||||
return tracker
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// WithTracker adds a tracker to the context
|
||||
func WithTracker(ctx context.Context, tracker *Tracker) context.Context {
|
||||
return context.WithValue(ctx, TrackerContextKey, tracker)
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package toolmanager
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"meshtastic_mqtt_server/internal/agenttool"
|
||||
)
|
||||
|
||||
// Manager manages loaded AI tools
|
||||
type Manager struct {
|
||||
tools map[string]agenttool.LoadedTool
|
||||
order []string
|
||||
}
|
||||
|
||||
// Load loads tools from the given directory
|
||||
// If directory doesn't exist or is empty, automatically loads all registered tools
|
||||
func Load(root string, options agenttool.LoadOptions) (*Manager, error) {
|
||||
manager := &Manager{tools: map[string]agenttool.LoadedTool{}}
|
||||
|
||||
// Try to read directory
|
||||
entries, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
if !os.IsNotExist(err) {
|
||||
return nil, fmt.Errorf("failed to read tools directory: %w", err)
|
||||
}
|
||||
// Directory doesn't exist, continue to load all registered tools
|
||||
entries = []os.DirEntry{}
|
||||
}
|
||||
|
||||
// Load tools from directory if they exist
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
name := strings.ToLower(strings.TrimSpace(entry.Name()))
|
||||
descriptor, ok := agenttool.Lookup(name)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
tool, err := descriptor.Load(filepath.Join(root, entry.Name()), options)
|
||||
if err != nil {
|
||||
manager.Close()
|
||||
return nil, fmt.Errorf("failed to load tool %s: %w", name, err)
|
||||
}
|
||||
if tool == nil {
|
||||
continue
|
||||
}
|
||||
toolName := strings.ToLower(strings.TrimSpace(tool.Name()))
|
||||
if toolName == "" {
|
||||
toolName = name
|
||||
}
|
||||
if _, ok := manager.tools[toolName]; ok {
|
||||
manager.Close()
|
||||
return nil, fmt.Errorf("duplicate tool name: %s", toolName)
|
||||
}
|
||||
manager.tools[toolName] = tool
|
||||
manager.order = append(manager.order, toolName)
|
||||
}
|
||||
|
||||
// If no tools loaded from directory, automatically load all registered tools
|
||||
if len(manager.tools) == 0 {
|
||||
registeredTools := agenttool.Names()
|
||||
for _, name := range registeredTools {
|
||||
descriptor, ok := agenttool.Lookup(name)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
// Use empty path for tools that don't require configuration files
|
||||
tool, err := descriptor.Load("", options)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if tool == nil {
|
||||
continue
|
||||
}
|
||||
toolName := strings.ToLower(strings.TrimSpace(tool.Name()))
|
||||
if toolName == "" {
|
||||
toolName = name
|
||||
}
|
||||
if _, ok := manager.tools[toolName]; ok {
|
||||
continue
|
||||
}
|
||||
manager.tools[toolName] = tool
|
||||
manager.order = append(manager.order, toolName)
|
||||
}
|
||||
}
|
||||
return manager, nil
|
||||
}
|
||||
|
||||
// NewForTest creates a manager with preloaded tools for testing
|
||||
func NewForTest(tools ...agenttool.LoadedTool) *Manager {
|
||||
manager := &Manager{tools: map[string]agenttool.LoadedTool{}}
|
||||
for _, tool := range tools {
|
||||
if tool == nil {
|
||||
continue
|
||||
}
|
||||
name := strings.ToLower(strings.TrimSpace(tool.Name()))
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := manager.tools[name]; !ok {
|
||||
manager.order = append(manager.order, name)
|
||||
}
|
||||
manager.tools[name] = tool
|
||||
}
|
||||
return manager
|
||||
}
|
||||
|
||||
// Tools returns all loaded tools
|
||||
func (m *Manager) Tools() []agenttool.LoadedTool {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
tools := make([]agenttool.LoadedTool, 0, len(m.order))
|
||||
for _, name := range m.order {
|
||||
if tool := m.tools[name]; tool != nil {
|
||||
tools = append(tools, tool)
|
||||
}
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
// Get returns a tool by name
|
||||
func (m *Manager) Get(name string) (agenttool.LoadedTool, bool) {
|
||||
if m == nil {
|
||||
return nil, false
|
||||
}
|
||||
tool, ok := m.tools[strings.ToLower(strings.TrimSpace(name))]
|
||||
return tool, ok
|
||||
}
|
||||
|
||||
// RawState returns the raw state of a tool
|
||||
func (m *Manager) RawState(name string) (any, bool) {
|
||||
tool, ok := m.Get(name)
|
||||
if !ok || tool == nil {
|
||||
return nil, false
|
||||
}
|
||||
return tool.RawState(), true
|
||||
}
|
||||
|
||||
// Close closes all tools
|
||||
func (m *Manager) Close() error {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
var errs []string
|
||||
for _, tool := range m.Tools() {
|
||||
if closer, ok := tool.(interface{ Close() error }); ok {
|
||||
if err := closer.Close(); err != nil {
|
||||
errs = append(errs, err.Error())
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(errs) > 0 {
|
||||
sort.Strings(errs)
|
||||
return errors.New(strings.Join(errs, "; "))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
package toolrouter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"meshtastic_mqtt_server/internal/completion"
|
||||
"meshtastic_mqtt_server/internal/llm"
|
||||
"meshtastic_mqtt_server/internal/message"
|
||||
"meshtastic_mqtt_server/internal/stream"
|
||||
"meshtastic_mqtt_server/internal/toolmanager"
|
||||
|
||||
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
||||
)
|
||||
|
||||
const maxAgentToolIterations = 6
|
||||
|
||||
// RunAgentToolLoop runs the agent tool calling loop
|
||||
// systemPrompt is the primary system prompt from LLM config
|
||||
func RunAgentToolLoop(ctx context.Context, state *State, profile *llm.Profile, systemPrompt string, chatMessages []message.ChatMessage, manager *toolmanager.Manager, emit stream.EmitFunc) ([]*model.ChatCompletionMessage, error) {
|
||||
finalMessages, err := buildArkMessages(chatMessages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
routerProfile := profile
|
||||
if state != nil {
|
||||
routerProfile = state.RouterProfile(profile)
|
||||
}
|
||||
tools := availableAgentTools(state, routerProfile, manager, emit)
|
||||
if len(tools) == 0 {
|
||||
// No tools available, add system prompt and return
|
||||
if strings.TrimSpace(systemPrompt) != "" {
|
||||
systemMessage := &model.ChatCompletionMessage{
|
||||
Role: "system",
|
||||
Content: &model.ChatCompletionMessageContent{
|
||||
StringValue: &systemPrompt,
|
||||
},
|
||||
}
|
||||
finalMessages = append([]*model.ChatCompletionMessage{systemMessage}, finalMessages...)
|
||||
}
|
||||
return finalMessages, nil
|
||||
}
|
||||
|
||||
decisionMessages := append([]*model.ChatCompletionMessage(nil), finalMessages...)
|
||||
|
||||
toolByName := make(map[string]AgentTool, len(tools))
|
||||
definitions := make([]*model.Tool, 0, len(tools))
|
||||
availableNames := make([]string, 0, len(tools))
|
||||
toolDescriptions := make([]string, 0, len(tools))
|
||||
for _, tool := range tools {
|
||||
toolByName[tool.name] = tool
|
||||
definitions = append(definitions, tool.definition)
|
||||
availableNames = append(availableNames, tool.name)
|
||||
if tool.definition != nil && tool.definition.Function != nil {
|
||||
toolDescriptions = append(toolDescriptions, fmt.Sprintf("%s: %s", tool.name, tool.definition.Function.Description))
|
||||
}
|
||||
}
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: "agent_tools", Stage: "prepare", Status: "success", Message: "已准备可用工具", Data: map[string]any{"tools": availableNames, "tool_descriptions": toolDescriptions}})
|
||||
}
|
||||
if state == nil || state.cfg == nil {
|
||||
// No tool router config, but we have tools - use primary system prompt
|
||||
if strings.TrimSpace(systemPrompt) != "" {
|
||||
systemMessage := &model.ChatCompletionMessage{
|
||||
Role: "system",
|
||||
Content: &model.ChatCompletionMessageContent{
|
||||
StringValue: &systemPrompt,
|
||||
},
|
||||
}
|
||||
finalMessages = append([]*model.ChatCompletionMessage{systemMessage}, finalMessages...)
|
||||
decisionMessages = append([]*model.ChatCompletionMessage{systemMessage}, decisionMessages...)
|
||||
}
|
||||
return finalMessages, nil
|
||||
}
|
||||
// Use tool router system prompt if available, otherwise fall back to primary system prompt
|
||||
prompt := strings.TrimSpace(state.cfg.SystemPrompt)
|
||||
if prompt == "" {
|
||||
prompt = strings.TrimSpace(systemPrompt)
|
||||
}
|
||||
if prompt != "" {
|
||||
systemMessage := &model.ChatCompletionMessage{
|
||||
Role: "system",
|
||||
Content: &model.ChatCompletionMessageContent{
|
||||
StringValue: &prompt,
|
||||
},
|
||||
}
|
||||
finalMessages = append([]*model.ChatCompletionMessage{systemMessage}, finalMessages...)
|
||||
decisionMessages = append([]*model.ChatCompletionMessage{systemMessage}, decisionMessages...)
|
||||
}
|
||||
for i := 0; i < maxAgentToolIterations; i++ {
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: "agent_tools", Stage: "request", Status: "running", Message: fmt.Sprintf("正在进行第 %d 轮工具判断", i+1), Data: map[string]any{"iteration": i + 1, "max_iterations": maxAgentToolIterations, "tools": availableNames}})
|
||||
}
|
||||
resp, err := completion.CompleteChat(ctx, routerProfile, model.CreateChatCompletionRequest{
|
||||
Model: routerProfile.Config.Model,
|
||||
Messages: decisionMessages,
|
||||
MaxTokens: &state.cfg.MaxTokens,
|
||||
Tools: definitions,
|
||||
ToolChoice: model.ToolChoiceStringTypeAuto,
|
||||
ParallelToolCalls: BoolPtr(false),
|
||||
}, time.Duration(state.cfg.Timeout)*time.Second)
|
||||
if err != nil {
|
||||
return finalMessages, err
|
||||
}
|
||||
if tracker := stream.TrackerFromContext(ctx); tracker != nil {
|
||||
tracker.AddTool(resp.Usage.PromptTokens, resp.Usage.CompletionTokens)
|
||||
}
|
||||
if len(resp.Choices) == 0 {
|
||||
return finalMessages, nil
|
||||
}
|
||||
choice := resp.Choices[0]
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: "agent_tools", Stage: "decision", Status: "success", Message: "工具判断响应已返回", Data: map[string]any{"iteration": i + 1}})
|
||||
}
|
||||
calls := choice.Message.ToolCalls
|
||||
if len(calls) == 0 && choice.Message.FunctionCall != nil {
|
||||
calls = []*model.ToolCall{{ID: "legacy_function_call", Type: model.ToolTypeFunction, Function: *choice.Message.FunctionCall}}
|
||||
}
|
||||
if len(calls) == 0 {
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: "agent_tools", Stage: "request", Status: "success", Message: "模型未请求工具，进入回答生成"})
|
||||
}
|
||||
return finalMessages, nil
|
||||
}
|
||||
callNames := make([]string, 0, len(calls))
|
||||
for _, call := range calls {
|
||||
if call != nil {
|
||||
callNames = append(callNames, call.Function.Name)
|
||||
}
|
||||
}
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: "agent_tools", Stage: "tool_calls", Status: "running", Message: fmt.Sprintf("模型请求调用 %d 个工具", len(calls)), Data: map[string]any{"tools": callNames, "iteration": i + 1}})
|
||||
}
|
||||
assistantMessage := &model.ChatCompletionMessage{Role: "assistant", ToolCalls: calls, Content: choice.Message.Content}
|
||||
finalMessages = append(finalMessages, assistantMessage)
|
||||
decisionMessages = append(decisionMessages, assistantMessage)
|
||||
for _, call := range calls {
|
||||
result := ExecuteAgentToolCall(ctx, call, toolByName, emit)
|
||||
resultContent := &model.ChatCompletionMessageContent{StringValue: &result}
|
||||
toolMessage := &model.ChatCompletionMessage{Role: "tool", ToolCallID: call.ID, Content: resultContent}
|
||||
finalMessages = append(finalMessages, toolMessage)
|
||||
decisionMessages = append(decisionMessages, toolMessage)
|
||||
}
|
||||
}
|
||||
limitText := "工具调用轮数已达到上限。请基于已有工具结果回答，并说明可能未完成全部工具调用。"
|
||||
limitMessage := &model.ChatCompletionMessage{Role: "system", Content: &model.ChatCompletionMessageContent{StringValue: &limitText}}
|
||||
finalMessages = append(finalMessages, limitMessage)
|
||||
return finalMessages, nil
|
||||
}
|
||||
|
||||
func buildArkMessages(chatMessages []message.ChatMessage) ([]*model.ChatCompletionMessage, error) {
|
||||
messages := make([]*model.ChatCompletionMessage, 0, len(chatMessages))
|
||||
for _, msg := range chatMessages {
|
||||
role := msg.Role
|
||||
if role == "" {
|
||||
role = "user"
|
||||
}
|
||||
content := &model.ChatCompletionMessageContent{StringValue: &msg.Content}
|
||||
messages = append(messages, &model.ChatCompletionMessage{
|
||||
Role: role,
|
||||
Content: content,
|
||||
})
|
||||
}
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
// BoolPtr returns a pointer to the given bool
|
||||
func BoolPtr(b bool) *bool {
|
||||
return &b
|
||||
}
|
||||
|
||||
// IntPtr returns a pointer to the given int
|
||||
func IntPtr(i int) *int {
|
||||
return &i
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package toolrouter
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"meshtastic_mqtt_server/internal/llm"
|
||||
)
|
||||
|
||||
// Config holds the tool router configuration
|
||||
type Config struct {
|
||||
Enabled bool
|
||||
OpenAIName string
|
||||
Timeout int
|
||||
MaxTokens int
|
||||
SystemPrompt string
|
||||
}
|
||||
|
||||
// State manages the tool router state
|
||||
type State struct {
|
||||
cfg *Config
|
||||
ai *llm.State
|
||||
}
|
||||
|
||||
// Option is a function that configures the State
|
||||
type Option func(*State)
|
||||
|
||||
// NewState creates a new tool router state
|
||||
func NewState(cfg *Config, ai *llm.State, options ...Option) (*State, error) {
|
||||
if cfg == nil {
|
||||
cfg = &Config{
|
||||
Enabled: true,
|
||||
Timeout: 30,
|
||||
MaxTokens: 512,
|
||||
SystemPrompt: "你可以按需直接调用可用工具来回答用户问题。\n每个工具的 description 描述了它的适用场景和调用条件。\n工具结果优先于模型内置知识；工具失败时必须如实说明，不要编造结果。\n只调用确实必要的工具。",
|
||||
}
|
||||
}
|
||||
if ai == nil {
|
||||
return nil, errors.New("tool router requires an LLM state")
|
||||
}
|
||||
if cfg.Enabled && strings.TrimSpace(cfg.OpenAIName) != "" {
|
||||
if _, err := ai.GetProfile(cfg.OpenAIName); err != nil {
|
||||
return nil, fmt.Errorf("invalid LLM provider name in tool router: %w", err)
|
||||
}
|
||||
}
|
||||
state := &State{cfg: cfg, ai: ai}
|
||||
for _, option := range options {
|
||||
option(state)
|
||||
}
|
||||
return state, nil
|
||||
}
|
||||
|
||||
// RouterProfile returns the LLM profile configured for the tool router
|
||||
func (s *State) RouterProfile(fallback *llm.Profile) *llm.Profile {
|
||||
if s == nil || s.cfg == nil || s.ai == nil {
|
||||
return fallback
|
||||
}
|
||||
name := strings.TrimSpace(s.cfg.OpenAIName)
|
||||
if name == "" {
|
||||
return fallback
|
||||
}
|
||||
profile, err := s.ai.GetProfile(name)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return profile
|
||||
}
|
||||
|
||||
// Config returns a copy of the current configuration
|
||||
func (s *State) Config() Config {
|
||||
if s == nil || s.cfg == nil {
|
||||
return Config{}
|
||||
}
|
||||
return *s.cfg
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package toolrouter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"meshtastic_mqtt_server/internal/agenttool"
|
||||
"meshtastic_mqtt_server/internal/llm"
|
||||
"meshtastic_mqtt_server/internal/stream"
|
||||
"meshtastic_mqtt_server/internal/toolmanager"
|
||||
|
||||
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
||||
)
|
||||
|
||||
// AgentTool represents a tool that can be called by the LLM
|
||||
type AgentTool struct {
|
||||
name string
|
||||
definition *model.Tool
|
||||
execute func(context.Context, string) (string, error)
|
||||
}
|
||||
|
||||
// NewAgentTool creates a new agent tool
|
||||
func NewAgentTool(name string, definition *model.Tool, execute func(context.Context, string) (string, error)) AgentTool {
|
||||
return AgentTool{name: name, definition: definition, execute: execute}
|
||||
}
|
||||
|
||||
// Name returns the tool name
|
||||
func (t AgentTool) Name() string { return t.name }
|
||||
|
||||
// Definition returns the tool definition
|
||||
func (t AgentTool) Definition() *model.Tool { return t.definition }
|
||||
|
||||
// AvailableAgentTools returns all available agent tools
|
||||
func AvailableAgentTools(state *State, profile *llm.Profile, manager *toolmanager.Manager, emit stream.EmitFunc) []AgentTool {
|
||||
return availableAgentTools(state, profile, manager, emit)
|
||||
}
|
||||
|
||||
func availableAgentTools(state *State, profile *llm.Profile, manager *toolmanager.Manager, emit stream.EmitFunc) []AgentTool {
|
||||
if state == nil || state.cfg == nil || !state.cfg.Enabled || manager == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
loaded := manager.Tools()
|
||||
tools := make([]AgentTool, 0, len(loaded))
|
||||
for _, tool := range loaded {
|
||||
if tool == nil {
|
||||
continue
|
||||
}
|
||||
if agentTool, ok := buildAgentTool(tool, "", profile, emit); ok {
|
||||
tools = append(tools, agentTool)
|
||||
}
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
func buildAgentTool(tool agenttool.LoadedTool, description string, profile *llm.Profile, emit stream.EmitFunc) (AgentTool, bool) {
|
||||
if tool == nil || !tool.Enabled() {
|
||||
return AgentTool{}, false
|
||||
}
|
||||
definition := tool.ToolDefinition(description)
|
||||
if definition == nil || definition.Function == nil {
|
||||
return AgentTool{}, false
|
||||
}
|
||||
name := tool.Name()
|
||||
return AgentTool{
|
||||
name: name,
|
||||
definition: definition,
|
||||
execute: func(ctx context.Context, args string) (string, error) {
|
||||
runtime := agenttool.Runtime{
|
||||
Profile: profile,
|
||||
Now: time.Now(),
|
||||
Emit: wrapAgentEmit(emit),
|
||||
}
|
||||
return tool.Execute(ctx, args, runtime)
|
||||
},
|
||||
}, true
|
||||
}
|
||||
|
||||
func wrapAgentEmit(emit stream.EmitFunc) agenttool.EmitFunc {
|
||||
if emit == nil {
|
||||
return nil
|
||||
}
|
||||
return func(frame any) {
|
||||
switch value := frame.(type) {
|
||||
case agenttool.Frame:
|
||||
emit(stream.Frame{Type: value.Type, Tool: value.Tool, Stage: value.Stage, Status: value.Status, Message: value.Message, Data: value.Data, Error: value.Error, Text: value.Text})
|
||||
case stream.Frame:
|
||||
emit(value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ExecuteAgentToolCall executes a tool call and returns the result
|
||||
func ExecuteAgentToolCall(ctx context.Context, call *model.ToolCall, tools map[string]AgentTool, emit stream.EmitFunc) string {
|
||||
if call == nil || call.Type != model.ToolTypeFunction {
|
||||
result := "工具调用无效：仅支持 function 类型工具。"
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: "agent_tools", Stage: "execute", Status: "error", Message: result})
|
||||
}
|
||||
return result
|
||||
}
|
||||
toolName := call.Function.Name
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: toolName, Stage: "arguments", Status: "running", Message: "准备执行工具", Data: map[string]any{"tool_call_id": call.ID, "arguments": call.Function.Arguments}})
|
||||
}
|
||||
tool, ok := tools[toolName]
|
||||
if !ok {
|
||||
result := fmt.Sprintf("工具调用失败：未知工具 %s。", toolName)
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: toolName, Stage: "execute", Status: "error", Message: result})
|
||||
}
|
||||
return result
|
||||
}
|
||||
started := time.Now()
|
||||
result, err := tool.execute(ctx, call.Function.Arguments)
|
||||
durationMs := time.Since(started).Milliseconds()
|
||||
if err != nil {
|
||||
messageText := fmt.Sprintf("工具 %s 执行失败：%v", tool.name, err)
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: tool.name, Stage: "execute", Status: "error", Message: "工具执行失败", Data: map[string]any{"tool_call_id": call.ID, "duration_ms": durationMs, "error": err.Error()}})
|
||||
}
|
||||
return messageText
|
||||
}
|
||||
if strings.TrimSpace(result) == "" {
|
||||
result = fmt.Sprintf("工具 %s 执行完成，但没有返回内容。", tool.name)
|
||||
}
|
||||
if emit != nil {
|
||||
emit(stream.Frame{Type: "trace", Tool: tool.name, Stage: "result", Status: "success", Message: "工具执行完成", Data: map[string]any{"tool_call_id": call.ID, "duration_ms": durationMs, "result_preview": truncateString(result, 1200)}})
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func truncateString(s string, max int) string {
|
||||
if len(s) <= max {
|
||||
return s
|
||||
}
|
||||
return s[:max] + "..."
|
||||
}
|
||||
Reference in New Issue
Block a user