geo/server/internal/tenant/app/online_generation_replay_test.go

package app

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

	sharedconfig "github.com/geo-platform/tenant-api/internal/shared/config"
	"github.com/geo-platform/tenant-api/internal/shared/llm"
	"github.com/geo-platform/tenant-api/internal/tenant/prompts"
)

type onlineGenerationFixture struct {
	GenerationTasks []onlineGenerationTask `json:"generation_tasks"`
}

type onlineGenerationTask struct {
	ID              int64                  `json:"id"`
	ArticleID       int64                  `json:"article_id"`
	InputParamsJSON map[string]interface{} `json:"input_params_json"`
}

type onlineKnowledgeFixture struct {
	KnowledgeItems []onlineKnowledgeItem `json:"knowledge_items"`
}

type onlineKnowledgeItem struct {
	ID              int64   `json:"id"`
	GroupID         int64   `json:"group_id"`
	Name            string  `json:"name"`
	SourceType      string  `json:"source_type"`
	SourceURI       *string `json:"source_uri"`
	Status          string  `json:"status"`
	ContentText     *string `json:"content_text"`
	MarkdownContent *string `json:"markdown_content"`
}

func TestReplayOnlineFailedArticleGeneration(t *testing.T) {
	if os.Getenv("RUN_ONLINE_GENERATION_REPLAY") != "1" {
		t.Skip("set RUN_ONLINE_GENERATION_REPLAY=1 to replay online failed generation against Ark")
	}

	fixtureDir := strings.TrimSpace(os.Getenv("ONLINE_GENERATION_FIXTURE_DIR"))
	if fixtureDir == "" {
		fixtureDir = filepath.Clean("../../../../tmp/online-generation-fixtures")
	}
	taskID := int64(2)
	if strings.TrimSpace(os.Getenv("ONLINE_GENERATION_TASK_ID")) == "1" {
		taskID = 1
	}

	var fixture onlineGenerationFixture
	readJSONFixture(t, filepath.Join(fixtureDir, "db-compact.json"), &fixture)
	task := findOnlineGenerationTask(t, fixture.GenerationTasks, taskID)

	var knowledge onlineKnowledgeFixture
	readJSONFixture(t, filepath.Join(fixtureDir, "knowledge-compact.json"), &knowledge)

	if promptsPath := strings.TrimSpace(os.Getenv("PROMPTS_CONFIG_FILE")); promptsPath == "" {
		prompts.SetConfigFile(filepath.Clean("../../../configs/prompts.yml"))
	}

	knowledgePrompt := buildReplayKnowledgePrompt(task.InputParamsJSON, knowledge)
	templateKey, templateName, promptTemplate := extractTemplateSnapshot(task.InputParamsJSON)
	prompt := buildGenerationPrompt(templateKey, templateName, promptTemplate, task.InputParamsJSON, knowledgePrompt)
	t.Logf("task_id=%d article_id=%d prompt_chars=%d knowledge_chars=%d", task.ID, task.ArticleID, len([]rune(prompt)), len([]rune(knowledgePrompt)))

	cfg, err := sharedconfig.Load(filepath.Clean("../../../configs/config.yaml"))
	if err != nil {
		t.Fatalf("load config: %v", err)
	}
	client := llm.New(cfg.LLM)
	if err := client.Validate(); err != nil {
		t.Fatalf("validate llm: %v", err)
	}

	start := time.Now()
	var outputChars int
	result, err := client.Generate(context.Background(), llm.GenerateRequest{
		Prompt:          prompt,
		Timeout:         cfg.Generation.ArticleTimeout,
		MaxOutputTokens: cfg.LLM.MaxOutputTokens,
	}, func(delta string) {
		outputChars += len([]rune(delta))
	})
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("replay generation failed after %s with %d streamed chars: %v", elapsed.Round(time.Second), outputChars, err)
	}
	t.Logf("replay generation completed after %s with %d chars, model=%s", elapsed.Round(time.Second), len([]rune(result.Content)), result.Model)
}

func readJSONFixture(t *testing.T, path string, target interface{}) {
	t.Helper()
	data, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read fixture %s: %v", path, err)
	}
	if err := json.Unmarshal(data, target); err != nil {
		t.Fatalf("parse fixture %s: %v", path, err)
	}
}

func findOnlineGenerationTask(t *testing.T, tasks []onlineGenerationTask, id int64) onlineGenerationTask {
	t.Helper()
	for _, task := range tasks {
		if task.ID == id {
			return task
		}
	}
	t.Fatalf("task %d not found in fixture", id)
	return onlineGenerationTask{}
}

func buildReplayKnowledgePrompt(params map[string]interface{}, fixture onlineKnowledgeFixture) string {
	groupIDs := make(map[int64]struct{})
	for _, id := range extractKnowledgeGroupIDs(params["knowledge_group_ids"]) {
		groupIDs[id] = struct{}{}
	}
	if len(groupIDs) == 0 {
		return ""
	}

	ctx := &KnowledgeContext{Snippets: []KnowledgeSnippet{}, PreciseFacts: []string{}}
	seenFacts := make(map[string]struct{})
	for _, item := range fixture.KnowledgeItems {
		if strings.TrimSpace(item.Status) != "completed" {
			continue
		}
		if _, ok := groupIDs[item.GroupID]; !ok {
			continue
		}
		text := strings.TrimSpace(derefKnowledgeString(item.MarkdownContent))
		if text == "" {
			text = strings.TrimSpace(derefKnowledgeString(item.ContentText))
		}
		if text == "" {
			continue
		}
		facts := extractKnowledgePreciseFacts(text)
		for _, fact := range facts {
			entry := strings.TrimSpace(fact)
			if entry == "" {
				continue
			}
			if _, exists := seenFacts[entry]; exists {
				continue
			}
			seenFacts[entry] = struct{}{}
			ctx.PreciseFacts = append(ctx.PreciseFacts, entry)
		}
		ctx.Snippets = append(ctx.Snippets, KnowledgeSnippet{
			GroupID:         item.GroupID,
			KnowledgeItemID: item.ID,
			ItemName:        item.Name,
			SourceType:      item.SourceType,
			SourceURI:       item.SourceURI,
			Text:            text,
			PreciseFacts:    facts,
		})
	}
	return (&KnowledgeService{}).RenderPromptSection(ctx)
}