Assistant/adapters/models/advisory.go

package models

import (
	"errors"
	"fmt"
	"unicode/utf8"

	"git.cin.su/m/Assistant/contracts/cognitive"
)

const (
	MaxAdvisoryMessageBytes  = 16 * 1024
	MaxAdvisoryExcerpts      = 8
	MaxAdvisoryExcerptBytes  = 8 * 1024
	MaxAdvisoryContextBytes  = 48 * 1024
	MaxAdvisoryResponseBytes = 16 * 1024
)

var ErrInvalidAdvisory = errors.New("invalid creator advisory")

func ValidateAdvisoryRequest(request cognitive.AdvisoryRequest) error {
	if request.Version != cognitive.CreatorAdvisoryV1 || request.RequestID == "" || request.ActorID == "" || request.PrincipalID == "" || request.ScopeID == "" || request.ProviderBindingID == "" || request.ProviderBindingRevision <= 0 {
		return fmt.Errorf("%w: missing or unsupported identity/binding fields", ErrInvalidAdvisory)
	}
	if request.Message == "" || !utf8.ValidString(request.Message) || len(request.Message) > MaxAdvisoryMessageBytes {
		return fmt.Errorf("%w: message", ErrInvalidAdvisory)
	}
	if request.Locale == "" || len(request.Locale) > 32 || !utf8.ValidString(request.Locale) || request.MaxOutputBytes <= 0 || request.MaxOutputBytes > MaxAdvisoryResponseBytes {
		return fmt.Errorf("%w: locale or response limit", ErrInvalidAdvisory)
	}
	if len(request.Context) > MaxAdvisoryExcerpts {
		return fmt.Errorf("%w: too many context excerpts", ErrInvalidAdvisory)
	}
	seen := make(map[cognitive.ContextExcerptID]struct{}, len(request.Context))
	total := 0
	for _, excerpt := range request.Context {
		if excerpt.ID == "" || excerpt.SourceKind == "" || excerpt.SourceID == "" || excerpt.Revision <= 0 || excerpt.Text == "" || !utf8.ValidString(excerpt.Text) || len(excerpt.Text) > MaxAdvisoryExcerptBytes {
			return fmt.Errorf("%w: context excerpt", ErrInvalidAdvisory)
		}
		if excerpt.Sensitivity != cognitive.SensitivityPublic && excerpt.Sensitivity != cognitive.SensitivityPrivateAllowed {
			return fmt.Errorf("%w: context sensitivity", ErrInvalidAdvisory)
		}
		if _, exists := seen[excerpt.ID]; exists {
			return fmt.Errorf("%w: duplicate context excerpt", ErrInvalidAdvisory)
		}
		seen[excerpt.ID] = struct{}{}
		total += len(excerpt.Text)
	}
	if total > MaxAdvisoryContextBytes {
		return fmt.Errorf("%w: total context", ErrInvalidAdvisory)
	}
	return nil
}

func ValidateAdvisoryResponse(request cognitive.AdvisoryRequest, response cognitive.AdvisoryResponse) error {
	if response.Version != cognitive.CreatorAdvisoryV1 || response.RequestID != request.RequestID || response.ProviderBindingID != request.ProviderBindingID || response.ProviderBindingRevision != request.ProviderBindingRevision {
		return fmt.Errorf("%w: response identity", ErrInvalidAdvisory)
	}
	if !validFinishReason(response.FinishReason) || !utf8.ValidString(response.AdvisoryText) || len(response.AdvisoryText) > request.MaxOutputBytes {
		return fmt.Errorf("%w: response content", ErrInvalidAdvisory)
	}
	if response.FinishReason == cognitive.FinishCompleted && response.AdvisoryText == "" {
		return fmt.Errorf("%w: completed response is empty", ErrInvalidAdvisory)
	}
	if response.FinishReason != cognitive.FinishCompleted && response.AdvisoryText != "" {
		return fmt.Errorf("%w: failure response contains advisory text", ErrInvalidAdvisory)
	}
	if response.FinishReason != cognitive.FinishCompleted && len(response.ReferencedExcerptIDs) != 0 {
		return fmt.Errorf("%w: failure response contains citations", ErrInvalidAdvisory)
	}
	allowed := make(map[cognitive.ContextExcerptID]struct{}, len(request.Context))
	for _, excerpt := range request.Context {
		allowed[excerpt.ID] = struct{}{}
	}
	seen := make(map[cognitive.ContextExcerptID]struct{}, len(response.ReferencedExcerptIDs))
	for _, id := range response.ReferencedExcerptIDs {
		if _, ok := allowed[id]; !ok {
			return fmt.Errorf("%w: unknown citation", ErrInvalidAdvisory)
		}
		if _, ok := seen[id]; ok {
			return fmt.Errorf("%w: duplicate citation", ErrInvalidAdvisory)
		}
		seen[id] = struct{}{}
	}
	if response.Usage != nil && (response.Usage.InputUnits < 0 || response.Usage.OutputUnits < 0) {
		return fmt.Errorf("%w: negative usage", ErrInvalidAdvisory)
	}
	return nil
}

func DeterministicAdvisoryUnavailable(request cognitive.AdvisoryRequest) (cognitive.AdvisoryResponse, error) {
	if err := ValidateAdvisoryRequest(request); err != nil {
		return cognitive.AdvisoryResponse{}, err
	}
	return cognitive.AdvisoryResponse{Version: cognitive.CreatorAdvisoryV1, RequestID: request.RequestID, ProviderBindingID: request.ProviderBindingID, ProviderBindingRevision: request.ProviderBindingRevision, FinishReason: cognitive.FinishProviderUnavailable}, nil
}

func validFinishReason(reason cognitive.FinishReason) bool {
	switch reason {
	case cognitive.FinishCompleted, cognitive.FinishLengthLimit, cognitive.FinishContentRejected,
		cognitive.FinishProviderUnavailable, cognitive.FinishModelUnavailable, cognitive.FinishRateLimited,
		cognitive.FinishAuthenticationFailed, cognitive.FinishBillingOrBudgetBlocked,
		cognitive.FinishContextTooLarge, cognitive.FinishUnsupportedCapability, cognitive.FinishTimeout,
		cognitive.FinishInvalidProviderResponse, cognitive.FinishInvalidStructuredOutput, cognitive.FinishDegraded:
		return true
	default:
		return false
	}
}