fix: bring back metrics on fantasy

Andrey Nering created 1 week ago

Change summary

internal/agent/agent.go |  9 +++++++
internal/agent/event.go | 50 +++++++++++++++++++++++++++++++++++++++++++
2 files changed, 59 insertions(+)

Detailed changes

internal/agent/agent.go

@@ -172,6 +172,9 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
 
 	history, files := a.preparePrompt(msgs, call.Attachments...)
 
+	startTime := time.Now()
+	a.eventPromptSent(call.SessionID)
+
 	var currentAssistant *message.Message
 	var shouldSummarize bool
 	result, err := agent.Stream(genCtx, fantasy.AgentStreamCall{
@@ -354,6 +357,9 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
 			},
 		},
 	})
+
+	a.eventPromptResponded(call.SessionID, time.Since(startTime).Truncate(time.Second))
+
 	if err != nil {
 		isCancelErr := errors.Is(err, context.Canceled)
 		isPermissionErr := errors.Is(err, permission.ErrorPermissionDenied)
@@ -684,6 +690,9 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
 		modelConfig.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
 		modelConfig.CostPer1MIn/1e6*float64(usage.InputTokens) +
 		modelConfig.CostPer1MOut/1e6*float64(usage.OutputTokens)
+
+	a.eventTokensUsed(session.ID, model, usage, cost)
+
 	session.Cost += cost
 	session.CompletionTokens = usage.OutputTokens + usage.CacheReadTokens
 	session.PromptTokens = usage.InputTokens + usage.CacheCreationTokens

internal/agent/event.go

@@ -0,0 +1,50 @@
+package agent
+
+import (
+	"time"
+
+	"charm.land/fantasy"
+	"github.com/charmbracelet/crush/internal/event"
+)
+
+func (a sessionAgent) eventPromptSent(sessionID string) {
+	event.PromptSent(
+		a.eventCommon(sessionID, a.largeModel)...,
+	)
+}
+
+func (a sessionAgent) eventPromptResponded(sessionID string, duration time.Duration) {
+	event.PromptResponded(
+		append(
+			a.eventCommon(sessionID, a.largeModel),
+			"prompt duration pretty", duration.String(),
+			"prompt duration in seconds", int64(duration.Seconds()),
+		)...,
+	)
+}
+
+func (a sessionAgent) eventTokensUsed(sessionID string, model Model, usage fantasy.Usage, cost float64) {
+	event.TokensUsed(
+		append(
+			a.eventCommon(sessionID, model),
+			"input tokens", usage.InputTokens,
+			"output tokens", usage.OutputTokens,
+			"cache read tokens", usage.CacheReadTokens,
+			"cache creation tokens", usage.CacheCreationTokens,
+			"total tokens", usage.InputTokens+usage.OutputTokens+usage.CacheReadTokens+usage.CacheCreationTokens,
+			"cost", cost,
+		)...,
+	)
+}
+
+func (a sessionAgent) eventCommon(sessionID string, model Model) []any {
+	m := model.ModelCfg
+
+	return []any{
+		"session id", sessionID,
+		"provider", m.Provider,
+		"model", m.Model,
+		"reasoning effort", m.ReasoningEffort,
+		"thinking mode", m.Think,
+	}
+}