Cyrene_v2/internal/runtime/memory/llmextractor_test.go

package memory

import (
	"context"
	"errors"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// mockAdapter is a scriptable model.Adapter for tests. With no fn it returns an
// empty candidate array, so a test that only needs a successful call can use it
// directly.
type mockAdapter struct {
	fn func(ctx context.Context, req model.Request, opts model.Options) (model.Response, error)
}

func (m *mockAdapter) Name() string { return "mock" }

func (m *mockAdapter) Chat(ctx context.Context, req model.Request, opts model.Options) (model.Response, error) {
	if m.fn != nil {
		return m.fn(ctx, req, opts)
	}
	return model.Response{Content: "[]"}, nil
}

func (m *mockAdapter) ChatStream(ctx context.Context, req model.Request, onEvent func(model.StreamEvent)) (model.Response, error) {
	return model.Response{}, model.ErrStream
}

func TestLLMExtractor_Extract_Parse_正常(t *testing.T) {
	t.Parallel()
	adapter := &mockAdapter{fn: func(_ context.Context, _ model.Request, _ model.Options) (model.Response, error) {
		return model.Response{Content: `[
			{"content":"用户喜欢冰美式","scope":"user","kind":"relationship","tags":["咖啡"],"confidence":0.9,"source_conversation":"conv.1"},
			{"content":"项目Alpha周三开会","scope":"project","project_id":"project.alpha","source_conversation":"conv.1"}
		]`}, nil
	}}
	ex := NewLLMExtractor(adapter, "persona.a", "role.1")

	items, err := ex.Extract(context.Background(), []model.Message{
		model.TextMessage(model.RoleUser, "你好"),
	})
	if err != nil {
		t.Fatalf("Extract: %v", err)
	}
	if len(items) != 2 {
		t.Fatalf("got %d items, want 2", len(items))
	}
	if items[0].Content != "用户喜欢冰美式" || items[0].Scope != ScopeUser || items[0].Confidence != 0.9 {
		t.Errorf("items[0] = %+v, want parsed user item", items[0])
	}
	if items[0].SourceConversation != "conv.1" {
		t.Errorf("items[0].SourceConversation = %q, want conv.1", items[0].SourceConversation)
	}
	if items[1].ProjectID != "project.alpha" || items[1].Scope != ScopeProject {
		t.Errorf("items[1] = %+v, want project-scoped item", items[1])
	}
	if items[0].PersonaID != "persona.a" || items[0].RoleID != "role.1" {
		t.Errorf("items[0] author = %q/%q, want persona.a/role.1", items[0].PersonaID, items[0].RoleID)
	}
}

func TestLLMExtractor_Extract_InvalidScopeKind_Fallback_正常(t *testing.T) {
	t.Parallel()
	adapter := &mockAdapter{fn: func(_ context.Context, _ model.Request, _ model.Options) (model.Response, error) {
		return model.Response{Content: `[{"content":"一个事实","scope":"bogus","kind":"bogus"}]`}, nil
	}}
	ex := NewLLMExtractor(adapter, "persona.a", "role.1")

	items, err := ex.Extract(context.Background(), []model.Message{
		model.TextMessage(model.RoleUser, "hi"),
	})
	if err != nil {
		t.Fatalf("Extract: %v", err)
	}
	if len(items) != 1 {
		t.Fatalf("got %d items, want 1", len(items))
	}
	// Bogus scope/kind fall back to the extractor's configured defaults.
	if items[0].Scope != ScopePersona || items[0].Kind != KindRelationship {
		t.Errorf("fallback scope/kind = %q/%q, want persona/relationship", items[0].Scope, items[0].Kind)
	}
}

func TestLLMExtractor_Extract_Dedup_正常(t *testing.T) {
	t.Parallel()
	// The model emits the same memory three times, differing only in case and
	// surrounding whitespace; they must collapse into one item.
	adapter := &mockAdapter{fn: func(_ context.Context, _ model.Request, _ model.Options) (model.Response, error) {
		return model.Response{Content: `[
			{"content":"记得喝咖啡"},
			{"content":"  记得喝咖啡  "},
			{"content":"记得喝咖啡"}
		]`}, nil
	}}
	ex := NewLLMExtractor(adapter, "persona.a", "role.1")

	items, err := ex.Extract(context.Background(), []model.Message{
		model.TextMessage(model.RoleUser, "hi"),
	})
	if err != nil {
		t.Fatalf("Extract: %v", err)
	}
	if len(items) != 1 {
		t.Fatalf("got %d items, want 1 after dedup", len(items))
	}
}

func TestLLMExtractor_Extract_Empty_正常(t *testing.T) {
	t.Parallel()
	ex := NewLLMExtractor(&mockAdapter{}, "persona.a", "role.1")

	items, err := ex.Extract(context.Background(), []model.Message{
		model.TextMessage(model.RoleUser, "hi"),
	})
	if err != nil {
		t.Fatalf("Extract: %v", err)
	}
	if len(items) != 0 {
		t.Fatalf("got %d items, want 0 for empty candidate response", len(items))
	}
}

func TestLLMExtractor_Extract_InvalidJSON_异常(t *testing.T) {
	t.Parallel()
	adapter := &mockAdapter{fn: func(_ context.Context, _ model.Request, _ model.Options) (model.Response, error) {
		return model.Response{Content: "{not-json}"}, nil
	}}
	ex := NewLLMExtractor(adapter, "persona.a", "role.1")

	if _, err := ex.Extract(context.Background(), []model.Message{
		model.TextMessage(model.RoleUser, "hi"),
	}); err == nil {
		t.Fatal("Extract on malformed JSON: expected error, got nil")
	}
}

func TestLLMExtractor_Extract_NilAdapter_异常(t *testing.T) {
	t.Parallel()
	ex := NewLLMExtractor(nil, "persona.a", "role.1")

	if _, err := ex.Extract(context.Background(), []model.Message{
		model.TextMessage(model.RoleUser, "hi"),
	}); !errors.Is(err, ErrNilAdapter) {
		t.Fatalf("Extract with nil adapter err = %v, want ErrNilAdapter", err)
	}
}

func TestLLMExtractor_Extract_Canceled_异常(t *testing.T) {
	t.Parallel()
	adapter := &mockAdapter{fn: func(ctx context.Context, _ model.Request, _ model.Options) (model.Response, error) {
		<-ctx.Done()
		return model.Response{}, ctx.Err()
	}}
	ex := NewLLMExtractor(adapter, "persona.a", "role.1")

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	if _, err := ex.Extract(ctx, []model.Message{
		model.TextMessage(model.RoleUser, "hi"),
	}); !errors.Is(err, context.Canceled) {
		t.Fatalf("Extract with canceled ctx err = %v, want context.Canceled", err)
	}
}

func TestLLMExtractor_ExtractAsync_异步_正常(t *testing.T) {
	t.Parallel()
	started := make(chan struct{})
	release := make(chan struct{})
	adapter := &mockAdapter{fn: func(_ context.Context, _ model.Request, _ model.Options) (model.Response, error) {
		close(started)
		<-release
		return model.Response{Content: `[{"content":"异步提取的记忆"}]`}, nil
	}}
	ex := NewLLMExtractor(adapter, "persona.a", "role.1")

	ch, err := ex.ExtractAsync(context.Background(), []model.Message{
		model.TextMessage(model.RoleUser, "hi"),
	})
	if err != nil {
		t.Fatalf("ExtractAsync: %v", err)
	}
	// ExtractAsync returned while the adapter is still blocked on release,
	// proving the extraction runs on a background goroutine.
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("extraction goroutine did not start in time")
	}
	close(release)

	res := <-ch
	if res.Err != nil {
		t.Fatalf("async result err = %v", res.Err)
	}
	if len(res.Items) != 1 || res.Items[0].Content != "异步提取的记忆" {
		t.Fatalf("async items = %+v, want one extracted item", res.Items)
	}
}

