Add files via upload

internal/multiagent/eino_stream_tool_call_index.go

@@ -0,0 +1,117 @@
+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
+		}
+	}
+}