LaodingBot/internal/agent/orchestrator.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strings"

	"laodingbot/internal/knowledge"
	"laodingbot/internal/llm"
	"laodingbot/internal/logger"
	"laodingbot/internal/memory"
	"laodingbot/internal/tools"
)

type Orchestrator struct {
	llm          llm.Client
	store        *memory.SQLiteStore
	tools        *tools.Registry
	soul         string
	skills       []knowledge.Skill
	skillsDoc    string
	reactMaxStep int
	log          *logger.Logger
}

func NewOrchestrator(
	llmClient llm.Client,
	store *memory.SQLiteStore,
	registry *tools.Registry,
	soul string,
	skills []knowledge.Skill,
	skillsDoc string,
	reactMaxStep int,
	log *logger.Logger,
) *Orchestrator {
	if reactMaxStep <= 0 {
		reactMaxStep = 4
	}
	return &Orchestrator{
		llm:          llmClient,
		store:        store,
		tools:        registry,
		soul:         soul,
		skills:       skills,
		skillsDoc:    skillsDoc,
		reactMaxStep: reactMaxStep,
		log:          log,
	}
}

func (o *Orchestrator) HandleMessage(ctx context.Context, chatID, userID, text string) (string, error) {
	if o.log != nil {
		o.log.Infof("handle message chat_id=%s user_id=%s text_len=%d", chatID, userID, len(text))
		o.log.Debugf("handle message text=%q", text)
	}
	if err := o.store.SaveMessage(chatID, userID, "user", text); err != nil {
		if o.log != nil {
			o.log.Errorf("save user message failed chat_id=%s err=%v", chatID, err)
		}
		return "", err
	}

	recent, err := o.store.LoadRecent(chatID, 16)
	if err != nil {
		if o.log != nil {
			o.log.Errorf("load recent failed chat_id=%s err=%v", chatID, err)
		}
		return "", err
	}
	compressed := memory.CompressForPrompt(recent, 6000)
	if o.log != nil {
		o.log.Debugf("prompt context prepared chat_id=%s recent_count=%d compressed_len=%d", chatID, len(recent), len(compressed))
	}

	matchedSkills := o.matchSkills(ctx, compressed, text)

	var response string
	if len(matchedSkills) == 0 {
		if o.log != nil {
			o.log.Infof("no skill matched; use direct llm chat_id=%s", chatID)
		}
		response, err = o.runDirectLLM(ctx, compressed, text)
	} else {
		if o.log != nil {
			names := make([]string, 0, len(matchedSkills))
			for _, s := range matchedSkills {
				names = append(names, s.Name)
				o.log.Infof("skill selected name=%s source=%s", s.Name, s.Source)
				o.log.Debugf("skill selected content name=%s content=%q", s.Name, s.Content)
			}
			o.log.Infof("skills matched chat_id=%s skills=%s", chatID, strings.Join(names, ","))
		}
		response, err = o.runReAct(ctx, compressed, text, matchedSkills)
	}
	if err != nil {
		if o.log != nil {
			o.log.Errorf("message generation failed chat_id=%s err=%v", chatID, err)
		}
		return "", err
	}

	if err := o.store.SaveMessage(chatID, userID, "assistant", response); err != nil {
		if o.log != nil {
			o.log.Errorf("save assistant response failed chat_id=%s err=%v", chatID, err)
		}
		return "", err
	}
	if o.log != nil {
		o.log.Infof("message handled chat_id=%s response_len=%d", chatID, len(response))
	}
	return response, nil
}

func (o *Orchestrator) runDirectLLM(ctx context.Context, compressedContext, userInput string) (string, error) {
	systemPrompt := strings.Join([]string{
		"你是一个个人自动化助手，必须遵循如下人格设定并保持一致：",
		o.soul,
		"",
		"如果当前问题没有匹配到已定义技能，请直接回答用户。",
		"当你判断必须依赖外部工具结果才能可靠回答时，请明确告知用户需要进一步操作信息。",
	}, "\n")

	userPrompt := strings.Join([]string{
		"历史上下文：",
		compressedContext,
		"",
		"用户问题：",
		userInput,
	}, "\n")

	return o.llm.Generate(ctx, systemPrompt, userPrompt)
}

type reactDecision struct {
	Thought     string `json:"thought"`
	Action      string `json:"action"`
	ActionInput string `json:"action_input"`
	Final       string `json:"final"`
}

