Function TestSessionFromEventsWithPerMessageUsage

pkg/evaluation/save_test.go:599–645

(t *testing.T)

}

func TestSessionFromEventsWithPerMessageUsage(t *testing.T) {
	t.Parallel()

	events := []map[string]any{
		{"type": "agent_choice", "content": "Hello!", "agent_name": "root"},
		{
			"type": "token_usage",
			"usage": map[string]any{
				"input_tokens":  float64(100),
				"output_tokens": float64(50),
				"cost":          0.005,
				"last_message": map[string]any{
					"input_tokens":        float64(100),
					"output_tokens":       float64(50),
					"cached_input_tokens": float64(25),
					"Model":               "gpt-4o",
					"Cost":                0.005,
				},
			},
		},
		{"type": "stream_stopped"},
	}

	sess := SessionFromEvents(events, "test", []string{"hi"})

	// Check session-level usage
	assert.Equal(t, int64(100), sess.InputTokens)
	assert.Equal(t, int64(50), sess.OutputTokens)
	assert.InDelta(t, 0.005, sess.Cost, 0.0001)

	// Find the assistant message
	var assistantMsg *session.Message
	for _, item := range sess.Messages {
		if item.Message != nil && item.Message.Message.Role == chat.MessageRoleAssistant {
			assistantMsg = item.Message
			break
		}
	}

	require.NotNil(t, assistantMsg, "should have assistant message")
	assert.Equal(t, "gpt-4o", assistantMsg.Message.Model)
	assert.InDelta(t, 0.005, assistantMsg.Message.Cost, 0.0001)
	require.NotNil(t, assistantMsg.Message.Usage)
	assert.Equal(t, int64(100), assistantMsg.Message.Usage.InputTokens)
	assert.Equal(t, int64(50), assistantMsg.Message.Usage.OutputTokens)
	assert.Equal(t, int64(25), assistantMsg.Message.Usage.CachedInputTokens)
}

func TestSessionFromEventsWithError(t *testing.T) {
	t.Parallel()