85 lines
5.6 KiB
Go
85 lines
5.6 KiB
Go
package memorytool
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
|
|
"data-terminal-cloud/internal/audit"
|
|
"data-terminal-cloud/internal/developmentembedding"
|
|
"data-terminal-cloud/internal/embeddingprovider"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
func TestSearchDerivesScopeAndRecordsActualMemoryUse(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()
|
|
contextID, conversationID, conversationPublicID := uuid.NewString(), uuid.NewString(), uuid.NewString()
|
|
taskRef := "memory-tool-task-" + uuid.NewString()
|
|
memoryID, memoryPublicID, versionID, versionPublicID := uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString()
|
|
defer func() {
|
|
_, _ = pool.Exec(ctx, `DELETE FROM audit_event WHERE request_id LIKE 'memory-tool-integration-%'`)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM memory_tool_call WHERE conversation_id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM conversation_memory_use WHERE conversation_id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM rtc_task WHERE conversation_id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM conversation WHERE id=$1::uuid`, conversationID)
|
|
_, _ = pool.Exec(ctx, `DELETE FROM context_snapshot WHERE id=$1::uuid`, contextID)
|
|
_, _ = 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 context_snapshot(id,tenant_id,subject_id,profile_version,consent_snapshot_id,memory_versions,prompt_bundle_version,engine_mode,content_digest,sanitized_summary)
|
|
VALUES($1::uuid,'10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',1,'10000000-0000-0000-0000-000000000702','[]','memory-tool-test','s2s_hybrid',$2,'')`, contextID, strings.Repeat("1", 64)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `INSERT INTO conversation(id,public_id,tenant_id,household_id,subject_id,device_id,consent_snapshot_id,context_snapshot_id,provider_task_ref,engine_mode,retention_mode,status,allocation_expires_at,started_at)
|
|
VALUES($1::uuid,$2::uuid,'10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000201','10000000-0000-0000-0000-000000000401','10000000-0000-0000-0000-000000000501','10000000-0000-0000-0000-000000000702',$3::uuid,$4,'s2s_hybrid','transcript_only','ACTIVE',now()+interval '1 hour',now())`, conversationID, conversationPublicID, contextID, taskRef); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := pool.Exec(ctx, `INSERT INTO rtc_task(conversation_id,provider,provider_task_ref,room_id,device_rtc_user_id,agent_rtc_user_id,status) VALUES($1::uuid,'mock',$2,$3,'device-user','agent-user','ACTIVE')`, conversationID, taskRef, "memory-room-"+conversationID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
content := "长者曾经在河边小学教书。"
|
|
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,content_digest) VALUES($1::uuid,$2::uuid,$3::uuid,1,$4,'[{"type":"story","story_id":"synthetic"}]','10000000-0000-0000-0000-000000000101',$5)`, versionID, versionPublicID, memoryID, content, strings.Repeat("2", 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,embedding,status,provider_ref) VALUES('10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',$1::uuid,$2::uuid,'pgvector',$3::vector,'ACTIVE','integration')`, memoryID, versionID, developmentembedding.Literal(developmentembedding.Vector(content))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
service, err := NewService(pool, embeddingprovider.Local{}, audit.NewRepository())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
claims := Claims{ConversationID: conversationPublicID, TaskRef: taskRef}
|
|
result, err := service.Search(ctx, claims, "河边小学", 3, "memory-tool-integration-search")
|
|
if err != nil || result.Degraded || len(result.Items) != 1 || result.Items[0].MemoryID != memoryPublicID || result.Items[0].MemoryVersionID != versionPublicID {
|
|
t.Fatalf("result=%+v err=%v", result, err)
|
|
}
|
|
var calls, uses int
|
|
if err := pool.QueryRow(ctx, `SELECT (SELECT count(*) FROM memory_tool_call WHERE conversation_id=$1::uuid AND status='SUCCEEDED'),(SELECT count(*) FROM conversation_memory_use WHERE conversation_id=$1::uuid AND use_type='in_call_tool')`, conversationID).Scan(&calls, &uses); err != nil || calls != 1 || uses != 1 {
|
|
t.Fatalf("calls=%d uses=%d err=%v", calls, uses, err)
|
|
}
|
|
if _, err := service.Search(ctx, Claims{ConversationID: conversationPublicID, TaskRef: "wrong-task"}, "河边小学", 3, "memory-tool-integration-wrong-task"); !errors.Is(err, ErrConversationUnavailable) {
|
|
t.Fatalf("wrong task error=%v", err)
|
|
}
|
|
}
|