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, preserveReasoning bool) (string, string, bool) {
	var sanitized bool
	content, sanitized = sanitizeKnownChatTemplateSentinels(content)
	reasoning, reasoningSanitized := sanitizeKnownChatTemplateSentinels(reasoning)
	sanitized = sanitized || reasoningSanitized
	if !policy.Strict {
		return content, reasoning, sanitized
	}
	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)
	}
	if preserveReasoning {
		return normalized, reasoning, sanitized || normalized != content
	}
	return normalized, "", sanitized || normalized != content || reasoning != ""
}

var knownChatTemplateSentinels = []string{
	"<|mask_end|>",
}

func sanitizeKnownChatTemplateSentinels(s string) (string, bool) {
	cleaned := s
	for _, token := range knownChatTemplateSentinels {
		cleaned = strings.ReplaceAll(cleaned, token, "")
	}
	if cleaned == s {
		return s, false
	}
	return strings.TrimSpace(cleaned), true
}

type streamSentinelFilter struct {
	pending string
}

func (f *streamSentinelFilter) Append(delta string) string {
	if delta == "" {
		return ""
	}
	f.pending += delta
	flushLen := streamSentinelSafeFlushLen(f.pending)
	out := f.pending[:flushLen]
	f.pending = f.pending[flushLen:]
	out, _ = sanitizeKnownChatTemplateSentinels(out)
	return out
}

func (f *streamSentinelFilter) Flush() string {
	out := f.pending
	f.pending = ""
	if isKnownSentinelPrefix(out) {
		return ""
	}
	out, _ = sanitizeKnownChatTemplateSentinels(out)
	return out
}

func streamSentinelSafeFlushLen(s string) int {
	hold := 0
	for _, token := range knownChatTemplateSentinels {
		max := len(token) - 1
		if max > len(s) {
			max = len(s)
		}
		for n := max; n > 0; n-- {
			if strings.HasSuffix(s, token[:n]) {
				if n > hold {
					hold = n
				}
				break
			}
		}
	}
	return len(s) - hold
}

func isKnownSentinelPrefix(s string) bool {
	if s == "" {
		return false
	}
	for _, token := range knownChatTemplateSentinels {
		if strings.HasPrefix(token, s) {
			return true
		}
	}
	return false
}

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()
}