content.go

package rig

import (
	"encoding/json"
	"fmt"
	"io"
	"regexp"
	"strings"
)

// maxDecodeScan bounds the total input the candidate scan may hand to
// json.Decoder. Each '{' offset starts a decode that can run to the end of the
// string, so a reply that never closes its braces -- or nests them thousands
// deep -- makes the scan quadratic while still fitting the response cap: a
// megabyte of `{"a":` costs minutes of one core. The reply is
// endpoint-controlled, so the work it can buy is capped instead of its length.
// A real answer spends a fraction of this; a reply that exhausts it gets the
// best candidate found so far, or the same error an unparseable one gets.
const maxDecodeScan = 1 << 22

// countingReader records what a decode actually read, which is the only
// honest way to charge a decode that failed.
type countingReader struct {
	r io.Reader
	n int64
}

func (c *countingReader) Read(p []byte) (int, error) {
	n, err := c.r.Read(p)
	c.n += int64(n)
	return n, err
}

// thinkBlockRe matches a closed inline reasoning block. Some backends leave
// the model's chain-of-thought in the message content instead of a separate
// field, and reasoning about a JSON answer routinely contains braces, which
// would defeat the brace scan in DecodeJSONObject. Reasoning drafts are also
// not the answer: a JSON object inside a think block must never be returned
// as the reply.
var thinkBlockRe = regexp.MustCompile(`(?s)<think>.*?</think>`)

// StripReasoning removes leaked chain-of-thought from assistant content.
// Closed <think>...</think> blocks are dropped. Some serving templates
// consume the opening tag, leaving bare reasoning that ends in </think>:
// everything through the last closing tag is reasoning, not reply. An opening
// tag left without a closing one opens reasoning that never ends, so
// everything from it on is dropped too -- a run that stopped inside the model's
// thinking has no reply, and passing the thinking through as one is worse than
// returning nothing. The result is trimmed of surrounding space.
func StripReasoning(s string) string {
	raw := strings.TrimSpace(thinkBlockRe.ReplaceAllString(s, ""))
	if i := strings.LastIndex(raw, "</think>"); i >= 0 {
		raw = strings.TrimSpace(raw[i+len("</think>"):])
	}
	if i := strings.Index(raw, "<think>"); i >= 0 {
		raw = strings.TrimSpace(raw[:i])
	}
	return raw
}

// DecodeJSONObject parses assistant content as a single JSON object,
// tolerating inline reasoning, stray prose, and code fences around it. When
// several candidates parse, the longest wins: the real reply contains every
// object nested inside it, while the schema sketches that leaked reasoning
// tends to include ("the shape is {\"stories\": [...]}") are short. Taking
// the first parseable candidate instead returned those sketches as the reply.
// Returns ErrNotJSON when no object is present at all, ErrInvalidJSON when
// braces are present but nothing parses.
func DecodeJSONObject(s string) (map[string]any, error) {
	raw := StripReasoning(s)
	var m map[string]any
	if err := json.Unmarshal([]byte(raw), &m); err == nil {
		return m, nil
	}
	var best map[string]any
	var bestLen, scanned int64
	for i := 0; i < len(raw) && scanned < maxDecodeScan; i++ {
		if raw[i] != '{' {
			continue
		}
		counter := &countingReader{r: strings.NewReader(raw[i:])}
		dec := json.NewDecoder(counter)
		var candidate map[string]any
		err := dec.Decode(&candidate)
		scanned += counter.n
		if err != nil {
			continue
		}
		end := dec.InputOffset()
		if end > bestLen {
			best, bestLen = candidate, end
		}
		// Every '{' inside a candidate that parsed opens an object nested in
		// it, and a nested object is shorter than the one containing it, so
		// the scan resumes past the whole candidate instead of re-parsing it
		// once per brace it holds.
		i += int(end) - 1
	}
	if best != nil {
		return best, nil
	}
	if !strings.Contains(raw, "{") {
		return nil, fmt.Errorf("decode json object: %w", ErrNotJSON)
	}
	return nil, fmt.Errorf("decode json object: %w", ErrInvalidJSON)
}