| 412 | } |
| 413 | |
| 414 | func TestRun_RecordsMetrics(t *testing.T) { |
| 415 | t.Parallel() |
| 416 | |
| 417 | reg := prometheus.NewRegistry() |
| 418 | metrics := chatloop.NewMetrics(reg) |
| 419 | |
| 420 | model := &chattest.FakeModel{ |
| 421 | ProviderName: "test-provider", |
| 422 | ModelName: "test-model", |
| 423 | StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { |
| 424 | return func(yield func(fantasy.StreamPart) bool) { |
| 425 | parts := []fantasy.StreamPart{ |
| 426 | {Type: fantasy.StreamPartTypeTextStart, ID: "t1"}, |
| 427 | {Type: fantasy.StreamPartTypeTextDelta, ID: "t1", Delta: "hello"}, |
| 428 | {Type: fantasy.StreamPartTypeTextEnd, ID: "t1"}, |
| 429 | {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, |
| 430 | } |
| 431 | for _, p := range parts { |
| 432 | if !yield(p) { |
| 433 | return |
| 434 | } |
| 435 | } |
| 436 | }, nil |
| 437 | }, |
| 438 | } |
| 439 | |
| 440 | err := chatloop.Run(context.Background(), chatloop.RunOptions{ |
| 441 | Model: model, |
| 442 | Messages: []fantasy.Message{ |
| 443 | { |
| 444 | Role: fantasy.MessageRoleUser, |
| 445 | Content: []fantasy.MessagePart{ |
| 446 | fantasy.TextPart{Text: "hello"}, |
| 447 | }, |
| 448 | }, |
| 449 | }, |
| 450 | MaxSteps: 1, |
| 451 | PersistStep: func(_ context.Context, _ chatloop.PersistedStep) error { |
| 452 | return nil |
| 453 | }, |
| 454 | Metrics: metrics, |
| 455 | }) |
| 456 | require.NoError(t, err) |
| 457 | |
| 458 | families, err := reg.Gather() |
| 459 | require.NoError(t, err) |
| 460 | |
| 461 | assertProviderModelLabels := func(t *testing.T, metric *dto.Metric) { |
| 462 | t.Helper() |
| 463 | labels := map[string]string{} |
| 464 | for _, lp := range metric.GetLabel() { |
| 465 | labels[lp.GetName()] = lp.GetValue() |
| 466 | } |
| 467 | assert.Equal(t, "test-provider", labels["provider"]) |
| 468 | assert.Equal(t, "test-model", labels["model"]) |
| 469 | } |
| 470 | |
| 471 | found := make(map[string]bool) |