xinmi/internal/application/service/wiki_ingest_cite_live_test.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/agent"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
)

type wikiLiveChatProvider interface{ chat.Chat }

type wikiLimitedLiveChat struct {
	wikiLiveChatProvider
	calls     int
	limit     int
	directory string
}

func (m *wikiLimitedLiveChat) Chat(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
	if m.calls >= m.limit {
		return nil, errors.New("live provider request budget exhausted")
	}
	m.calls++
	response, err := m.wikiLiveChatProvider.Chat(ctx, messages, options)
	if err == nil && response != nil {
		path := filepath.Join(m.directory, "response-private-"+time.Now().UTC().Format("20060102T150405.000000000")+".json")
		if os.WriteFile(path, []byte(response.Content), 0600) != nil {
			return nil, errors.New("could not preserve private live diagnostic response")
		}
	}
	return response, err
}

// Explicit opt-in only. The private fixture contains the configured model and
// one read-only source snapshot. It makes one provider invocation by default;
// separately authorized recovery mode has a hard two-invocation limit. There
// is no database writer, and output is counts only. Never enable this through
// the default test suite or pass model secrets on the command line.
func TestWikiCitationLiveSingleBatch(t *testing.T) {
	path := os.Getenv("WIKI_CITATION_LIVE_FIXTURE")
	if path == "" {
		t.Skip("requires an explicitly authorized private source/model fixture")
	}
	var fixture struct {
		Model      types.Model        `json:"model"`
		Candidates combinedExtraction `json:"candidates"`
		Chunks     []*types.Chunk     `json:"chunks"`
		AESKey     string             `json:"aes_key"`
	}
	data, err := os.ReadFile(path)
	if err != nil || json.Unmarshal(data, &fixture) != nil {
		t.Fatal("could not read the private live-test fixture")
	}
	if fixture.AESKey != "" {
		t.Setenv("SYSTEM_AES_KEY", fixture.AESKey)
	}
	parameters, _ := json.Marshal(fixture.Model.Parameters)
	if fixture.Model.Parameters.Scan(parameters) != nil || fixture.Model.Parameters.APIKey == "" {
		t.Fatal("live-test configured model credentials unavailable")
	}
	remote, err := chat.NewRemoteChat(chat.ConfigFromModel(&fixture.Model, "", ""))
	if err != nil {
		t.Fatal("live-test model initialization failed; details suppressed")
	}
	candidates := renderCandidateSlugsXML(fixture.Candidates.Entities, fixture.Candidates.Concepts)
	encoded, handles := encodeWikiCitationCandidates(candidates)
	batches := splitChunksIntoCitationBatches(fixture.Chunks)
	if len(batches) == 0 {
		t.Fatal("live-test fixture has no citation batches")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
	defer cancel()
	started := time.Now()
	model := &wikiLimitedLiveChat{wikiLiveChatProvider: remote, limit: 1, directory: filepath.Dir(path)}
	var parsed citationBatchResult
	if os.Getenv("WIKI_CITATION_LIVE_ALLOW_REPAIR") == "1" {
		model.limit = 2
		parsed, err = (&wikiIngestService{}).classifyCitationBatch(ctx, model, candidates, batches[0], "Chinese", true)
	} else {
		var raw string
		raw, err = (&wikiIngestService{}).generateWithTemplateRequest(ctx, model, agent.WikiChunkCitationPrompt, map[string]string{
			"CandidateSlugs": encoded, "ChunksXML": renderChunksXML(batches[0]), "Language": "Chinese",
		}, "", 1)
		if err == nil {
			parsed, err = parseWikiCitationBatch(raw, batches[0], wikiCandidateSlugSet(candidates), handles)
		}
	}
	if err != nil {
		t.Fatalf("single live citation validation failed (provider_requests=%d candidates=%d first_batch_chunks=%d first_batch_chars=%d); private diagnostic retained", model.calls, handles.Len(), len(batches[0].chunks), batches[0].totalRuneLen)
	}
	t.Logf("single_batch_live_pass=true provider_requests=%d provider_request_limit=%d database_writes=0 candidates=%d candidate_chars=%d planned_batches=%d first_batch_chunks=%d first_batch_chars=%d citation_keys=%d declared_new_slugs=%d elapsed_seconds=%.1f",
		model.calls, model.limit, handles.Len(), len([]rune(encoded)), len(batches), len(batches[0].chunks), batches[0].totalRuneLen, len(parsed.Citations), len(parsed.NewSlugs), time.Since(started).Seconds())
}