133 lines
3.3 KiB
Go
133 lines
3.3 KiB
Go
package processing
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestStripJSONFences(t *testing.T) {
|
|
in := "```json\n{\"a\":1}\n```"
|
|
got := StripJSONFences(in)
|
|
if got != `{"a":1}` {
|
|
t.Fatalf("got=%q", got)
|
|
}
|
|
}
|
|
|
|
func TestParseJSONObject_fenceAndProse(t *testing.T) {
|
|
obj, err := ParseJSONObject("Here you go:\n```\n{\"name\":\"X\",\"description\":\"Y\"}\n```")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if obj["name"] != "X" {
|
|
t.Fatalf("%v", obj)
|
|
}
|
|
}
|
|
|
|
func TestParseJSONObject_arrayOfObjects(t *testing.T) {
|
|
obj, err := ParseJSONObject(`[{"name":"N","description":"D"},{"name":"Other"}]`)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if obj["name"] != "N" || obj["description"] != "D" {
|
|
t.Fatalf("%v", obj)
|
|
}
|
|
}
|
|
|
|
func TestParseJSONObject_arrayFirstObjectFallback(t *testing.T) {
|
|
obj, err := ParseJSONObject(`[{"foo":1},{"name":"N"}]`)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if obj["name"] != "N" {
|
|
t.Fatalf("%v", obj)
|
|
}
|
|
}
|
|
|
|
func TestCompactAttrs_priorityAndCap(t *testing.T) {
|
|
attrs := map[string]any{
|
|
"zzz": "late", "brand": "Acme", "color": "Red",
|
|
"a": "1", "b": "2", "c": "3", "d": "4", "e": "5", "f": "6", "g": "7", "h": "8",
|
|
}
|
|
got := CompactAttrs(attrs, 5)
|
|
if len(got) > 5 {
|
|
t.Fatalf("len=%d", len(got))
|
|
}
|
|
if got["brand"] != "Acme" {
|
|
t.Fatalf("brand missing: %v", got)
|
|
}
|
|
}
|
|
|
|
func TestCompleteJSON_retriesOnBadJSON(t *testing.T) {
|
|
calls := 0
|
|
c := stubCompleter{fn: func(_, _ string) (Completion, error) {
|
|
calls++
|
|
if calls == 1 {
|
|
return Completion{Text: "not json", TotalTokens: 2}, nil
|
|
}
|
|
return Completion{Text: `{"name":"N","description":"D"}`, TotalTokens: 3}, nil
|
|
}}
|
|
comp, obj, err := CompleteJSON(context.Background(), c, "sys", "user", CompleteOptions{MaxTokens: 50})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if calls != 2 {
|
|
t.Fatalf("calls=%d", calls)
|
|
}
|
|
if obj["name"] != "N" {
|
|
t.Fatalf("%v", obj)
|
|
}
|
|
if comp.TotalTokens != 5 {
|
|
t.Fatalf("tokens=%d", comp.TotalTokens)
|
|
}
|
|
}
|
|
|
|
func TestCompleteJSON_returnsRetryError(t *testing.T) {
|
|
calls := 0
|
|
retryErr := context.DeadlineExceeded
|
|
c := stubCompleter{fn: func(_, _ string) (Completion, error) {
|
|
calls++
|
|
if calls == 1 {
|
|
return Completion{Text: "not json", TotalTokens: 2}, nil
|
|
}
|
|
return Completion{}, retryErr
|
|
}}
|
|
|
|
comp, obj, err := CompleteJSON(context.Background(), c, "sys", "user", CompleteOptions{MaxTokens: 50})
|
|
if err != retryErr {
|
|
t.Fatalf("err=%v want=%v", err, retryErr)
|
|
}
|
|
if obj != nil {
|
|
t.Fatalf("obj=%v", obj)
|
|
}
|
|
if comp.Text != "not json" {
|
|
t.Fatalf("comp=%+v", comp)
|
|
}
|
|
if calls != 2 {
|
|
t.Fatalf("calls=%d", calls)
|
|
}
|
|
}
|
|
|
|
func TestProductEnhanceUser_truncates(t *testing.T) {
|
|
long := strings.Repeat("x", MaxProductDescRunes+500)
|
|
u := ProductEnhanceUser("Cat", "Name", long, map[string]any{"brand": "B"})
|
|
// Category/Name/Attrs overhead + MaxProductDescRunes of description.
|
|
if len([]rune(u)) > MaxProductDescRunes+200 {
|
|
t.Fatalf("user too long: %d", len([]rune(u)))
|
|
}
|
|
descIdx := strings.Index(u, "Description: ")
|
|
if descIdx < 0 {
|
|
t.Fatalf("missing Description: in %q", u)
|
|
}
|
|
descPart := u[descIdx+len("Description: "):]
|
|
if i := strings.Index(descPart, "\n"); i >= 0 {
|
|
descPart = descPart[:i]
|
|
}
|
|
if len([]rune(descPart)) > MaxProductDescRunes {
|
|
t.Fatalf("description runes=%d want <= %d", len([]rune(descPart)), MaxProductDescRunes)
|
|
}
|
|
if !strings.Contains(u, "brand") {
|
|
t.Fatalf("%s", u)
|
|
}
|
|
}
|