func (o *Orchestrator) runReAct(ctx context.Context, compressedContext, userInput string, selectedSkills []knowledge.Skill) (string, error) {
	selectedSkillsDoc := formatSkills(selectedSkills)
	toolDoc := o.formatToolDoc()
	if o.log != nil {
		names := make([]string, 0, len(selectedSkills))
		for _, s := range selectedSkills {
			names = append(names, s.Name)
		}
		o.log.Infof("react start steps=%d skills=%s", o.reactMaxStep, strings.Join(names, ","))
		o.log.Debugf("react selected_skills_doc=%q", selectedSkillsDoc)
		o.log.Debugf("react tools_doc=%q", toolDoc)
	}

	systemPrompt := strings.Join([]string{
		"你是一个个人自动化助手，必须遵循如下人格设定并保持一致：",
		o.soul,
		"",
		"已匹配到的 skills（只可按下列技能执行）：",
		selectedSkillsDoc,
		"",
		"可用工具：",
		toolDoc,
		"",
		"你必须使用 ReAct 模式做决策。",
		"只有当技能明确需要工具能力时才调用工具。",
		"如果问题可直接回答，不要调用工具。",
		"你的输出必须是 JSON，对象字段为 thought, action, action_input, final。",
		"规则：",
		"1) 当需要调工具时：final 置空，action 必须是可用工具之一，action_input 为工具输入。",
		"2) 当可以最终回答时：action 置 none，action_input 置空，final 填最终回复。",
		"3) 不要输出 JSON 之外内容。",
	}, "\n")

	scratchpad := ""
	for step := 1; step <= o.reactMaxStep; step++ {
		if o.log != nil {
			o.log.Infof("react step start step=%d/%d", step, o.reactMaxStep)
			o.log.Debugf("react scratchpad_before step=%d content=%q", step, scratchpad)
		}
		prompt := strings.Join([]string{
			"历史上下文：",
			compressedContext,
			"",
			"用户问题：",
			userInput,
			"",
			"当前推理记录（按时间顺序）：",
			scratchpad,
			"",
			fmt.Sprintf("请输出下一步 JSON 决策。当前步骤: %d/%d", step, o.reactMaxStep),
		}, "\n")

		raw, err := o.llm.Generate(ctx, systemPrompt, prompt)
		if err != nil {
			return "", err
		}
		if o.log != nil {
			o.log.Infof("react step llm output step=%d raw=%q", step, raw)
		}
		decision, err := parseDecision(raw)
		if err != nil {
			if o.log != nil {
				o.log.Warnf("react parse failed, use raw as final err=%v", err)
			}
			return strings.TrimSpace(raw), nil
		}
		if o.log != nil {
			o.log.Infof("react step decision step=%d thought=%q action=%q action_input=%q final=%q", step, decision.Thought, decision.Action, decision.ActionInput, decision.Final)
		}

		action := strings.ToLower(strings.TrimSpace(decision.Action))
		if action == "" {
			action = "none"
		}

		if action == "none" {
			finalText := strings.TrimSpace(decision.Final)
			if finalText == "" {
				finalText = "我已完成思考，但当前没有足够信息给出稳定结论。"
			}
			if o.log != nil {
				o.log.Infof("react final step=%d final=%q", step, finalText)
			}
			return finalText, nil
		}

		tool, ok := o.tools.Get(action)
		if !ok {
			if o.log != nil {
				o.log.Warnf("react step tool missing step=%d tool=%s", step, action)
			}
			scratchpad += fmt.Sprintf("Step %d Thought: %s\nStep %d Observation: tool %s 不存在\n", step, decision.Thought, step, action)
			continue
		}

		toolOut, toolErr := tool.Call(ctx, decision.ActionInput)
		if o.log != nil {
			o.log.Infof("react step tool call step=%d tool=%s input=%q", step, action, decision.ActionInput)
		}
		obs := strings.TrimSpace(toolOut)
		if obs == "" {
			obs = "(empty output)"
		}
		if toolErr != nil {
			obs = obs + "\nERROR: " + toolErr.Error()
		}
		if o.log != nil {
			o.log.Infof("react step observation step=%d tool=%s observation=%q", step, action, obs)
		}
		if len(obs) > 2000 {
			obs = obs[:2000]
		}
		scratchpad += fmt.Sprintf("Step %d Thought: %s\nStep %d Action: %s\nStep %d ActionInput: %s\nStep %d Observation: %s\n", step, decision.Thought, step, action, step, decision.ActionInput, step, obs)
	}

	return "我尝试了多轮思考与工具调用，但仍未得到稳定结论。请给我更具体的约束或允许我继续尝试。", nil
}

