story/internal/extraction/json_generator.go

package extraction

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/kaptinlin/jsonrepair"
)

var ErrModelAccountOverdue = errors.New("model account overdue")

// A whole memoir needs more output space than a single extraction result.
// Keep its limit separate so ordinary chat extraction retains its existing cap.
func NewArkJSONGenerator(baseURL, apiKey, model string, timeout time.Duration, maxTokens int) (*ArkGenerator, error) {
	if maxTokens < 256 || maxTokens > 262144 {
		return nil, errors.New("invalid structured output limit")
	}
	g, err := NewArkGenerator(baseURL, apiKey, model, timeout)
	if err == nil {
		g.jsonMaxTokens = maxTokens
	}
	return g, err
}

// GenerateJSON reuses the configured tenant model for an explicitly separate
// structured task. The caller validates schema, provenance and authorization.
func (g *ArkGenerator) GenerateJSON(ctx context.Context, content, prompt string) (json.RawMessage, error) {
	return g.generateJSON(ctx, content, prompt, nil)
}

// Schema mode keeps long book prose inside correctly escaped paragraph arrays.
func (g *ArkGenerator) GenerateJSONSchema(ctx context.Context, content, prompt string, schema json.RawMessage) (json.RawMessage, error) {
	if !json.Valid(schema) {
		return nil, errors.New("invalid structured output schema")
	}
	return g.generateJSON(ctx, content, prompt, schema)
}

func (g *ArkGenerator) generateJSON(ctx context.Context, content, prompt string, schema json.RawMessage) (json.RawMessage, error) {
	if strings.TrimSpace(content) == "" || len(content) > 2<<20 {
		return nil, errors.New("invalid structured input")
	}
	maxTokens := g.jsonMaxTokens
	if maxTokens == 0 {
		maxTokens = 8192
	}
	payload := map[string]any{"model": g.model, "messages": []map[string]string{{"role": "system", "content": prompt}, {"role": "user", "content": content}}, "temperature": 0, "max_tokens": maxTokens}
	if g.provider == "ark" {
		payload["thinking"] = map[string]string{"type": "disabled"}
	}
	if schema != nil {
		payload["response_format"] = map[string]any{"type": "json_schema", "json_schema": map[string]any{"name": "memoir_book", "strict": true, "schema": schema}}
	} else if providerSupportsJSONMode(g.provider) {
		payload["response_format"] = map[string]string{"type": "json_object"}
	}
	raw, e := json.Marshal(payload)
	if e != nil {
		return nil, e
	}
	req, e := http.NewRequestWithContext(ctx, "POST", g.endpoint, bytes.NewReader(raw))
	if e != nil {
		return nil, e
	}
	req.Header.Set("Authorization", "Bearer "+g.apiKey)
	req.Header.Set("Content-Type", "application/json")
	resp, e := g.client.Do(req)
	if e != nil {
		return nil, errors.New("structured model unavailable")
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		var failure struct {
			Error struct {
				Code string `json:"code"`
			} `json:"error"`
		}
		_ = json.NewDecoder(io.LimitReader(resp.Body, 64<<10)).Decode(&failure)
		if failure.Error.Code == "AccountOverdueError" {
			return nil, ErrModelAccountOverdue
		}
		return nil, errors.New("structured model request rejected")
	}
	var out struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
	}
	if e = json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&out); e != nil || len(out.Choices) != 1 || out.Choices[0].FinishReason == "length" {
		return nil, errors.New("structured model result incomplete")
	}
	result := strings.TrimSpace(extractJSONObject(out.Choices[0].Message.Content))
	// Some book models emit prose with an unescaped quotation mark or newline,
	// even when schema output is requested. Repair syntax locally; the book
	// validator still requires every original story ID and its unchanged date.
	// A length-truncated completion was already rejected above.
	if schema != nil && !json.Valid([]byte(result)) {
		result = repairSchemaArrayBraces(result, schema)
		if !json.Valid([]byte(result)) {
			if repaired, err := jsonrepair.Repair(result); err == nil && json.Valid([]byte(repaired)) {
				result = repaired
			}
		}
	}
	if !json.Valid([]byte(result)) {
		var parsed any
		parseErr := json.Unmarshal([]byte(result), &parsed)
		return json.RawMessage(result), fmt.Errorf("structured model result invalid (content bytes=%d, extracted bytes=%d, finish=%s): %w", len(out.Choices[0].Message.Content), len(result), out.Choices[0].FinishReason, parseErr)
	}
	return json.RawMessage(result), nil
}

// The Seed book endpoint can surround array values with braces, for example
// "paragraphs":{"first paragraph","second paragraph"}. The schema gives an
// unambiguous repair: change only those container delimiters, keeping every
// byte of prose, identifier and date. Braces inside strings remain untouched.
func repairSchemaArrayBraces(input string, schema json.RawMessage) string {
	var root map[string]any
	if json.Unmarshal(schema, &root) != nil {
		return input
	}
	arrays := map[string]bool{}
	var visit func(map[string]any)
	visit = func(n map[string]any) {
		if properties, ok := n["properties"].(map[string]any); ok {
			for key, value := range properties {
				if child, ok := value.(map[string]any); ok {
					if child["type"] == "array" {
						arrays[key] = true
					}
					visit(child)
				}
			}
		}
		if items, ok := n["items"].(map[string]any); ok {
			visit(items)
		}
	}
	visit(root)
	output := []byte(input)
	opens := map[int]bool{}
	missingOpens := map[int]bool{}
	stack := []bool{}
	space := func(b byte) bool { return b == ' ' || b == '\n' || b == '\r' || b == '\t' }
	for i := 0; i < len(input); i++ {
		if input[i] == '"' {
			start := i
			for i++; i < len(input); i++ {
				if input[i] == '\\' {
					i++
					continue
				}
				if input[i] == '"' {
					break
				}
			}
			if i >= len(input) {
				break
			}
			var key string
			if json.Unmarshal([]byte(input[start:i+1]), &key) == nil && arrays[key] {
				j := i + 1
				for j < len(input) && space(input[j]) {
					j++
				}
				if j < len(input) && input[j] == ':' {
					j++
					for j < len(input) && space(input[j]) {
						j++
					}
					if j < len(input) && input[j] == '{' {
						opens[j] = true
					}
					// Also repair a missing opening bracket in a string sequence:
					// "paragraphs":"first","second"]. Its closing ] must exist.
					if j < len(input) && input[j] == '"' {
						quoted := false
						for k := j; k < len(input); k++ {
							if input[k] == '\\' && quoted {
								k++
								continue
							}
							if input[k] == '"' {
								quoted = !quoted
								continue
							}
							if !quoted && input[k] == '}' {
								break
							}
							if !quoted && input[k] == ']' {
								missingOpens[j] = true
								break
							}
						}
					}
				}
			}
			continue
		}
		if input[i] == '{' {
			array := opens[i]
			stack = append(stack, array)
			if array {
				output[i] = '['
			}
		}
		if input[i] == '}' && len(stack) > 0 {
			array := stack[len(stack)-1]
			stack = stack[:len(stack)-1]
			if array {
				output[i] = ']'
			}
		}
	}
	if len(missingOpens) == 0 {
		return string(output)
	}
	var repaired strings.Builder
	for i, b := range output {
		if missingOpens[i] {
			repaired.WriteByte('[')
		}
		repaired.WriteByte(b)
	}
	return repaired.String()
}