promptkit/internal/prompt/requirements.go

package prompt

import (
	"context"
	"sort"
	"text/template"
	"text/template/parse"

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

type requirementReachability uint8

const (
	requirementDefinite requirementReachability = iota
	requirementPossible
)

type dotIdentity uint8

const (
	dotRoot dotIdentity = iota
	dotNonRoot
)

type inputEvidence struct {
	name        string
	kind        domain.RequirementKind
	declared    bool
	contentType *string
	description *string
	explanation string
}

type variableEvidence struct {
	name        string
	kind        domain.RequirementKind
	explanation string
}

type requirementIssueKey struct {
	code         domain.RequirementIssueCode
	inputName    string
	variableName string
}

type requirementWalker struct {
	ctx          context.Context
	vars         map[string]string
	inputs       map[string]*inputEvidence
	variables    map[string]*variableEvidence
	issues       map[requirementIssueKey]domain.RequirementIssue
	complete     bool
	rootKnown    bool
	dependencies int
	templateSet  *template.Template
	visitedTrees map[*parse.Tree]bool
}

// AnalyzeRequirements collects the prompt dependencies recognizable from its
// captured syntax. Branch bodies remain conditional until their conditions are
// classified; unsupported dependency-bearing syntax makes the result partial.
func AnalyzeRequirements(
	ctx context.Context,
	definition *domain.PromptDefinition,
	snapshot RequirementSnapshot,
	vars map[string]string,
	sessionBypassed bool,
) (domain.RequirementInspection, error) {
	if err := ctx.Err(); err != nil {
		return domain.RequirementInspection{}, err
	}
	result := domain.RequirementInspection{Complete: true}
	if definition != nil {
		result.PromptID = definition.ID
		result.PromptVersion = definition.Version
	}
	walker := &requirementWalker{
		ctx:       ctx,
		vars:      vars,
		inputs:    make(map[string]*inputEvidence, len(snapshot.Inputs)),
		variables: make(map[string]*variableEvidence),
		issues:    make(map[requirementIssueKey]domain.RequirementIssue),
		complete:  true,
		rootKnown: true,
	}
	for _, input := range snapshot.Inputs {
		if err := ctx.Err(); err != nil {
			return domain.RequirementInspection{}, err
		}
		kind := domain.RequirementOptional
		if input.Required {
			kind = domain.RequirementRequired
		}
		walker.inputs[input.Name] = &inputEvidence{
			name:        input.Name,
			kind:        kind,
			declared:    true,
			contentType: copyString(input.ContentType),
			description: copyString(input.Description),
		}
		if kind == domain.RequirementOptional {
			walker.inputs[input.Name].explanation = "Declared as an optional input."
		}
	}
	for _, message := range snapshot.Templates.Messages {
		if err := walker.walkTemplate(message); err != nil {
			return domain.RequirementInspection{}, err
		}
	}
	if !sessionBypassed && snapshot.Templates.Session != nil {
		if err := walker.walkTemplate(snapshot.Templates.Session); err != nil {
			return domain.RequirementInspection{}, err
		}
	}
	if err := ctx.Err(); err != nil {
		return domain.RequirementInspection{}, err
	}
	return walker.result(result), nil
}

func (walker *requirementWalker) walkTemplate(compiled *template.Template) error {
	if err := walker.ctx.Err(); err != nil {
		return err
	}
	if compiled == nil || compiled.Tree == nil || compiled.Tree.Root == nil {
		walker.addIssue("", "", "A captured template has no parsed root.")
		return nil
	}
	previousSet := walker.templateSet
	previousVisited := walker.visitedTrees
	walker.templateSet = compiled
	walker.visitedTrees = map[*parse.Tree]bool{compiled.Tree: true}
	walker.rootKnown = true
	err := walker.visit(compiled.Tree.Root, requirementDefinite, dotRoot)
	walker.templateSet = previousSet
	walker.visitedTrees = previousVisited
	return err
}

func (walker *requirementWalker) visit(node parse.Node, reach requirementReachability, dot dotIdentity) error {
	if err := walker.ctx.Err(); err != nil {
		return err
	}
	if node == nil {
		return nil
	}
	switch node := node.(type) {
	case *parse.ListNode:
		if node == nil {
			return nil
		}
		for _, child := range node.Nodes {
			if err := walker.visit(child, reach, dot); err != nil {
				return err
			}
		}
	case *parse.ActionNode:
		return walker.visit(node.Pipe, reach, dot)
	case *parse.IfNode:
		if err := walker.visit(node.Pipe, reach, dot); err != nil {
			return err
		}
		if err := walker.visit(node.List, requirementPossible, dot); err != nil {
			return err
		}
		return walker.visit(node.ElseList, requirementPossible, dot)
	case *parse.WithNode:
		if err := walker.visit(node.Pipe, reach, dot); err != nil {
			return err
		}
		if err := walker.visit(node.List, requirementPossible, dotNonRoot); err != nil {
			return err
		}
		return walker.visit(node.ElseList, requirementPossible, dot)
	case *parse.RangeNode:
		if err := walker.visit(node.Pipe, reach, dot); err != nil {
			return err
		}
		before := walker.dependencies
		if err := walker.visit(node.List, requirementPossible, dotNonRoot); err != nil {
			return err
		}
		if err := walker.visit(node.ElseList, requirementPossible, dot); err != nil {
			return err
		}
		if walker.dependencies != before {
			walker.addIssue("", "", "Dependencies inside a range are only conservatively enumerated.")
		}
	case *parse.TemplateNode:
		if err := walker.visit(node.Pipe, requirementPossible, dot); err != nil {
			return err
		}
		called := walker.templateSet.Lookup(node.Name)
		if called == nil || called.Tree == nil || called.Tree.Root == nil {
			walker.addIssue("", "", "An invoked named template is unavailable for syntax inspection.")
			return nil
		}
		if walker.visitedTrees[called.Tree] {
			walker.addIssue("", "", "An invoked named template is scanned conservatively.")
			return nil
		}
		walker.visitedTrees[called.Tree] = true
		rootKnown := walker.rootKnown
		walker.rootKnown = false
		err := walker.visit(called.Tree.Root, requirementPossible, dotNonRoot)
		walker.rootKnown = rootKnown
		if err != nil {
			return err
		}
		walker.addIssue("", "", "An invoked named template is scanned conservatively.")
	case *parse.PipeNode:
		return walker.visitPipe(node, reach, dot)
	case *parse.CommandNode:
		return walker.visitCommand(node, reach, dot)
	case *parse.FieldNode:
		if dot == dotRoot && len(node.Ident) > 0 {
			walker.recordVariable(node.Ident[0], reach)
			if len(node.Ident) > 1 {
				walker.addIssue("", node.Ident[0], "A chained field access is outside the direct-key scope.")
			}
		}
	case *parse.VariableNode:
		walker.visitVariable(node, reach)
	case *parse.ChainNode:
		before := walker.dependencies
		if err := walker.visit(node.Node, requirementPossible, dot); err != nil {
			return err
		}
		if walker.dependencies != before {
			walker.addIssue("", "", "A chained expression is not interpreted.")
		}
	case *parse.BreakNode, *parse.ContinueNode:
		walker.addIssue("", "", "Loop control is outside the supported syntax scope.")
	}
	return nil
}

func (walker *requirementWalker) visitPipe(pipe *parse.PipeNode, reach requirementReachability, dot dotIdentity) error {
	if pipe == nil {
		return nil
	}
	if len(pipe.Cmds) == 1 && len(pipe.Decl) == 0 && !pipe.IsAssign {
		return walker.visitCommand(pipe.Cmds[0], reach, dot)
	}
	before := walker.dependencies
	for _, command := range pipe.Cmds {
		if err := walker.visitCommand(command, requirementPossible, dot); err != nil {
			return err
		}
	}
	if walker.dependencies != before {
		walker.addIssue("", "", "A pipeline or variable assignment is not interpreted.")
	}
	for _, declaration := range pipe.Decl {
		if len(declaration.Ident) > 0 && declaration.Ident[0] == "$" {
			walker.rootKnown = false
			walker.addIssue("", "", "The root template binding is reassigned.")
		}
	}
	return nil
}

func (walker *requirementWalker) visitCommand(command *parse.CommandNode, reach requirementReachability, dot dotIdentity) error {
	if command == nil || len(command.Args) == 0 {
		return nil
	}
	identifier, isIdentifier := command.Args[0].(*parse.IdentifierNode)
	if isIdentifier && identifier.Ident == "input" {
		if len(command.Args) == 2 {
			if name, ok := command.Args[1].(*parse.StringNode); ok {
				walker.recordInput(name.Text, reach)
				return nil
			}
		}
		for _, argument := range command.Args[1:] {
			if err := walker.visit(argument, requirementPossible, dot); err != nil {
				return err
			}
		}
		walker.addIssue("", "", "An input name is not a direct string literal.")
		return nil
	}
	if isIdentifier && identifier.Ident == "index" && len(command.Args) == 3 {
		if _, ok := command.Args[2].(*parse.StringNode); ok && isRootMap(command.Args[1], dot, walker.rootKnown) {
			key := command.Args[2].(*parse.StringNode).Text
			walker.recordOptionalVariable(key, reach)
			return nil
		}
	}
	if !isIdentifier && len(command.Args) == 1 {
		return walker.visit(command.Args[0], reach, dot)
	}
	before := walker.dependencies
	for index, argument := range command.Args {
		if isIdentifier && index == 0 {
			continue
		}
		if err := walker.visit(argument, requirementPossible, dot); err != nil {
			return err
		}
	}
	if walker.dependencies != before {
		walker.addIssue("", "", "A function call or computed expression is not interpreted.")
	}
	return nil
}

func (walker *requirementWalker) visitVariable(variable *parse.VariableNode, reach requirementReachability) {
	if len(variable.Ident) < 2 || variable.Ident[0] != "$" {
		return
	}
	if !walker.rootKnown {
		walker.addIssue("", "", "The root template binding is not known.")
		return
	}
	walker.recordVariable(variable.Ident[1], reach)
	if len(variable.Ident) > 2 {
		walker.addIssue("", variable.Ident[1], "A chained field access is outside the direct-key scope.")
	}
}

func isRootMap(node parse.Node, dot dotIdentity, rootKnown bool) bool {
	switch node := node.(type) {
	case *parse.DotNode:
		return dot == dotRoot
	case *parse.VariableNode:
		return rootKnown && len(node.Ident) == 1 && node.Ident[0] == "$"
	default:
		return false
	}
}

func (walker *requirementWalker) recordInput(name string, reach requirementReachability) {
	evidence := walker.inputs[name]
	if evidence == nil {
		evidence = &inputEvidence{
			name: name, kind: domain.RequirementOptional,
			explanation: "Declared as an optional input.",
		}
		walker.inputs[name] = evidence
	}
	walker.promoteInput(evidence, kindForReach(reach))
	if evidence.kind == domain.RequirementRequired {
		evidence.explanation = ""
	} else if reach == requirementPossible {
		evidence.explanation = "This reference is not known to be reached on every execution path."
		walker.addIssue(name, "", "A dependency is conservatively classified.")
	}
	walker.dependencies++
}

func (walker *requirementWalker) recordVariable(name string, reach requirementReachability) {
	evidence := walker.variables[name]
	if evidence == nil {
		evidence = &variableEvidence{name: name, kind: domain.RequirementOptional}
		walker.variables[name] = evidence
	}
	walker.promoteVariable(evidence, kindForReach(reach))
	if evidence.kind == domain.RequirementRequired {
		evidence.explanation = ""
	} else if reach == requirementPossible {
		evidence.explanation = "This reference is not known to be reached on every execution path."
		walker.addIssue("", name, "A dependency is conservatively classified.")
	}
	walker.dependencies++
}

func (walker *requirementWalker) recordOptionalVariable(name string, _ requirementReachability) {
	evidence := walker.variables[name]
	if evidence == nil {
		evidence = &variableEvidence{name: name, kind: domain.RequirementOptional}
		walker.variables[name] = evidence
	}
	if evidence.kind == domain.RequirementOptional {
		evidence.explanation = "This key is read through an optional map lookup."
	}
	walker.dependencies++
}

func (walker *requirementWalker) promoteInput(evidence *inputEvidence, kind domain.RequirementKind) {
	if requirementRank(kind) > requirementRank(evidence.kind) {
		evidence.kind = kind
	}
	if evidence.kind == domain.RequirementRequired {
		evidence.explanation = ""
	}
}

func (walker *requirementWalker) promoteVariable(evidence *variableEvidence, kind domain.RequirementKind) {
	if requirementRank(kind) > requirementRank(evidence.kind) {
		evidence.kind = kind
	}
	if evidence.kind == domain.RequirementRequired {
		evidence.explanation = ""
	}
}

func kindForReach(reach requirementReachability) domain.RequirementKind {
	if reach == requirementPossible {
		return domain.RequirementConditional
	}
	return domain.RequirementRequired
}

func requirementRank(kind domain.RequirementKind) int {
	switch kind {
	case domain.RequirementRequired:
		return 3
	case domain.RequirementConditional:
		return 2
	default:
		return 1
	}
}

func (walker *requirementWalker) addIssue(inputName, variableName, description string) {
	walker.complete = false
	issue := domain.RequirementIssue{
		Code:         domain.RequirementIssueUnsupportedConstruct,
		InputName:    inputName,
		VariableName: variableName,
		Description:  description,
	}
	key := requirementIssueKey{code: issue.Code, inputName: inputName, variableName: variableName}
	walker.issues[key] = issue
}

func (walker *requirementWalker) result(result domain.RequirementInspection) domain.RequirementInspection {
	result.Complete = walker.complete
	result.Inputs = make([]domain.InputRequirement, 0, len(walker.inputs))
	for _, evidence := range walker.inputs {
		result.Inputs = append(result.Inputs, domain.InputRequirement{
			Name: evidence.name, Kind: evidence.kind, Declared: evidence.declared,
			ContentType: copyString(evidence.contentType), Description: copyString(evidence.description),
			Explanation: evidence.explanation,
		})
	}
	sort.Slice(result.Inputs, func(i, j int) bool { return result.Inputs[i].Name < result.Inputs[j].Name })
	result.Variables = make([]domain.VariableRequirement, 0, len(walker.variables))
	for _, evidence := range walker.variables {
		_, supplied := walker.vars[evidence.name]
		result.Variables = append(result.Variables, domain.VariableRequirement{
			Name: evidence.name, Kind: evidence.kind, Supplied: supplied, Explanation: evidence.explanation,
		})
	}
	sort.Slice(result.Variables, func(i, j int) bool { return result.Variables[i].Name < result.Variables[j].Name })
	result.Issues = make([]domain.RequirementIssue, 0, len(walker.issues))
	for _, issue := range walker.issues {
		result.Issues = append(result.Issues, issue)
	}
	sort.Slice(result.Issues, func(i, j int) bool {
		left, right := result.Issues[i], result.Issues[j]
		if left.Code != right.Code {
			return left.Code < right.Code
		}
		if left.InputName != right.InputName {
			return left.InputName < right.InputName
		}
		return left.VariableName < right.VariableName
	})
	return result
}

func copyString(value *string) *string {
	if value == nil {
		return nil
	}
	copy := *value
	return &copy
}