initial commit

backend/internal/types/message_stats.go

@@ -0,0 +1,50 @@
+package types
+
+import (
+	"time"
+
+	"reichard.io/aethera/pkg/ptr"
+)
+
+type MessageStats struct {
+	StartTime time.Time  `json:"start_time"`
+	EndTime   *time.Time `json:"end_time,omitempty"`
+
+	PromptTokens    *int32 `json:"prompt_tokens"`
+	GeneratedTokens *int32 `json:"generated_tokens"`
+
+	PromptPerSec    *float32 `json:"prompt_per_second"`
+	GeneratedPerSec *float32 `json:"generated_per_second"`
+
+	TimeToFirstToken *int32 `json:"time_to_first_token,omitempty"`
+	TimeToLastToken  *int32 `json:"time_to_last_token,omitempty"`
+}
+
+func (s *MessageStats) RecordFirstToken() {
+	if s.TimeToFirstToken == nil {
+		s.TimeToFirstToken = ptr.Of(int32(time.Since(s.StartTime).Milliseconds()))
+	}
+}
+
+func (s *MessageStats) RecordLastToken() {
+	s.TimeToLastToken = ptr.Of(int32(time.Since(s.StartTime).Milliseconds()))
+}
+
+func (s *MessageStats) CalculateDerived() {
+	// Populate PromptPerSec
+	if s.PromptPerSec == nil && s.TimeToFirstToken != nil && s.PromptTokens != nil {
+		ttft := *s.TimeToFirstToken
+		pt := *s.PromptTokens
+		if ttft > 0 && pt > 0 {
+			s.PromptPerSec = ptr.Of(float32(1000 * pt / ttft))
+		}
+	}
+
+	// Populate GeneratedPerSec
+	if s.GeneratedPerSec == nil && s.TimeToFirstToken != nil && s.TimeToLastToken != nil && s.GeneratedTokens != nil {
+		genTimeMS := *s.TimeToLastToken - *s.TimeToFirstToken
+		if genTimeMS > 0 && *s.GeneratedTokens > 0 {
+			s.GeneratedPerSec = ptr.Of(float32(1000 * float32(*s.GeneratedTokens) / float32(genTimeMS)))
+		}
+	}
+}