mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-05-15 12:58:01 +02:00
公明
85 lines
2.2 KiB
Go
85 lines
2.2 KiB
Go
package multiagent
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"sync"
|
||
|
||
"github.com/cloudwego/eino/adk"
|
||
)
|
||
|
||
// modelFacingTraceHolder 保存「即将送入 ChatModel」的消息快照（已走 summarization / reduction / orphan 修剪等），
|
||
// 用于 last_react_input 落库，使续跑与「上下文压缩后」的模型视角一致，而非仅依赖事件流 append 的 runAccumulatedMsgs。
|
||
type modelFacingTraceHolder struct {
|
||
mu sync.Mutex
|
||
// msgs 为深拷贝后的切片，避免框架后续原地修改污染快照
|
||
msgs []adk.Message
|
||
}
|
||
|
||
func newModelFacingTraceHolder() *modelFacingTraceHolder {
|
||
return &modelFacingTraceHolder{}
|
||
}
|
||
|
||
// Snapshot 返回当前快照的再一次深拷贝（供序列化落库，避免与 holder 互斥长期持锁）。
|
||
func (h *modelFacingTraceHolder) Snapshot() []adk.Message {
|
||
if h == nil {
|
||
return nil
|
||
}
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
return cloneADKMessagesForTrace(h.msgs)
|
||
}
|
||
|
||
func (h *modelFacingTraceHolder) storeFromState(state *adk.ChatModelAgentState) {
|
||
if h == nil || state == nil || len(state.Messages) == 0 {
|
||
return
|
||
}
|
||
cloned := cloneADKMessagesForTrace(state.Messages)
|
||
if len(cloned) == 0 {
|
||
return
|
||
}
|
||
h.mu.Lock()
|
||
h.msgs = cloned
|
||
h.mu.Unlock()
|
||
}
|
||
|
||
func cloneADKMessagesForTrace(msgs []adk.Message) []adk.Message {
|
||
if len(msgs) == 0 {
|
||
return nil
|
||
}
|
||
b, err := json.Marshal(msgs)
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
var out []adk.Message
|
||
if err := json.Unmarshal(b, &out); err != nil {
|
||
return nil
|
||
}
|
||
return out
|
||
}
|
||
|
||
// modelFacingTraceMiddleware 必须在 Handlers 链中处于 **BeforeModel 最后**（telemetry 之后），
|
||
// 此时 state.Messages 即为本次 LLM 调用的最终入参。
|
||
type modelFacingTraceMiddleware struct {
|
||
adk.BaseChatModelAgentMiddleware
|
||
holder *modelFacingTraceHolder
|
||
}
|
||
|
||
func newModelFacingTraceMiddleware(holder *modelFacingTraceHolder) adk.ChatModelAgentMiddleware {
|
||
if holder == nil {
|
||
return nil
|
||
}
|
||
return &modelFacingTraceMiddleware{holder: holder}
|
||
}
|
||
|
||
func (m *modelFacingTraceMiddleware) BeforeModelRewriteState(
|
||
ctx context.Context,
|
||
state *adk.ChatModelAgentState,
|
||
mc *adk.ModelContext,
|
||
) (context.Context, *adk.ChatModelAgentState, error) {
|
||
if m.holder != nil && state != nil {
|
||
m.holder.storeFromState(state)
|
||
}
|
||
return ctx, state, nil
|
||
}
|