up

stream/stream.go

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