chore: establish local platform baseline
This commit is contained in:
@@ -0,0 +1,136 @@
|
||||
package memoryindex
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"data-terminal-cloud/internal/embeddingprovider"
|
||||
"data-terminal-cloud/internal/outbox"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
var ErrInvalidIndexEvent = errors.New("invalid memory index event")
|
||||
var ErrMemoryChangedDuringIndex = errors.New("memory changed while embedding was generated")
|
||||
|
||||
type Processor struct {
|
||||
pool *pgxpool.Pool
|
||||
provider embeddingprovider.Provider
|
||||
}
|
||||
|
||||
type indexSnapshot struct {
|
||||
content, status, currentVersionID, sensitivity string
|
||||
consentEffective bool
|
||||
}
|
||||
|
||||
func NewProcessor(pool *pgxpool.Pool, providers ...embeddingprovider.Provider) (*Processor, error) {
|
||||
if pool == nil {
|
||||
return nil, errors.New("database pool is required")
|
||||
}
|
||||
var provider embeddingprovider.Provider = embeddingprovider.Local{}
|
||||
if len(providers) > 1 || (len(providers) == 1 && providers[0] == nil) {
|
||||
return nil, errors.New("one valid embedding provider is required")
|
||||
}
|
||||
if len(providers) == 1 {
|
||||
provider = providers[0]
|
||||
}
|
||||
return &Processor{pool: pool, provider: provider}, nil
|
||||
}
|
||||
|
||||
func (p *Processor) Handle(ctx context.Context, event outbox.Event) error {
|
||||
var payload struct {
|
||||
MemoryVersionID string `json:"memory_version_id"`
|
||||
}
|
||||
if err := json.Unmarshal(event.Payload, &payload); err != nil || strings.TrimSpace(payload.MemoryVersionID) == "" {
|
||||
return ErrInvalidIndexEvent
|
||||
}
|
||||
|
||||
before, err := loadIndexSnapshot(ctx, p.pool, payload.MemoryVersionID, false)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !eligible(before, payload.MemoryVersionID) {
|
||||
return p.markStale(ctx, payload.MemoryVersionID, "memory is not currently eligible")
|
||||
}
|
||||
result, err := p.provider.Embed(ctx, before.content)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
literal, err := embeddingprovider.Literal(result.Vector)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Embedding is generated without an open transaction. Only the short
|
||||
// revalidation and vector update below hold memory/index row locks.
|
||||
tx, err := p.pool.BeginTx(ctx, pgx.TxOptions{})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback(context.Background())
|
||||
current, err := loadIndexSnapshot(ctx, tx, payload.MemoryVersionID, true)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !eligible(current, payload.MemoryVersionID) {
|
||||
if _, err = tx.Exec(ctx, `UPDATE memory_index_entry SET embedding=NULL,status='STALE',last_attempt_at=now(),last_error='memory is not currently eligible',updated_at=now() WHERE memory_version_id=$1::uuid AND index_type='pgvector'`, payload.MemoryVersionID); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
if current.content != before.content {
|
||||
return outbox.Retryable(ErrMemoryChangedDuringIndex)
|
||||
}
|
||||
if _, err = tx.Exec(ctx, `
|
||||
UPDATE memory_index_entry SET embedding=$2::vector,status='ACTIVE',provider_ref=$3,
|
||||
indexed_at=now(),last_attempt_at=now(),last_error=NULL,updated_at=now()
|
||||
WHERE memory_version_id=$1::uuid AND index_type='pgvector'`,
|
||||
payload.MemoryVersionID, literal, result.ProviderRef); err != nil {
|
||||
return fmt.Errorf("activate memory index entry: %w", err)
|
||||
}
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
type indexQueryer interface {
|
||||
QueryRow(context.Context, string, ...any) pgx.Row
|
||||
}
|
||||
|
||||
func loadIndexSnapshot(ctx context.Context, q indexQueryer, memoryVersionID string, lock bool) (indexSnapshot, error) {
|
||||
lockClause := ""
|
||||
if lock {
|
||||
lockClause = " FOR UPDATE OF mi,cm,sp"
|
||||
}
|
||||
var snapshot indexSnapshot
|
||||
err := q.QueryRow(ctx, `
|
||||
SELECT mv.content,cm.status,COALESCE(cm.current_version_id::text,''),cm.sensitivity,
|
||||
EXISTS(SELECT 1 FROM consent_record cr
|
||||
WHERE cr.tenant_id=cm.tenant_id AND cr.subject_id=cm.subject_id
|
||||
AND cr.consent_type='personalized_memory' AND cr.withdrawn_at IS NULL
|
||||
AND (cr.expires_at IS NULL OR cr.expires_at>now()))
|
||||
FROM memory_index_entry mi
|
||||
JOIN memory_version mv ON mv.id=mi.memory_version_id
|
||||
JOIN curated_memory cm ON cm.id=mi.memory_id
|
||||
JOIN subject_profile sp ON sp.id=cm.subject_id AND sp.tenant_id=cm.tenant_id AND sp.status IN ('active','suspended')
|
||||
WHERE mi.memory_version_id=$1::uuid AND mi.index_type='pgvector'`+lockClause, memoryVersionID).
|
||||
Scan(&snapshot.content, &snapshot.status, &snapshot.currentVersionID, &snapshot.sensitivity, &snapshot.consentEffective)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return indexSnapshot{}, ErrInvalidIndexEvent
|
||||
}
|
||||
if err != nil {
|
||||
return indexSnapshot{}, fmt.Errorf("load memory index entry: %w", err)
|
||||
}
|
||||
return snapshot, nil
|
||||
}
|
||||
|
||||
func eligible(snapshot indexSnapshot, memoryVersionID string) bool {
|
||||
return snapshot.status == "ACTIVE" && snapshot.currentVersionID == memoryVersionID && snapshot.consentEffective && snapshot.sensitivity == "normal"
|
||||
}
|
||||
|
||||
func (p *Processor) markStale(ctx context.Context, memoryVersionID, reason string) error {
|
||||
_, err := p.pool.Exec(ctx, `UPDATE memory_index_entry SET embedding=NULL,status='STALE',last_attempt_at=now(),last_error=$2,updated_at=now() WHERE memory_version_id=$1::uuid AND index_type='pgvector'`, memoryVersionID, reason)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package memoryindex
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"data-terminal-cloud/internal/developmentembedding"
|
||||
"data-terminal-cloud/internal/embeddingprovider"
|
||||
"data-terminal-cloud/internal/outbox"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type indexLockProbe struct {
|
||||
pool *pgxpool.Pool
|
||||
versionID string
|
||||
}
|
||||
|
||||
func (p indexLockProbe) Embed(ctx context.Context, content string) (embeddingprovider.Result, error) {
|
||||
updateContext, cancel := context.WithTimeout(ctx, time.Second)
|
||||
defer cancel()
|
||||
if _, err := p.pool.Exec(updateContext, `UPDATE memory_index_entry SET updated_at=updated_at WHERE memory_version_id=$1::uuid`, p.versionID); err != nil {
|
||||
return embeddingprovider.Result{}, err
|
||||
}
|
||||
return embeddingprovider.Result{Vector: developmentembedding.Vector(content), ProviderRef: "integration-lock-probe"}, nil
|
||||
}
|
||||
|
||||
func TestProcessorActivatesAndInvalidatesMemoryIndex(t *testing.T) {
|
||||
databaseURL := os.Getenv("TEST_DATABASE_URL")
|
||||
if databaseURL == "" {
|
||||
t.Skip("TEST_DATABASE_URL is not set")
|
||||
}
|
||||
ctx := context.Background()
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer pool.Close()
|
||||
if _, err := NewProcessor(nil); err == nil {
|
||||
t.Fatal("nil database pool was accepted")
|
||||
}
|
||||
processor, err := NewProcessor(pool)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
memoryID, memoryPublicID := uuid.NewString(), uuid.NewString()
|
||||
versionID, versionPublicID := uuid.NewString(), uuid.NewString()
|
||||
defer func() {
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM memory_index_entry WHERE memory_id=$1::uuid`, memoryID)
|
||||
_, _ = pool.Exec(ctx, `UPDATE curated_memory SET current_version_id=NULL WHERE id=$1::uuid`, memoryID)
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM memory_version WHERE memory_id=$1::uuid`, memoryID)
|
||||
_, _ = pool.Exec(ctx, `DELETE FROM curated_memory WHERE id=$1::uuid`, memoryID)
|
||||
}()
|
||||
if _, err := pool.Exec(ctx, `
|
||||
INSERT INTO curated_memory(id,public_id,tenant_id,subject_id,status,sensitivity,valid_from)
|
||||
VALUES($1::uuid,$2::uuid,'10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401','ACTIVE','normal',now())`, memoryID, memoryPublicID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, `
|
||||
INSERT INTO memory_version(id,public_id,memory_id,version_no,content,source_refs,confirmed_by,change_reason,content_digest)
|
||||
VALUES($1::uuid,$2::uuid,$3::uuid,1,$4,'[{"type":"integration"}]','10000000-0000-0000-0000-000000000101','合成索引测试',$5)`,
|
||||
versionID, versionPublicID, memoryID, "长者在河边小学教书。", strings.Repeat("a", 64)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, `UPDATE curated_memory SET current_version_id=$1::uuid WHERE id=$2::uuid`, versionID, memoryID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, `
|
||||
INSERT INTO memory_index_entry(tenant_id,subject_id,memory_id,memory_version_id,index_type,status)
|
||||
VALUES('10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',$1::uuid,$2::uuid,'pgvector','PENDING')`, memoryID, versionID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := processor.Handle(ctx, outbox.Event{Payload: []byte(`{}`)}); !errors.Is(err, ErrInvalidIndexEvent) {
|
||||
t.Fatalf("invalid event error=%v", err)
|
||||
}
|
||||
missingID := uuid.NewString()
|
||||
if err := processor.Handle(ctx, outbox.Event{Payload: []byte(`{"memory_version_id":"` + missingID + `"}`)}); !errors.Is(err, ErrInvalidIndexEvent) {
|
||||
t.Fatalf("missing version error=%v", err)
|
||||
}
|
||||
event := outbox.Event{Payload: []byte(`{"memory_version_id":"` + versionID + `"}`)}
|
||||
processor, err = NewProcessor(pool, indexLockProbe{pool: pool, versionID: versionID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := processor.Handle(ctx, event); err != nil {
|
||||
t.Fatalf("activate index: %v", err)
|
||||
}
|
||||
var status, providerRef string
|
||||
var dimensions int
|
||||
if err := pool.QueryRow(ctx, `SELECT status,provider_ref,vector_dims(embedding) FROM memory_index_entry WHERE memory_version_id=$1::uuid`, versionID).Scan(&status, &providerRef, &dimensions); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if status != "ACTIVE" || providerRef != "integration-lock-probe" || dimensions != 1536 {
|
||||
t.Fatalf("index status=%s provider=%s dimensions=%d", status, providerRef, dimensions)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, `UPDATE curated_memory SET status='DISABLED' WHERE id=$1::uuid`, memoryID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, `UPDATE memory_index_entry SET status='PENDING' WHERE memory_version_id=$1::uuid`, versionID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := processor.Handle(ctx, event); err != nil {
|
||||
t.Fatalf("invalidate index: %v", err)
|
||||
}
|
||||
var embeddingIsNull bool
|
||||
if err := pool.QueryRow(ctx, `SELECT status,embedding IS NULL FROM memory_index_entry WHERE memory_version_id=$1::uuid`, versionID).Scan(&status, &embeddingIsNull); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if status != "STALE" || !embeddingIsNull {
|
||||
t.Fatalf("ineligible index status=%s embedding_null=%v", status, embeddingIsNull)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package memoryindex
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"data-terminal-cloud/internal/developmentembedding"
|
||||
)
|
||||
|
||||
func TestDevelopmentEmbeddingIsDeterministicAndNormalized(t *testing.T) {
|
||||
first := developmentembedding.Vector("长者谈到小时候在河边散步。")
|
||||
second := developmentembedding.Vector("长者谈到小时候在河边散步。")
|
||||
other := developmentembedding.Vector("完全不同的记忆内容。")
|
||||
if len(first) != developmentembedding.Dimensions || len(second) != len(first) {
|
||||
t.Fatalf("embedding dimensions = %d, want %d", len(first), developmentembedding.Dimensions)
|
||||
}
|
||||
var squared float64
|
||||
var different bool
|
||||
for index := range first {
|
||||
if first[index] != second[index] {
|
||||
t.Fatalf("embedding differs at %d", index)
|
||||
}
|
||||
if first[index] != other[index] {
|
||||
different = true
|
||||
}
|
||||
squared += first[index] * first[index]
|
||||
}
|
||||
if !different {
|
||||
t.Fatal("different content produced the same embedding")
|
||||
}
|
||||
if math.Abs(math.Sqrt(squared)-1) > 1e-9 {
|
||||
t.Fatalf("embedding norm = %.12f, want 1", math.Sqrt(squared))
|
||||
}
|
||||
literal := developmentembedding.Literal(first)
|
||||
if !strings.HasPrefix(literal, "[") || !strings.HasSuffix(literal, "]") || strings.Count(literal, ",") != developmentembedding.Dimensions-1 {
|
||||
t.Fatal("vector literal has an invalid shape")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user