Agent-System/internal/app/pi_segment_test.go

package app

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

	"chat-agent-system/internal/agent"
	"chat-agent-system/internal/app/intentflow"
	"chat-agent-system/internal/config"
)

type recordingSegmentBackend struct {
	last    agent.Request
	payload []byte
	body    string
	fail    string
}

func (r *recordingSegmentBackend) Name() string { return "recording-segment" }

func (r *recordingSegmentBackend) Execute(_ context.Context, req agent.Request) (*agent.Result, error) {
	r.last = req
	if b, err := os.ReadFile(filepath.Join(req.WorkDir, "input.json")); err == nil {
		r.payload = b
	}
	if r.fail != "" {
		return &agent.Result{Success: false, ErrorReason: r.fail}, nil
	}
	body := r.body
	if body == "" {
		body = `{"schema_version":"intentflow.segment_proposal.v1","digest":"d1","spans":[{"start":0,"end":3,"text":"abc"}]}`
	}
	return &agent.Result{Success: true, Patches: map[string][]byte{"spans.md": []byte(body)}}, nil
}

func TestPiSegmentProposer_NoToolsAndDigest(t *testing.T) {
	rec := &recordingSegmentBackend{}
	p := &PiSegmentProposer{Backend: rec, Model: "test"}
	got, err := p.Propose(context.Background(), intentflow.SegmentProposalRequest{
		Text: "abc", RequestDigest: "d1",
	})
	if err != nil {
		t.Fatal(err)
	}
	if got == nil || len(got.Spans) != 1 {
		t.Fatalf("got=%+v", got)
	}
	if !strings.Contains(strings.ToLower(rec.last.Instructions), "zero tool") {
		t.Fatalf("instructions=%s", rec.last.Instructions)
	}
	var in map[string]any
	if err := json.Unmarshal(rec.payload, &in); err != nil {
		t.Fatal(err)
	}
	if in["digest"] != "d1" {
		t.Fatalf("input=%v", in)
	}
}

func TestPiSegmentProposer_MalformedAndTimeout(t *testing.T) {
	p := &PiSegmentProposer{Backend: &recordingSegmentBackend{body: "not-json"}}
	if _, err := p.Propose(context.Background(), intentflow.SegmentProposalRequest{Text: "abc", RequestDigest: "d1"}); err != intentflow.ErrSegmentProposalMalformed {
		t.Fatalf("malformed: %v", err)
	}
	p = &PiSegmentProposer{Backend: &recordingSegmentBackend{fail: "TIMEOUT"}}
	if _, err := p.Propose(context.Background(), intentflow.SegmentProposalRequest{Text: "abc", RequestDigest: "d1"}); err != intentflow.ErrSegmentProposalTimeout {
		t.Fatalf("timeout: %v", err)
	}
	p = &PiSegmentProposer{Backend: &recordingSegmentBackend{body: `{"schema_version":"intentflow.segment_proposal.v1","digest":"d1","spans":[{"start":0,"end":99,"text":"nope"}]}`}}
	if _, err := p.Propose(context.Background(), intentflow.SegmentProposalRequest{Text: "abc", RequestDigest: "d1"}); err != intentflow.ErrSegmentProposalOutOfBounds {
		t.Fatalf("oob: %v", err)
	}
}

func TestPiSegmentProposer_LiveMiniMax(t *testing.T) {
	if os.Getenv("RUN_PI_E2E") != "1" {
		t.Skip("set RUN_PI_E2E=1 for live PI segment proposal")
	}
	if testing.Short() {
		t.Skip("live PI segment is not a -short test")
	}
	if err := os.Chdir(filepath.Clean(filepath.Join("..", ".."))); err != nil {
		t.Fatal(err)
	}
	cfg, err := config.Load("")
	if err != nil || cfg == nil {
		t.Fatalf("config: %v", err)
	}
	if strings.TrimSpace(cfg.PiAgentRunnerAuthToken) == "" {
		if tok := runnerTokenFromDotenv(); tok != "" {
			cfg.PiAgentRunnerAuthToken = tok
		}
	}
	if strings.TrimSpace(cfg.PiAgentRunnerAuthToken) == "" {
		t.Fatal("runner auth token missing")
	}
	cfg.PiAgentRunnerURL = "http://127.0.0.1:18865"
	dataDir, err := filepath.Abs("data/dev")
	if err != nil {
		t.Fatal(err)
	}
	p := &PiSegmentProposer{
		Backend: agent.NewBackend("pi-agent", cfg),
		Model:   cfg.PiAgentModel,
		Timeout: 4 * time.Minute,
		DataDir: dataDir,
	}
	text := "帮我看看电气安全条款生成立项建议"
	start := time.Now()
	got, err := p.Propose(context.Background(), intentflow.SegmentProposalRequest{
		Text: text, RequestDigest: "live-u9-digest",
	})
	if err != nil {
		t.Fatalf("live propose: %v", err)
	}
	t.Logf("elapsed=%s proposal=%+v", time.Since(start).Round(time.Millisecond), got)
	if got.Digest != "live-u9-digest" {
		t.Fatalf("digest=%s", got.Digest)
	}
}