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