-
-func TestReadUsageFileSortsByCallStart(t *testing.T) {
- path := filepath.Join(t.TempDir(), "usage.jsonl")
- usage := ModelUsageLog{
- path: path,
- }
-
- for _, date := range []string{"2026-01-01T12:00:02Z", "2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z"} {
- require.NoError(t, usage.Write(ModelUsageRecord{
- Date: date,
- }))
- }
-
- records, err := probeReadUsage(path)
- require.NoError(t, err)
- require.Equal(t, []string{"2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z", "2026-01-01T12:00:02Z"}, []string{records[0].Date, records[1].Date, records[2].Date})
-}
-
-func TestTimeTextCountsTheCallsAtOnce(t *testing.T) {
- start := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC)
- record := func(offset, seconds float64) ModelUsageRecord {
- return ModelUsageRecord{
- Date: start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano),
- Seconds: seconds,
- }
- }
-
- text := probeTimeText([]ModelUsageRecord{record(0, 10), record(2, 4), record(3, 2), record(10, 10)})
- require.Equal(t, "time: 20 s in all, 12.0 calls per minute, 6.5 s per call (median 10.0, max 10.0), calls at once: at most 3, 1.3 on average", text)
-}
-
-// ProbeWriter writes readable made-up messages.
-type ProbeWriter struct {
- random *rand.Rand
-}
-
-func (g *ProbeWriter) message() (string, string) {
- roll := g.random.IntN(100)
- switch {
- case roll < 5:
- return "user", g.fill(g.userLine, 10+g.random.IntN(200))
- case roll < 10:
- return "reply", g.fill(g.replyLine, 100+g.random.IntN(1500))
- case roll < 12:
- return "note", g.fill(g.noteLine, 50+g.random.IntN(600))
- case roll < 13:
- return "work", "[" + g.pick(probePeople) + "] report: " + g.fill(g.replyLine, 300+g.random.IntN(3000))
- case roll < 56:
- return "tool", g.toolCall(60 + g.random.IntN(400))
- case roll < 57:
- return "echo", g.fill(g.outputLine, 30000+g.random.IntN(30000))
- default:
- return "echo", g.fill(g.outputLine, 20+g.random.IntN(4000))
- }
-}
-
-// fill joins lines of write until the text has about size bytes.
-func (g *ProbeWriter) fill(write func() string, size int) string {
- var text strings.Builder
- for text.Len() < size {
- if text.Len() > 0 {
- text.WriteByte(' ')
- }
-
- text.WriteString(write())
- }
-
- return text.String()
-}
-
-func (g *ProbeWriter) userLine() string {
- return g.pick([]string{
- fmt.Sprintf("Please check why the %s %s.", g.pick(probeThings), g.pick(probeProblems)),
- fmt.Sprintf("Move the %s to %s before %s.", g.pick(probeThings), g.pick(probePlaces), g.pick(probeDays)),
- fmt.Sprintf("Thanks, that works. Next, %s.", g.pick(probeTasks)),
- fmt.Sprintf("No, keep the old %s. %s asked for it.", g.pick(probeThings), g.pick(probePeople)),
- fmt.Sprintf("How many %s orders came in on %s?", g.pick(probeBreads), g.pick(probeDays)),
- })
-}
-
-func (g *ProbeWriter) replyLine() string {
- return g.pick([]string{
- fmt.Sprintf("I looked at %s. The %s %s because %s.", g.pick(probeFiles), g.pick(probeThings), g.pick(probeProblems), g.pick(probeCauses)),
- fmt.Sprintf("I changed %s, and the tests pass now.", g.pick(probeChanges)),
- fmt.Sprintf("There were %d %s orders on %s, %d of them late.", 10+g.random.IntN(90), g.pick(probeBreads), g.pick(probeDays), g.random.IntN(10)),
- fmt.Sprintf("Open question: should %s get %s?", g.pick(probePeople), g.pick(probeTasks)),
- })
-}
-
-func (g *ProbeWriter) noteLine() string {
- return g.pick([]string{
- fmt.Sprintf("%s prefers %s.", g.pick(probePeople), g.pick(probePreferences)),
- fmt.Sprintf("The bakery opens at %d and closes at %d on %s.", 5+g.random.IntN(3), 17+g.random.IntN(4), g.pick(probeDays)),
- })
-}
-
-func (g *ProbeWriter) pick(words []string) string {
- return words[g.random.IntN(len(words))]
-}
-
-// toolCall returns a tool name and its JSON input of about size bytes.
-func (g *ProbeWriter) toolCall(size int) string {
- name := g.pick([]string{"shell", "read", "write", "edit"})
- var input map[string]string
- switch name {
- case "shell":
- input = map[string]string{"command": "grep -n '" + g.pick(probeThings) + "' " + g.pick(probeFiles)}
- case "read":
- input = map[string]string{"path": g.pick(probeFiles)}
- default:
- input = map[string]string{
- "path": g.pick(probeFiles),
- "text": g.fill(g.replyLine, max(0, size-80)),
- }
- }
-
- data, err := json.Marshal(input)
- if err != nil {
- panic(err)
- }
-
- return name + " " + string(data)
-}
-
-func (g *ProbeWriter) outputLine() string {
- return g.pick([]string{
- fmt.Sprintf("%02d:%02d order %d for %d %s, packed by %s.", 6+g.random.IntN(12), g.random.IntN(60), 1000+g.random.IntN(9000), 1+g.random.IntN(24), g.pick(probeBreads), g.pick(probePeople)),
- fmt.Sprintf("%s:%d: %s %s.", g.pick(probeFiles), 1+g.random.IntN(400), g.pick(probeThings), g.pick(probeProblems)),
- fmt.Sprintf("ok bakery/%s %d.%03ds", g.pick(probeModules), g.random.IntN(3), g.random.IntN(1000)),
- })
-}
-
-// probeLog returns count messages of made-up work on a small bakery order
-// system, with the mix of kinds and sizes of syntheticMessage. The same count
-// always gives the same log.
-func probeLog(count int) []ProbeEntry {
- g := ProbeWriter{
- random: rand.New(rand.NewPCG(7, 11)),
- }
-
- messages := make([]ProbeEntry, count)
- for id := range messages {
- kind, text := g.message()
- messages[id] = ProbeEntry{
- Kind: kind,
- Text: text,
- }
- }
-
- return messages
-}
-
-func newTestProbeRun(t *testing.T, server *ModelFake) *ProbeRun {
- t.Helper()
-
- root := t.TempDir()
- usagePath := filepath.Join(root, "usage.jsonl")
- client := newModelClient(server.URL, ModelAPIKey(modelTestKey), "onechat-"+probeName, usagePath)
- memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probePrompt, defaultConfig().Compactor))
- require.NoError(t, err)
-
- return &ProbeRun{
- memory: memory,
- client: client,
- usagePath: usagePath,
- }
-}
-
-var (
- probePeople = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
- probeBreads = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}
- probeDays = []string{"Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"}
- probeThings = []string{"order form", "price list", "delivery route", "stock count", "invoice export", "oven schedule"}
- probeProblems = []string{"shows the wrong total", "skips the last row", "is slow at noon", "loses the customer note", "rounds the tax down"}
- probePlaces = []string{"the front page", "the weekly report", "the shared folder", "the admin menu"}
- probeTasks = []string{"add a field for allergies", "send the invoices on Friday", "split the route in two", "print labels for the shelf"}
- probeCauses = []string{"the date has no time zone", "the cache keeps old prices", "the loop stops one row early", "the form trims the note"}
- probeChanges = []string{"the total to use cents", "the route to sort by street", "the export to add the tax line", "the form to keep the note"}
- probePreferences = []string{"short answers", "a call before a change to prices", "invoices as PDF", "reports in the morning"}
- probeFiles = []string{"orders/form.go", "prices/list.go", "routes/plan.go", "stock/count.go", "export/invoice.go", "ovens/schedule.go"}
- probeModules = []string{"orders", "prices", "routes", "stock", "export", "ovens"}
-)
projects / onechat / blobdiff