iop/apps/edge/internal/openai/strict_output.go

package openai

import (
	"fmt"
	"html"
	"regexp"
	"strings"
)

func strictOutputContractInstruction(policy strictOutputPolicy) string {
	if !policy.ContractInstruction {
		return ""
	}
	resultTag := policy.XMLResultTag
	if resultTag == "" {
		resultTag = "result"
	}
	return fmt.Sprintf("Strict output contract: follow the client's XML tool protocol. When you complete the task, output exactly one <%s> block and no text outside that XML block. The <%s> content must be the direct final answer that should be shown to the end user. Do not summarize that you answered, completed, explained, greeted, or performed the task. Do not write phrases like \"I completed\", \"I explained\", \"I greeted\", \"summary\", or \"main points\" unless the user explicitly asks for a summary. Put the actual user-facing response inside <%s>.", policy.XMLCompletionTool, resultTag, resultTag)
}

func prependSystemMessage(messages []chatMessage, content string) []chatMessage {
	out := make([]chatMessage, 0, len(messages)+1)
	out = append(out, chatMessage{Role: "system", Content: content})
	out = append(out, messages...)
	return out
}

func normalizeCompletionOutput(policy strictOutputPolicy, content, reasoning string) (string, string, bool) {
	if !policy.Strict {
		return content, reasoning, false
	}
	normalized := normalizeStrictAgentContent(content)
	if strings.TrimSpace(normalized) == "" && strings.TrimSpace(reasoning) != "" {
		normalized = normalizeStrictAgentContent(reasoning)
	}
	if strings.TrimSpace(normalized) == "" {
		normalized = "Model returned no assistant content."
	}
	if policy.XMLCompletionTool != "" && !startsWithXMLRoot(normalized) {
		normalized = wrapXMLCompletion(policy, normalized)
	}
	return normalized, "", normalized != content || reasoning != ""
}

var (
	agentThinkingBlockRE       = regexp.MustCompile(`(?is)<think(?:ing)?\b[^>]*>.*?</think(?:ing)?\s*>`)
	agentCodeFenceLineRE       = regexp.MustCompile(`(?m)^\s*` + "```" + `[A-Za-z0-9_-]*\s*$`)
	agentXMLOpenTagRE          = regexp.MustCompile(`(?is)<\s*([A-Za-z_][A-Za-z0-9_.:-]*)\b[^>]*>`)
	agentXMLCloseTagRE         = regexp.MustCompile(`(?is)</\s*([A-Za-z_][A-Za-z0-9_.:\s-]*)\s*>`)
	completedUseToolRE         = regexp.MustCompile("(?is)completed.{0,300}?use\\s+(?:the\\s+)?`?([A-Za-z_][A-Za-z0-9_.:-]*)`?\\s+tool")
	xmlToolWithResultREPattern = `(?is)<\s*%s\b[^>]*>\s*<\s*([A-Za-z_][A-Za-z0-9_.:-]*)\b`
)

func normalizeStrictAgentContent(content string) string {
	cleaned := strings.TrimSpace(content)
	cleaned = agentThinkingBlockRE.ReplaceAllString(cleaned, "")
	cleaned = agentCodeFenceLineRE.ReplaceAllString(cleaned, "")
	cleaned = strings.TrimSpace(cleaned)
	for strings.HasPrefix(cleaned, "---") {
		cleaned = strings.TrimSpace(strings.TrimPrefix(cleaned, "---"))
	}
	cleaned = strings.TrimSpace(cleaned)
	if block, ok := firstXMLRootBlock(cleaned); ok {
		return block
	}
	return cleaned
}

func startsWithXMLRoot(content string) bool {
	trimmed := strings.TrimSpace(content)
	open := agentXMLOpenTagRE.FindStringIndex(trimmed)
	return open != nil && open[0] == 0
}

func inferXMLCompletionContract(prompt string) (string, string) {
	tool := ""
	if match := completedUseToolRE.FindStringSubmatch(prompt); len(match) == 2 {
		tool = match[1]
	}
	if tool == "" {
		return "", ""
	}
	resultTag := inferXMLResultTag(prompt, tool)
	if resultTag == "" {
		resultTag = "result"
	}
	return tool, resultTag
}

func inferXMLResultTag(prompt, tool string) string {
	toolPattern := regexp.QuoteMeta(tool)
	re := regexp.MustCompile(fmt.Sprintf(xmlToolWithResultREPattern, toolPattern))
	if match := re.FindStringSubmatch(prompt); len(match) == 2 {
		return match[1]
	}
	return ""
}

func wrapXMLCompletion(policy strictOutputPolicy, content string) string {
	resultTag := policy.XMLResultTag
	if resultTag == "" {
		resultTag = "result"
	}
	return "<" + policy.XMLCompletionTool + ">\n<" + resultTag + ">" + html.EscapeString(strings.TrimSpace(content)) + "</" + resultTag + ">\n</" + policy.XMLCompletionTool + ">"
}

func firstXMLRootBlock(content string) (string, bool) {
	open := agentXMLOpenTagRE.FindStringSubmatchIndex(content)
	if open == nil {
		return "", false
	}
	root := content[open[2]:open[3]]
	search := content[open[1]:]
	for _, close := range agentXMLCloseTagRE.FindAllStringSubmatchIndex(search, -1) {
		name := search[close[2]:close[3]]
		if canonicalXMLName(name) != canonicalXMLName(root) {
			continue
		}
		closeStart := open[1] + close[0]
		block := content[open[0]:closeStart] + "</" + root + ">"
		return strings.TrimSpace(block), true
	}
	return "", false
}

func canonicalXMLName(name string) string {
	var b strings.Builder
	for _, r := range strings.ToLower(name) {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') {
			b.WriteRune(r)
		}
	}
	return b.String()
}