func TestLLMExtractor_ExtractAsync_Idempotent_正常(t *testing.T) {
	t.Parallel()
	adapter := &mockAdapter{fn: func(_ context.Context, _ model.Request, _ model.Options) (model.Response, error) {
		return model.Response{Content: `[{"content":"习惯：早起"}]`}, nil
	}}
	ex := NewLLMExtractor(adapter, "persona.a", "role.1")

	conv := []model.Message{model.TextMessage(model.RoleUser, "hi")}

	first, err := ex.ExtractAsync(context.Background(), conv)
	if err != nil {
		t.Fatalf("ExtractAsync first: %v", err)
	}
	if res := <-first; res.Err != nil || len(res.Items) != 1 {
		t.Fatalf("first async result = %+v, want 1 item", res)
	}

	// A second async extraction of the same memory deduplicates to zero.
	second, err := ex.ExtractAsync(context.Background(), conv)
	if err != nil {
		t.Fatalf("ExtractAsync second: %v", err)
	}
	if res := <-second; res.Err != nil || len(res.Items) != 0 {
		t.Fatalf("second async result = %+v, want 0 items (idempotent)", res)
	}
}

func TestLLMExtractor_ExtractAsync_NilAdapter_异常(t *testing.T) {
	t.Parallel()
	ex := NewLLMExtractor(nil, "persona.a", "role.1")

	if _, err := ex.ExtractAsync(context.Background(), []model.Message{
		model.TextMessage(model.RoleUser, "hi"),
	}); !errors.Is(err, ErrNilAdapter) {
		t.Fatalf("ExtractAsync with nil adapter err = %v, want ErrNilAdapter", err)
	}
}