majordomo/provider/fake/fake.go

// Package fake provides an in-memory llm.Provider for hermetic tests.
//
// Why: the resolver, env-DSN loader, chain executor, health tracker, agent
// loop, and skill composition must all be testable with no live API calls.
// The fake provider scripts responses and errors per model id, records every
// request it receives, and supports tools, structured output, and streaming
// well enough to drive those layers deterministically.
package fake

import (
	"context"
	"fmt"
	"io"
	"sync"

	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

// Step is one scripted outcome: either a response or an error.
type Step struct {
	Response *llm.Response
	Err      error
}

// Reply scripts a successful text response.
func Reply(text string) Step {
	return Step{Response: &llm.Response{
		Parts:        []llm.Part{llm.Text(text)},
		FinishReason: llm.FinishStop,
		Usage:        llm.Usage{InputTokens: 1, OutputTokens: 1},
	}}
}

// ReplyWith scripts an arbitrary successful response.
func ReplyWith(resp llm.Response) Step { return Step{Response: &resp} }

// Fail scripts an error outcome.
func Fail(err error) Step { return Step{Err: err} }

// Call records one request received by the fake, with the model id it was
// addressed to.
type Call struct {
	ModelID string
	Request llm.Request
}

// Provider is a scriptable in-memory llm.Provider.
//
// Outcomes are enqueued per model id with Enqueue. A model whose queue is
// empty falls back to the provider default response (a fixed text reply).
// All methods are safe for concurrent use.
type Provider struct {
	name string

	mu        sync.Mutex
	caps      llm.Capabilities
	modelCaps map[string]llm.Capabilities
	queues    map[string][]Step
	calls     []Call
	defaultFn func(modelID string, req llm.Request) Step
}

// Option configures the fake provider.
type Option func(*Provider)

// WithCapabilities sets the provider-default capabilities.
func WithCapabilities(caps llm.Capabilities) Option {
	return func(p *Provider) { p.caps = caps }
}

// WithModelCapabilities overrides capabilities for one model id.
func WithModelCapabilities(modelID string, caps llm.Capabilities) Option {
	return func(p *Provider) { p.modelCaps[modelID] = caps }
}

// WithDefault sets the outcome used when a model's queue is empty.
func WithDefault(fn func(modelID string, req llm.Request) Step) Option {
	return func(p *Provider) { p.defaultFn = fn }
}

// New creates a fake provider with the given registry name.
func New(name string, opts ...Option) *Provider {
	p := &Provider{
		name:      name,
		modelCaps: make(map[string]llm.Capabilities),
		queues:    make(map[string][]Step),
		caps: llm.Capabilities{
			SupportsTools:      true,
			SupportsStructured: true,
			SupportsStreaming:  true,
			MaxImagesPerReq:    4,
		},
		defaultFn: func(modelID string, _ llm.Request) Step {
			return Reply(fmt.Sprintf("fake response from %s", modelID))
		},
	}
	for _, opt := range opts {
		opt(p)
	}
	return p
}

// Name implements llm.Provider.
func (p *Provider) Name() string { return p.name }

// Model implements llm.Provider. Any id is accepted.
func (p *Provider) Model(id string, opts ...llm.ModelOption) (llm.Model, error) {
	cfg := llm.ApplyModelOptions(opts)
	return &model{provider: p, id: id, cfg: cfg}, nil
}

// Enqueue appends scripted outcomes for a model id.
func (p *Provider) Enqueue(modelID string, steps ...Step) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.queues[modelID] = append(p.queues[modelID], steps...)
}

// Calls returns a copy of every request received so far.
func (p *Provider) Calls() []Call {
	p.mu.Lock()
	defer p.mu.Unlock()
	out := make([]Call, len(p.calls))
	copy(out, p.calls)
	return out
}

// CallCount returns the number of requests received for one model id.
func (p *Provider) CallCount(modelID string) int {
	p.mu.Lock()
	defer p.mu.Unlock()
	n := 0
	for _, c := range p.calls {
		if c.ModelID == modelID {
			n++
		}
	}
	return n
}

// next records the call and pops the next scripted outcome.
func (p *Provider) next(modelID string, req llm.Request) Step {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.calls = append(p.calls, Call{ModelID: modelID, Request: req})
	q := p.queues[modelID]
	if len(q) == 0 {
		return p.defaultFn(modelID, req)
	}
	step := q[0]
	p.queues[modelID] = q[1:]
	return step
}

func (p *Provider) capsFor(modelID string, cfg llm.ModelConfig) llm.Capabilities {
	if cfg.Capabilities != nil {
		return *cfg.Capabilities
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if caps, ok := p.modelCaps[modelID]; ok {
		return caps
	}
	return p.caps
}

type model struct {
	provider *Provider
	id       string
	cfg      llm.ModelConfig
}

func (m *model) Generate(ctx context.Context, req llm.Request, opts ...llm.Option) (*llm.Response, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	req = req.Apply(opts...)
	step := m.provider.next(m.id, req)
	if step.Err != nil {
		return nil, step.Err
	}
	resp := *step.Response
	if resp.Model == "" {
		resp.Model = m.provider.name + "/" + m.id
	}
	return &resp, nil
}

func (m *model) Stream(ctx context.Context, req llm.Request, opts ...llm.Option) (llm.Stream, error) {
	resp, err := m.Generate(ctx, req, opts...)
	if err != nil {
		return nil, err
	}
	// Deliver the response as a small sequence of events: one text delta per
	// part, one event per tool call, then the final response.
	var events []llm.StreamEvent
	for _, part := range resp.Parts {
		if t, ok := part.(llm.TextPart); ok {
			events = append(events, llm.StreamEvent{TextDelta: t.Text})
		}
	}
	for i := range resp.ToolCalls {
		events = append(events, llm.StreamEvent{ToolCall: &resp.ToolCalls[i]})
	}
	events = append(events, llm.StreamEvent{Response: resp})
	return &stream{events: events}, nil
}

func (m *model) Capabilities() llm.Capabilities {
	return m.provider.capsFor(m.id, m.cfg)
}

type stream struct {
	mu     sync.Mutex
	events []llm.StreamEvent
	pos    int
}

func (s *stream) Next() (llm.StreamEvent, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.pos >= len(s.events) {
		return llm.StreamEvent{}, io.EOF
	}
	ev := s.events[s.pos]
	s.pos++
	return ev, nil
}

func (s *stream) Close() error { return nil }