story/internal/extraction/json_generator_test.go

package extraction

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func TestStructuredBookUsesStrictSchemaAndLongOutput(t *testing.T) {
	schema := json.RawMessage(`{"type":"object","properties":{"paragraphs":{"type":"array","items":{"type":"string"}}},"required":["paragraphs"],"additionalProperties":false}`)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var request struct {
			MaxTokens      int `json:"max_tokens"`
			ResponseFormat struct {
				Type       string `json:"type"`
				JSONSchema struct {
					Strict bool            `json:"strict"`
					Schema json.RawMessage `json:"schema"`
				} `json:"json_schema"`
			} `json:"response_format"`
		}
		if e := json.NewDecoder(r.Body).Decode(&request); e != nil {
			t.Error(e)
		}
		if request.MaxTokens != 65536 || request.ResponseFormat.Type != "json_schema" || !request.ResponseFormat.JSONSchema.Strict || string(request.ResponseFormat.JSONSchema.Schema) != string(schema) {
			t.Errorf("unexpected structured request: %+v", request)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"paragraphs\":[\"我记得那一天。\"]}"},"finish_reason":"stop"}]}`))
	}))
	defer server.Close()
	g, e := NewArkJSONGenerator(server.URL, "test-key", "doubao-seed-evolving", time.Second, 65536)
	if e != nil {
		t.Fatal(e)
	}
	out, e := g.GenerateJSONSchema(context.Background(), `{"stories":[]}`, "写回忆录", schema)
	if e != nil || !json.Valid(out) {
		t.Fatalf("output=%s error=%v", out, e)
	}
}

func TestStructuredBookReportsProviderOverdue(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusForbidden)
		_, _ = w.Write([]byte(`{"error":{"code":"AccountOverdueError"}}`))
	}))
	defer server.Close()
	g, _ := NewArkGenerator(server.URL, "test-key", "doubao-seed-evolving", time.Second)
	_, e := g.GenerateJSON(context.Background(), `{"stories":[]}`, "写回忆录")
	if !errors.Is(e, ErrModelAccountOverdue) {
		t.Fatalf("wrong error: %v", e)
	}
}

func TestBookSchemaRepairsProseEscapingButRejectsTruncatedCompletion(t *testing.T) {
	for _, finish := range []string{"stop", "length"} {
		t.Run(finish, func(t *testing.T) {
			broken := "{\"paragraphs\":[\"她说：\"先剪线头\"，我点了点头。\n记了下来。\"],\"story_id\":\"source-one\",\"event_time\":\"1970年9月\"}"
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]string{"content": broken}, "finish_reason": finish}}})
			}))
			defer server.Close()
			g, _ := NewArkJSONGenerator(server.URL, "test-key", "doubao-seed-evolving", time.Second, 65536)
			raw, err := g.GenerateJSONSchema(context.Background(), "故事原料", "写回忆录", json.RawMessage(`{"type":"object"}`))
			if finish == "length" {
				if err == nil {
					t.Fatal("truncated book accepted")
				}
				return
			}
			var out struct {
				Paragraphs []string `json:"paragraphs"`
				StoryID    string   `json:"story_id"`
				EventTime  string   `json:"event_time"`
			}
			if err != nil || json.Unmarshal(raw, &out) != nil {
				t.Fatalf("repair failed: %v", err)
			}
			if out.StoryID != "source-one" || out.EventTime != "1970年9月" || len(out.Paragraphs) != 1 || out.Paragraphs[0] != "她说：\"先剪线头\"，我点了点头。\n记了下来。" {
				t.Fatalf("repair changed prose or provenance: %+v", out)
			}
		})
	}
}

func TestBookSchemaArrayBracesPreserveAllSourceMetadataAndLiteralText(t *testing.T) {
	schema := json.RawMessage(`{"type":"object","properties":{"chapters":{"type":"array","items":{"type":"object","properties":{"paragraphs":{"type":"array","items":{"type":"string"}}}}}}}`)
	broken := `{"title":"旧事","chapters":{{"story_id":"one","event_time":"1970年","paragraphs":{"杯口有缺口，记号是{蓝边}。","她说：别碰坏了。"}},{"story_id":"two","event_time":"1971年","paragraphs":{"我学做手套。"}}}}`
	fixed := repairSchemaArrayBraces(broken, schema)
	var out struct {
		Chapters []struct {
			StoryID    string   `json:"story_id"`
			EventTime  string   `json:"event_time"`
			Paragraphs []string `json:"paragraphs"`
		} `json:"chapters"`
	}
	if json.Unmarshal([]byte(fixed), &out) != nil || len(out.Chapters) != 2 || out.Chapters[0].StoryID != "one" || out.Chapters[1].EventTime != "1971年" || out.Chapters[0].Paragraphs[0] != "杯口有缺口，记号是{蓝边}。" {
		t.Fatalf("schema array repair changed content: %s", fixed)
	}
}

func TestBookSchemaMissingParagraphOpeningBracket(t *testing.T) {
	schema := json.RawMessage(`{"type":"object","properties":{"paragraphs":{"type":"array","items":{"type":"string"}}}}`)
	broken := `{"paragraphs":"她说：别碰坏了。","杯口有缺口，记号是{蓝边}。"],"story_id":"one"}`
	fixed := repairSchemaArrayBraces(broken, schema)
	if fixed != `{"paragraphs":["她说：别碰坏了。","杯口有缺口，记号是{蓝边}。"],"story_id":"one"}` || !json.Valid([]byte(fixed)) {
		t.Fatalf("unexpected repair: %s", fixed)
	}
}