mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-07 03:18:39 +02:00
118 lines
3.2 KiB
Go
118 lines
3.2 KiB
Go
package multiagent
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/cloudwego/eino/components/model"
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
// streamToolCallIndexRepairModel isolates an OpenAI-compatible streaming
|
|
// protocol defect before Eino concatenates response chunks. Some providers
|
|
// reuse a tool-call index for different non-empty tool-call IDs in one stream.
|
|
// Eino correctly rejects that shape because one index represents one call.
|
|
//
|
|
// The wrapper keeps valid streams untouched. When it sees the conflicting
|
|
// shape, it assigns each distinct ID a stable, stream-local index so Eino can
|
|
// retain all calls instead of aborting the agent run.
|
|
type streamToolCallIndexRepairModel struct {
|
|
base model.ToolCallingChatModel
|
|
}
|
|
|
|
func newStreamToolCallIndexRepairModel(base model.ToolCallingChatModel) model.ToolCallingChatModel {
|
|
if base == nil {
|
|
return nil
|
|
}
|
|
return &streamToolCallIndexRepairModel{base: base}
|
|
}
|
|
|
|
func (m *streamToolCallIndexRepairModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
return m.base.Generate(ctx, input, opts...)
|
|
}
|
|
|
|
func (m *streamToolCallIndexRepairModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
stream, err := m.base.Stream(ctx, input, opts...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
state := newStreamToolCallIndexRepairState()
|
|
return schema.StreamReaderWithConvert(stream, state.repairMessage), nil
|
|
}
|
|
|
|
func (m *streamToolCallIndexRepairModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
|
|
withTools, err := m.base.WithTools(tools)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return newStreamToolCallIndexRepairModel(withTools), nil
|
|
}
|
|
|
|
type streamToolCallIndexRepairState struct {
|
|
indexByID map[string]int
|
|
idByIndex map[int]string
|
|
nextFreeIndex int
|
|
}
|
|
|
|
func newStreamToolCallIndexRepairState() *streamToolCallIndexRepairState {
|
|
return &streamToolCallIndexRepairState{
|
|
indexByID: make(map[string]int),
|
|
idByIndex: make(map[int]string),
|
|
}
|
|
}
|
|
|
|
func (s *streamToolCallIndexRepairState) repairMessage(msg *schema.Message) (*schema.Message, error) {
|
|
if msg == nil || len(msg.ToolCalls) == 0 {
|
|
return msg, nil
|
|
}
|
|
|
|
var calls []schema.ToolCall
|
|
changed := false
|
|
for i := range msg.ToolCalls {
|
|
call := msg.ToolCalls[i]
|
|
if call.Index == nil || call.ID == "" {
|
|
continue
|
|
}
|
|
|
|
sourceIndex := *call.Index
|
|
if sourceIndex >= s.nextFreeIndex {
|
|
s.nextFreeIndex = sourceIndex + 1
|
|
}
|
|
|
|
assigned, known := s.indexByID[call.ID]
|
|
if !known {
|
|
assigned = sourceIndex
|
|
if owner, occupied := s.idByIndex[assigned]; occupied && owner != call.ID {
|
|
assigned = s.takeFreeIndex()
|
|
}
|
|
s.indexByID[call.ID] = assigned
|
|
s.idByIndex[assigned] = call.ID
|
|
}
|
|
if assigned == sourceIndex {
|
|
continue
|
|
}
|
|
if calls == nil {
|
|
calls = append([]schema.ToolCall(nil), msg.ToolCalls...)
|
|
}
|
|
index := assigned
|
|
calls[i].Index = &index
|
|
changed = true
|
|
}
|
|
|
|
if !changed {
|
|
return msg, nil
|
|
}
|
|
out := *msg
|
|
out.ToolCalls = calls
|
|
return &out, nil
|
|
}
|
|
|
|
func (s *streamToolCallIndexRepairState) takeFreeIndex() int {
|
|
for {
|
|
candidate := s.nextFreeIndex
|
|
s.nextFreeIndex++
|
|
if _, occupied := s.idByIndex[candidate]; !occupied {
|
|
return candidate
|
|
}
|
|
}
|
|
}
|