func (o *Orchestrator) matchSkills(ctx context.Context, compressedContext, userInput string) []knowledge.Skill {
	if len(o.skills) == 0 {
		return nil
	}

	type skillChoice struct {
		Skills []string `json:"skills"`
	}

	systemPrompt := strings.Join([]string{
		"你是技能路由器。",
		"任务：根据用户问题，从候选技能中选择 0-2 个最相关技能名称。",
		"输出必须是 JSON：{\"skills\":[\"name1\",\"name2\"]}",
		"如果没有匹配技能，返回 {\"skills\":[]}。",
		"不要输出 JSON 之外内容。",
	}, "\n")

	userPrompt := strings.Join([]string{
		"候选技能：",
		formatSkillCatalog(o.skills),
		"",
		"历史上下文：",
		compressedContext,
		"",
		"用户问题：",
		userInput,
	}, "\n")

	raw, err := o.llm.Generate(ctx, systemPrompt, userPrompt)
	if err != nil {
		if o.log != nil {
			o.log.Warnf("skill match llm failed err=%v", err)
		}
		return nil
	}
	if o.log != nil {
		o.log.Infof("skill router output raw=%q", raw)
	}

	raw = normalizeJSON(raw)
	choice := skillChoice{}
	if err := json.Unmarshal([]byte(raw), &choice); err != nil {
		if o.log != nil {
			o.log.Warnf("skill match parse failed err=%v", err)
		}
		return nil
	}

	picked := make([]knowledge.Skill, 0, 2)
	seen := map[string]struct{}{}
	for _, name := range choice.Skills {
		name = strings.TrimSpace(strings.ToLower(name))
		if name == "" {
			continue
		}
		if _, ok := seen[name]; ok {
			continue
		}
		for _, skill := range o.skills {
			if strings.ToLower(strings.TrimSpace(skill.Name)) == name {
				picked = append(picked, skill)
				seen[name] = struct{}{}
				break
			}
		}
		if len(picked) >= 2 {
			break
		}
	}
	if o.log != nil {
		names := make([]string, 0, len(picked))
		for _, s := range picked {
			names = append(names, s.Name)
		}
		o.log.Infof("skill router selected skills=%s", strings.Join(names, ","))
	}

	return picked
}

func parseDecision(raw string) (reactDecision, error) {
	raw = normalizeJSON(raw)
	start := strings.Index(raw, "{")
	end := strings.LastIndex(raw, "}")
	if start < 0 || end < start {
		return reactDecision{}, fmt.Errorf("no json object found")
	}
	raw = raw[start : end+1]

	var out reactDecision
	if err := json.Unmarshal([]byte(raw), &out); err != nil {
		return reactDecision{}, err
	}
	return out, nil
}

func normalizeJSON(raw string) string {
	raw = strings.TrimSpace(raw)
	raw = strings.TrimPrefix(raw, "```json")
	raw = strings.TrimPrefix(raw, "```")
	raw = strings.TrimSuffix(raw, "```")
	return strings.TrimSpace(raw)
}

func formatSkills(skills []knowledge.Skill) string {
	b := strings.Builder{}
	for _, skill := range skills {
		b.WriteString("## ")
		b.WriteString(skill.Name)
		b.WriteString("\n")
		b.WriteString(skill.Content)
		b.WriteString("\n\n")
	}
	return strings.TrimSpace(b.String())
}

func formatSkillCatalog(skills []knowledge.Skill) string {
	b := strings.Builder{}
	for _, skill := range skills {
		summary := strings.ReplaceAll(skill.Content, "\n", " ")
		summary = strings.TrimSpace(summary)
		if len(summary) > 220 {
			summary = summary[:220]
		}
		b.WriteString("- ")
		b.WriteString(skill.Name)
		if summary != "" {
			b.WriteString(": ")
			b.WriteString(summary)
		}
		b.WriteString("\n")
	}
	return strings.TrimSpace(b.String())
}

func (o *Orchestrator) formatToolDoc() string {
	list := o.tools.List()
	if len(list) == 0 {
		return "(none)"
	}
	sort.Slice(list, func(i, j int) bool {
		return list[i].Name() < list[j].Name()
	})
	b := strings.Builder{}
	for _, t := range list {
		b.WriteString("- ")
		b.WriteString(t.Name())
		b.WriteString(": ")
		b.WriteString(t.Description())
		b.WriteString("\n")
	}
	return strings.TrimSpace(b.String())
}