cogito/usage_counter_internal_test.go at main · walcz-de/cogito · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
package cogito

import (
	"context"
	"testing"

	"github.com/sashabaranov/go-openai"
)

// fakeLLM is a minimal LLM that returns a fixed usage per CreateChatCompletion
// call and records a fixed usage on the fragment it returns from Ask.
type fakeLLM struct {
	ccUsage  LLMUsage
	askUsage LLMUsage
}

func (f *fakeLLM) CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (LLMReply, LLMUsage, error) {
	return LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{
		Choices: []openai.ChatCompletionChoice{{Message: openai.ChatCompletionMessage{Role: "assistant"}}},
	}}, f.ccUsage, nil
}

func (f *fakeLLM) Ask(ctx context.Context, frag Fragment) (Fragment, error) {
	out := Fragment{Status: &Status{}}
	out.Status.LastUsage = f.askUsage
	return out, nil
}

func TestCountingLLMAccumulatesBothPaths(t *testing.T) {
	inner := &fakeLLM{
		ccUsage:  LLMUsage{PromptTokens: 10, CompletionTokens: 5, TotalTokens: 15},
		askUsage: LLMUsage{PromptTokens: 7, CompletionTokens: 3, TotalTokens: 10},
	}
	counter := &usageCounter{}
	llm := newCountingLLM(inner, counter)

	if _, _, err := llm.CreateChatCompletion(context.Background(), openai.ChatCompletionRequest{}); err != nil {
		t.Fatalf("CreateChatCompletion: %v", err)
	}
	if _, _, err := llm.CreateChatCompletion(context.Background(), openai.ChatCompletionRequest{}); err != nil {
		t.Fatalf("CreateChatCompletion: %v", err)
	}
	if _, err := llm.Ask(context.Background(), NewEmptyFragment()); err != nil {
		t.Fatalf("Ask: %v", err)
	}

	got := counter.snapshot()
	if got.TotalTokens != 40 { // 15 + 15 + 10
		t.Errorf("TotalTokens = %d, want 40", got.TotalTokens)
	}
	if got.PromptTokens != 27 { // 10 + 10 + 7
		t.Errorf("PromptTokens = %d, want 27", got.PromptTokens)
	}
	if got.CompletionTokens != 13 { // 5 + 5 + 3
		t.Errorf("CompletionTokens = %d, want 13", got.CompletionTokens)
	}
}

// streamingFake additionally implements StreamingLLM.
type streamingFake struct{ fakeLLM }

func (s *streamingFake) CreateChatCompletionStream(ctx context.Context, req openai.ChatCompletionRequest) (<-chan StreamEvent, error) {
	ch := make(chan StreamEvent, 1)
	ch <- StreamEvent{Type: StreamEventDone, Usage: LLMUsage{TotalTokens: 99}}
	close(ch)
	return ch, nil
}

func TestNewCountingLLMPreservesStreaming(t *testing.T) {
	plain := newCountingLLM(&fakeLLM{}, &usageCounter{})
	if _, ok := plain.(StreamingLLM); ok {
		t.Error("wrapping a non-streaming LLM must not yield a StreamingLLM")
	}

	streaming := newCountingLLM(&streamingFake{}, &usageCounter{})
	if _, ok := streaming.(StreamingLLM); !ok {
		t.Error("wrapping a StreamingLLM must yield a StreamingLLM")
	}
}