rony-llm-agent/pkg/llm/mock/mock.go

package mock

import (
	"context"
	"iter"
	"strings"

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
)

// MockLLMClient is a deterministic implementation of llm.LLMClient for testing.
type MockLLMClient struct {
	GenerateFunc     func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error)
	StreamFunc       func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error]
	NameFunc         func() string
	CapabilitiesFunc func() llm.ProviderCapabilities
}

func (m *MockLLMClient) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
	if m.GenerateFunc != nil {
		return m.GenerateFunc(ctx, req)
	}
	return llm.CompletionResponse{
		Content:    "default response",
		StopReason: llm.StopReasonEndTurn,
	}, nil
}

func (m *MockLLMClient) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
	if m.StreamFunc != nil {
		return m.StreamFunc(ctx, req)
	}
	return func(yield func(llm.StreamChunk, error) bool) {
		yield(llm.StreamChunk{
			Delta: "default response",
		}, nil)
	}
}

func (m *MockLLMClient) Name() string {
	if m.NameFunc != nil {
		return m.NameFunc()
	}
	return "mock"
}

func (m *MockLLMClient) Capabilities() llm.ProviderCapabilities {
	if m.CapabilitiesFunc != nil {
		return m.CapabilitiesFunc()
	}
	return llm.ProviderCapabilities{
		SupportsTools:    true,
		SupportsVision:   false,
		SupportsJSON:     true,
		MaxContextWindow: 128000,
	}
}

// New returns a MockLLMClient with default behaviors.
func New() *MockLLMClient {
	return &MockLLMClient{}
}

// NewWithGenerate returns a MockLLMClient that returns the given response.
func NewWithGenerate(resp llm.CompletionResponse) *MockLLMClient {
	return &MockLLMClient{
		GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
			return resp, nil
		},
	}
}

// NewWithStream returns a MockLLMClient that streams the given chunks.
func NewWithStream(chunks []llm.StreamChunk) *MockLLMClient {
	return &MockLLMClient{
		StreamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
			return func(yield func(llm.StreamChunk, error) bool) {
				for _, c := range chunks {
					if !yield(c, nil) {
						return
					}
				}
			}
		},
	}
}

// NewWithMatch returns a MockLLMClient that matches input messages against patterns.
func NewWithMatch(responses []MatchResponse) *MockLLMClient {
	return &MockLLMClient{
		GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
			for _, r := range responses {
				if r.Match == "*" || matchesAny(r.Match, req.Messages) {
					return llm.CompletionResponse{
						Content:    r.Response,
						StopReason: llm.StopReasonEndTurn,
					}, nil
				}
			}
			return llm.CompletionResponse{
				Content:    "no match",
				StopReason: llm.StopReasonEndTurn,
			}, nil
		},
	}
}

// MatchResponse pairs a pattern with a response string.
type MatchResponse struct {
	Match    string // exact string or "*" for wildcard
	Response string
}

func matchesAny(pattern string, msgs []llm.Message) bool {
	for _, m := range msgs {
		if m.Role == llm.RoleUser && strings.Contains(m.Content, pattern) {
			return true
		}
	}
	return false
}