74 lines
2.2 KiB
Go
74 lines
2.2 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/ConfigKit"
|
|
"gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
type fakeLLMClient struct {
|
|
response *promptkit.GenerateResponse
|
|
responses []*promptkit.GenerateResponse
|
|
err error
|
|
requests []promptkit.GenerateRequest
|
|
}
|
|
|
|
func (f *fakeLLMClient) Generate(_ context.Context, request promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
f.requests = append(f.requests, request)
|
|
if f.err != nil {
|
|
return nil, f.err
|
|
}
|
|
if len(f.responses) > 0 {
|
|
index := len(f.requests) - 1
|
|
if index >= len(f.responses) {
|
|
return nil, fmt.Errorf("no response configured for generation %d", index+1)
|
|
}
|
|
return f.responses[index], nil
|
|
}
|
|
return f.response, nil
|
|
}
|
|
|
|
type countingLLMClient struct{ calls *atomic.Int64 }
|
|
|
|
func (client countingLLMClient) Generate(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
if client.calls != nil {
|
|
client.calls.Add(1)
|
|
}
|
|
return &promptkit.GenerateResponse{Content: "ok"}, nil
|
|
}
|
|
|
|
type roundTripFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { return f(request) }
|
|
|
|
// newSimpleEngine exercises the public constructor with a complete keyed
|
|
// document. Tests that need layering or resources select their own inputs.
|
|
func newSimpleEngine(t *testing.T, content string, options ...promptkit.Option) *promptkit.Engine {
|
|
t.Helper()
|
|
engine, err := promptkit.NewEngine(simpleConfig(content), options...)
|
|
if err != nil {
|
|
t.Fatalf("construct consumer test engine: %v", err)
|
|
}
|
|
return engine
|
|
}
|
|
|
|
func simpleConfig(content string) promptkit.Config {
|
|
return promptkit.Config{Sources: promptkit.ConfigurationSources{
|
|
Main: &configkit.Input{Source: configkit.Mapping("consumer test", map[string]any{
|
|
"profiles": map[string]any{"profile": map[string]any{
|
|
"endpoint": "http://example.test/v1", "model": "model",
|
|
}},
|
|
"prompts": map[string]any{"prompt": map[string]any{"1": map[string]any{
|
|
"default_profile": "profile",
|
|
"messages": []any{map[string]any{"role": "user", "content": content}},
|
|
"output": map[string]any{"format": "text", "validation_mode": "none"},
|
|
}}},
|
|
})},
|
|
}}
|
|
}
|