promptkit/internal/usecase/smoke_inference_case.go

package usecase

import (
	"context"
	"embed"
	"fmt"
	"strings"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"gitea.maximumdirect.net/eric/promptkit/internal/validate"
)

const (
	smokeTextExpectedContent       = "PROMPTKIT_OK"
	smokeTextPrompt                = "Reply with exactly the token PROMPTKIT_OK and no other text."
	smokeStructuredExpectedContent = `{"status":"ok"}`
	smokeStructuredPromptPrefix    = "Return exactly this JSON object: "
	smokeStructuredSchemaPath      = "smoke_schema.json"
)

//go:embed smoke_schema.json
var smokeSchemaFS embed.FS

type smokeInferenceCase struct {
	kind             domain.SmokeTestKind
	request          domain.GenerateRequest
	expectedContent  string
	jsonValidator    validate.Validator
	schemaValidation validate.PreparedValidation
}

func prepareSmokeInferenceCase(
	ctx context.Context,
	kind domain.SmokeTestKind,
	target domain.SmokeTestTarget,
	transientRetries uint64,
) (*smokeInferenceCase, error) {
	switch kind {
	case domain.SmokeTestTextInference:
		return &smokeInferenceCase{
			kind:            kind,
			expectedContent: smokeTextExpectedContent,
			request: newGenerationRequest(
				domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: smokeTextPrompt}}},
				"",
				target.EffectiveSettings,
				target.SettingsPresence,
				nil,
				transientRetries,
				1,
			),
		}, nil
	case domain.SmokeTestStructuredOutput:
		validator := validate.NewFSValidator(smokeSchemaFS, ".")
		preparer, ok := validator.(validate.ValidationPreparer)
		if !ok {
			return nil, fmt.Errorf("%w: private validator does not support prepared validation", ErrValidation)
		}
		plan, err := preparer.PrepareValidation(ctx, domain.OutputContract{
			Format:         domain.FormatJSON,
			ValidationMode: domain.ValidationJSONSchema,
			SchemaPath:     smokeStructuredSchemaPath,
		})
		if err != nil {
			return nil, fmt.Errorf("%w: prepare private structured-output schema: %w", ErrValidation, err)
		}
		structuredOutput := structuredOutputSpec(&domain.PromptDefinition{ID: "promptkit_smoke"}, plan.SchemaDocument())
		return &smokeInferenceCase{
			kind:             kind,
			expectedContent:  smokeStructuredExpectedContent,
			jsonValidator:    validate.NewStandardValidator(""),
			schemaValidation: plan,
			request: newGenerationRequest(
				domain.RenderedPrompt{Messages: []domain.RenderedMessage{{Role: "user", Content: smokeStructuredPromptPrefix + smokeStructuredExpectedContent}}},
				"",
				target.EffectiveSettings,
				target.SettingsPresence,
				structuredOutput,
				transientRetries,
				1,
			),
		}, nil
	default:
		return nil, fmt.Errorf("%w: smoke inference case does not support kind %q", ErrInvalidRequest, kind)
	}
}

func (c *smokeInferenceCase) evaluate(ctx context.Context, response *domain.GenerateResponse) (domain.SmokeTestFacts, error) {
	var facts domain.SmokeTestFacts
	if c == nil || response == nil {
		return facts, nil
	}

	switch c.kind {
	case domain.SmokeTestTextInference:
		content := strings.TrimSpace(response.Content)
		textReturned := content != ""
		matched := content == c.expectedContent
		facts.TextReturned = boolPointer(textReturned)
		facts.ExpectedTextMatched = boolPointer(matched)
		return facts, nil
	case domain.SmokeTestStructuredOutput:
		artifact := domain.Artifact{Body: []byte(response.Content)}
		jsonResult, err := c.jsonValidator.Validate(ctx, &artifact, domain.OutputContract{
			Format:         domain.FormatJSON,
			ValidationMode: domain.ValidationJSON,
		})
		if err != nil {
			return facts, fmt.Errorf("%w: evaluate private smoke-test JSON: %w", ErrValidation, err)
		}
		validJSON := jsonResult.IsValid
		facts.ValidJSON = boolPointer(validJSON)
		if !validJSON {
			return facts, nil
		}
		schemaResult, err := c.schemaValidation.Validate(ctx, &artifact)
		if err != nil {
			return facts, fmt.Errorf("%w: evaluate private smoke-test schema: %w", ErrValidation, err)
		}
		facts.SchemaMatched = boolPointer(schemaResult.IsValid)
		return facts, nil
	default:
		return facts, fmt.Errorf("%w: smoke inference case does not support kind %q", ErrInvalidRequest, c.kind)
	}
}

func (c *smokeInferenceCase) structuredOutputRequested() bool {
	return c != nil && c.request.StructuredOutput != nil
}