59 lines
1.5 KiB
Go
59 lines
1.5 KiB
Go
package stream
|
|
|
|
import "context"
|
|
|
|
// Frame represents an event in the stream
|
|
type Frame struct {
|
|
Type string `json:"type,omitempty"`
|
|
Tool string `json:"tool,omitempty"`
|
|
Stage string `json:"stage,omitempty"`
|
|
Status string `json:"status,omitempty"`
|
|
Message string `json:"message,omitempty"`
|
|
Data map[string]any `json:"data,omitempty"`
|
|
Error string `json:"error,omitempty"`
|
|
Text string `json:"text,omitempty"`
|
|
}
|
|
|
|
// EmitFunc is a function that emits frames
|
|
type EmitFunc func(frame Frame)
|
|
|
|
// ContextKey type for context keys
|
|
type contextKey string
|
|
|
|
const (
|
|
// TrackerContextKey is the key for the stream tracker in context
|
|
TrackerContextKey contextKey = "stream_tracker"
|
|
)
|
|
|
|
// Tracker tracks token usage during streaming
|
|
type Tracker struct {
|
|
PromptTokens int
|
|
CompletionTokens int
|
|
ToolCalls int
|
|
}
|
|
|
|
// NewTracker creates a new stream tracker
|
|
func NewTracker() *Tracker {
|
|
return &Tracker{}
|
|
}
|
|
|
|
// AddTool adds tool call token usage
|
|
func (t *Tracker) AddTool(promptTokens, completionTokens int) {
|
|
t.PromptTokens += promptTokens
|
|
t.CompletionTokens += completionTokens
|
|
t.ToolCalls++
|
|
}
|
|
|
|
// TrackerFromContext retrieves the tracker from context
|
|
func TrackerFromContext(ctx context.Context) *Tracker {
|
|
if tracker, ok := ctx.Value(TrackerContextKey).(*Tracker); ok {
|
|
return tracker
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// WithTracker adds a tracker to the context
|
|
func WithTracker(ctx context.Context, tracker *Tracker) context.Context {
|
|
return context.WithValue(ctx, TrackerContextKey, tracker)
|
|
}
|