77 lines
1.8 KiB
Go
77 lines
1.8 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"sync/atomic"
|
|
|
|
"ollie/backend"
|
|
)
|
|
|
|
// mockResponse defines a canned response for sequentialStream.
|
|
type mockResponse struct {
|
|
content string
|
|
toolCalls []backend.ToolCall
|
|
stopReason string
|
|
}
|
|
|
|
// sequentialStream returns a respond function that plays back responses in order.
|
|
func sequentialStream(responses []mockResponse) func(context.Context, []backend.Message, []backend.Tool, backend.GenerationParams) (<-chan backend.StreamEvent, error) {
|
|
var n int32
|
|
return func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) {
|
|
i := int(atomic.AddInt32(&n, 1)) - 1
|
|
if i >= len(responses) {
|
|
return textStream("done"), nil
|
|
|
|
}
|
|
r := responses[i]
|
|
ch := make(chan backend.StreamEvent, 1)
|
|
ch <- backend.StreamEvent{
|
|
Content: r.content,
|
|
ToolCalls: r.toolCalls,
|
|
Done: true,
|
|
StopReason: r.stopReason,
|
|
}
|
|
close(ch)
|
|
return ch, nil
|
|
|
|
}
|
|
}
|
|
|
|
// simpleState is a minimal state implementation for direct run() calls.
|
|
type simpleState struct {
|
|
msgs []backend.Message
|
|
}
|
|
|
|
func newState() *simpleState {
|
|
return &simpleState{}
|
|
}
|
|
|
|
func (s *simpleState) history() []backend.Message {
|
|
return s.msgs
|
|
}
|
|
|
|
func (s *simpleState) taskState() *TaskState { return nil }
|
|
|
|
func (s *simpleState) updateTaskState(TaskState) {}
|
|
|
|
func (s *simpleState) update(msg backend.Message, results []toolResult) {
|
|
s.msgs = append(s.msgs, msg)
|
|
for _, r := range results {
|
|
s.msgs = append(s.msgs, backend.Message{
|
|
Role: "tool",
|
|
Content: r.Content,
|
|
ToolCallID: r.ToolCallID,
|
|
})
|
|
}
|
|
}
|
|
|
|
func (s *simpleState) estimateTokens() int {
|
|
chars := 0
|
|
for _, m := range s.msgs {
|
|
chars += len(m.Content)
|
|
}
|
|
return chars / 4
|
|
}
|
|
|
|
func (s *simpleState) stripCold(_ context.Context, _ backend.Backend) {}
|