orchard

 1package main
 2
 3import (
 4	"context"
 5	"fmt"
 6	"os"
 7	"strings"
 8	"testing"
 9	"time"
10
11	"chat.bythewood.me/tools"
12)
13
14// Drives real turns against the real model and the real snapshot, so it runs
15// only when told to. Set PROMPT_SMOKE=1, LLM_KEY, and optionally WIKI_URL.
16func TestPromptsSmoke(t *testing.T) {
17	if os.Getenv("PROMPT_SMOKE") == "" {
18		t.Skip("set PROMPT_SMOKE=1 to drive the real model")
19	}
20	if base := os.Getenv("WIKI_URL"); base != "" {
21		tools.WikiBase = base
22	}
23	llm := NewLLM(env("LLM_URL", "http://orchard-llm:8000"), env("LLM_MODEL", "local"), os.Getenv("LLM_KEY"))
24	eng := NewEngine(llm, env("LLM_NAME", "Ornith 1.5 9B"))
25	// main wires the markdown renderer, and the streaming path calls it.
26	eng.Render = func(md string) string { return md }
27
28	prompts := []string{
29		"who is the prime minister of japan",
30		"what is a goodyear welt",
31	}
32
33	for _, p := range prompts {
34		ctx, cancel := context.WithTimeout(context.Background(), 4*time.Minute)
35		start := time.Now()
36		tr := NewTrace(nil)
37		reply, used, _, _, _, err := eng.Run(ctx, nil, p, "", "", tr, func(Event) {})
38		cancel()
39		if err != nil {
40			t.Errorf("%q: %v", p, err)
41			continue
42		}
43		var calls []string
44		for _, u := range used {
45			s := u.Name
46			if u.Err != "" {
47				s += "(error: " + u.Err + ")"
48			}
49			calls = append(calls, s)
50		}
51		answer := strings.TrimSpace(reply.Content)
52		var steps []string
53		for _, st := range tr.Steps() {
54			flag := ""
55			if st.Bad {
56				flag = "!"
57			}
58			steps = append(steps, st.Kind+flag)
59		}
60		t.Logf("\n--- %q  [%s]\ntools: %s\nsteps: %s\n%s\n",
61			p, time.Since(start).Round(time.Millisecond),
62			strings.Join(calls, ", "), strings.Join(steps, " > "), truncate(answer, 400))
63	}
64}
65
66func truncate(s string, n int) string {
67	if len(s) <= n {
68		return s
69	}
70	return s[:n] + fmt.Sprintf("... [%d more chars]", len(s)-n)
71}