opsxdb/server/internal/ai/harness.go

package ai

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"strings"
)

// HarnessConfig Agent 运行时预算（桌面裁剪版；不做 Evolver）。
type HarnessConfig struct {
	MaxRounds           int // 默认 8，硬帽 16
	MaxToolResultChars  int // 回灌截断
	MaxContextTokens    int // trim 预算（粗估字符/4）
	SameCallLimit       int // 同工具同参上限；0=关闭
	MinRetainMessages   int // trim 至少保留条数（不含 system）
}

// DefaultHarness 与历史 agent.go 常量对齐。
func DefaultHarness() HarnessConfig {
	return HarnessConfig{
		MaxRounds:          defaultAgentMaxRounds,
		MaxToolResultChars: maxToolResultFeedChars,
		MaxContextTokens:   8000,
		SameCallLimit:      3,
		MinRetainMessages:  4,
	}
}

func (h HarnessConfig) ClampRounds(n int) int {
	if n <= 0 {
		n = h.MaxRounds
		if n <= 0 {
			n = defaultAgentMaxRounds
		}
	}
	if n > 16 {
		n = 16
	}
	return n
}

// NewRunID 生成短 run_id，贯穿审计与步骤。
func NewRunID() string {
	var b [8]byte
	_, _ = rand.Read(b[:])
	return "run_" + hex.EncodeToString(b[:])
}

// EstimateTokens 粗估（字符/4）；含 tool_calls JSON。
func EstimateTokens(msgs []Message) int {
	n := 0
	for _, m := range msgs {
		n += len(m.Content)
		if len(m.ToolCalls) > 0 {
			if b, err := json.Marshal(m.ToolCalls); err == nil {
				n += len(b)
			}
		}
		n += len(m.Name) + len(m.ToolCallID) + 8
	}
	if n < 0 {
		return 0
	}
	return n / 4
}

// TrimMessages 保留 system + 最近消息，控制在 MaxContextTokens*90%；
// 并按 tool_call_id 整组保留 assistant.tool_calls 与对应 tool。
func TrimMessages(msgs []Message, h HarnessConfig) []Message {
	if len(msgs) == 0 {
		return msgs
	}
	budget := h.MaxContextTokens
	if budget <= 0 {
		budget = 8000
	}
	limit := budget * 90 / 100
	minRetain := h.MinRetainMessages
	if minRetain <= 0 {
		minRetain = 4
	}

	var system []Message
	var rest []Message
	for _, m := range msgs {
		if m.Role == "system" {
			system = append(system, m)
		} else {
			rest = append(rest, m)
		}
	}
	if EstimateTokens(msgs) <= limit {
		return ensureToolPairing(msgs)
	}

	// 从尾部向前取，直到接近预算
	var kept []Message
	for i := len(rest) - 1; i >= 0; i-- {
		candidate := append([]Message{}, rest[i:]...)
		candidate = append(append([]Message{}, system...), candidate...)
		if EstimateTokens(candidate) > limit && len(kept) >= minRetain {
			break
		}
		kept = rest[i:]
		if len(kept) >= minRetain && EstimateTokens(candidate) <= limit {
			break
		}
	}
	if len(kept) == 0 && len(rest) > 0 {
		// 至少留最后一条
		kept = rest[len(rest)-1:]
	}
	out := append(append([]Message{}, system...), kept...)
	return ensureToolPairing(out)
}

// ensureToolPairing 丢弃没有对应 assistant tool_calls 的孤立 tool，以及裁切后半截的 tool_calls。
func ensureToolPairing(msgs []Message) []Message {
	toolIDs := map[string]bool{}
	for _, m := range msgs {
		if m.Role == "assistant" {
			for _, tc := range m.ToolCalls {
				if tc.ID != "" {
					toolIDs[tc.ID] = true
				}
			}
		}
	}
	out := make([]Message, 0, len(msgs))
	for _, m := range msgs {
		if m.Role == "tool" {
			if m.ToolCallID == "" || !toolIDs[m.ToolCallID] {
				continue
			}
		}
		if m.Role == "assistant" && len(m.ToolCalls) > 0 {
			// 若后续缺少任一 tool 结果，整组 assistant 仍保留（由调用方补 incomplete）；
			// 这里只滤 tool 侧孤儿。
		}
		out = append(out, m)
	}
	return out
}

// ToolCallFingerprint 同参循环检测键。
func ToolCallFingerprint(name, argsJSON string) string {
	name = strings.TrimSpace(name)
	argsJSON = strings.Join(strings.Fields(argsJSON), " ")
	return name + "\n" + argsJSON
}

// DetectRepeatedToolCall 若同一指纹已出现 limit 次（不含本次），返回 true。
func DetectRepeatedToolCall(history []string, fp string, limit int) bool {
	if limit <= 0 || fp == "" {
		return false
	}
	n := 0
	for _, h := range history {
		if h == fp {
			n++
		}
	}
	return n >= limit
}

// CheckpointHint 中段检查点（第 4 / 6 轮，1-based）；空串表示本轮无注入。
func CheckpointHint(round1Based int) string {
	switch round1Based {
	case 4:
		return "[检查点] 若已有足够信息请直接作答；避免重复相同参数的工具调用。"
	case 6:
		return "[检查点] 请尽快收尾给出结论，勿再扩大探索范围。"
	default:
		return ""
	}
}