210 lines
7.8 KiB
Go
210 lines
7.8 KiB
Go
package vikingmemory
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"data-terminal-cloud/internal/outbox"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
const SyncEvent = "conversation.viking_sync_requested"
|
|
const eligibleSQL = `SELECT EXISTS(SELECT 1 FROM subject_profile sp WHERE sp.id=$2::uuid AND sp.tenant_id=$1::uuid AND sp.status='active'
|
|
AND NOT EXISTS(SELECT 1 FROM privacy_request pr WHERE pr.tenant_id=sp.tenant_id AND pr.subject_id=sp.id AND pr.request_type='DELETE')
|
|
AND EXISTS(SELECT 1 FROM consent_record cr WHERE cr.tenant_id=sp.tenant_id AND cr.subject_id=sp.id AND cr.consent_type='personalized_memory' AND cr.withdrawn_at IS NULL AND (cr.expires_at IS NULL OR cr.expires_at>now()))
|
|
AND (NOT EXISTS(SELECT 1 FROM consent_record cr WHERE cr.tenant_id=sp.tenant_id AND cr.subject_id=sp.id AND cr.consent_type='transcript_storage')
|
|
OR EXISTS(SELECT 1 FROM consent_record cr WHERE cr.tenant_id=sp.tenant_id AND cr.subject_id=sp.id AND cr.consent_type='transcript_storage' AND cr.withdrawn_at IS NULL AND (cr.expires_at IS NULL OR cr.expires_at>now()))))`
|
|
|
|
type Service struct {
|
|
pool *pgxpool.Pool
|
|
client *Client
|
|
}
|
|
|
|
func NewService(pool *pgxpool.Pool, client *Client) *Service {
|
|
return &Service{pool: pool, client: client}
|
|
}
|
|
func (s *Service) Eligible(ctx context.Context, tenant, subject string) (bool, error) {
|
|
var ok bool
|
|
err := s.pool.QueryRow(ctx, eligibleSQL, tenant, subject).Scan(&ok)
|
|
return ok, err
|
|
}
|
|
func (s *Service) Search(ctx context.Context, tenant, subject, query string, limit int) ([]Item, error) {
|
|
ok, err := s.Eligible(ctx, tenant, subject)
|
|
if err != nil || !ok {
|
|
return []Item{}, err
|
|
}
|
|
items, err := s.client.Search(ctx, UserID(tenant, subject), query, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// Recheck after network I/O so a concurrent withdrawal cannot expose results.
|
|
ok, err = s.Eligible(ctx, tenant, subject)
|
|
if err != nil || !ok {
|
|
return []Item{}, err
|
|
}
|
|
return items, nil
|
|
}
|
|
func (s *Service) PrepareForDevice(ctx context.Context, tenant, device string) (string, string, error) {
|
|
var subject string
|
|
err := s.pool.QueryRow(ctx, `SELECT b.subject_id::text FROM device_binding b JOIN device d ON d.id=b.device_id AND d.tenant_id=b.tenant_id WHERE b.device_id=$1::uuid AND b.tenant_id=$2::uuid AND b.unbound_at IS NULL AND d.status='registered'`, device, tenant).Scan(&subject)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return "", "", nil
|
|
}
|
|
if err != nil {
|
|
return "", "", err
|
|
}
|
|
// Empty query retrieves recent memories and the user profile at session start.
|
|
items, err := s.Search(ctx, tenant, subject, "", 4)
|
|
if err != nil {
|
|
return "", "", err
|
|
}
|
|
parts := []string{}
|
|
for _, item := range items {
|
|
parts = append(parts, item.Content)
|
|
}
|
|
return subject, strings.Join(parts, "\n"), nil
|
|
}
|
|
func (s *Service) Enqueue(ctx context.Context, event outbox.Event) error {
|
|
tx, err := s.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(context.Background())
|
|
var tenant string
|
|
var generation int64
|
|
err = tx.QueryRow(ctx, `SELECT c.tenant_id::text,COALESCE(max(ct.sequence_no),0) FROM conversation c LEFT JOIN conversation_turn ct ON ct.conversation_id=c.id WHERE c.id=$1::uuid AND c.status IN ('COMPLETED','FAILED') GROUP BY c.tenant_id`, event.AggregateID).Scan(&tenant, &generation)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if generation == 0 {
|
|
return nil
|
|
}
|
|
payload, _ := json.Marshal(map[string]int64{"generation_version": generation})
|
|
key := fmt.Sprintf("viking-v1:%s:%d", event.AggregateID, generation)
|
|
_, err = outbox.NewStore(s.pool).Enqueue(ctx, tx, outbox.NewEvent{TenantID: &tenant, AggregateType: "conversation", AggregateID: event.AggregateID, EventType: SyncEvent, DeduplicationKey: &key, Payload: payload})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (s *Service) Handle(ctx context.Context, event outbox.Event) error {
|
|
var tenant, subject, publicID string
|
|
var when time.Time
|
|
var frozen []byte
|
|
err := s.pool.QueryRow(ctx, `SELECT c.tenant_id::text,c.subject_id::text,c.public_id::text,COALESCE(c.started_at,c.created_at),cs.effective_consents FROM conversation c JOIN consent_snapshot cs ON cs.id=c.consent_snapshot_id AND cs.tenant_id=c.tenant_id WHERE c.id=$1::uuid AND c.status IN ('COMPLETED','FAILED')`, event.AggregateID).Scan(&tenant, &subject, &publicID, &when, &frozen)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
consents := map[string]bool{}
|
|
if json.Unmarshal(frozen, &consents) != nil || !consents["personalized_memory"] || !consents["transcript_storage"] {
|
|
return nil
|
|
}
|
|
ok, err := s.Eligible(ctx, tenant, subject)
|
|
if err != nil || !ok {
|
|
return err
|
|
}
|
|
// Serialize external writes and erasure per user, including retries from another worker.
|
|
conn, err := s.pool.Acquire(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer conn.Release()
|
|
user := UserID(tenant, subject)
|
|
if _, err = conn.Exec(ctx, `SELECT pg_advisory_lock(hashtextextended($1,0))`, user); err != nil {
|
|
return err
|
|
}
|
|
defer conn.Exec(context.Background(), `SELECT pg_advisory_unlock(hashtextextended($1,0))`, user)
|
|
ok, err = s.Eligible(ctx, tenant, subject)
|
|
if err != nil || !ok {
|
|
return err
|
|
}
|
|
rows, err := conn.Query(ctx, `SELECT speaker,final_text,start_at FROM conversation_turn WHERE conversation_id=$1::uuid ORDER BY sequence_no,id`, event.AggregateID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
messages := []Message{}
|
|
hasUser, hasAssistant := false, false
|
|
for rows.Next() {
|
|
var speaker, text string
|
|
var at time.Time
|
|
if err = rows.Scan(&speaker, &text, &at); err != nil {
|
|
rows.Close()
|
|
return err
|
|
}
|
|
role, id := "", user
|
|
if speaker == "subject" {
|
|
role = "user"
|
|
hasUser = true
|
|
} else if speaker == "agent" {
|
|
role = "assistant"
|
|
id = AssistantID
|
|
hasAssistant = true
|
|
}
|
|
if role != "" && strings.TrimSpace(text) != "" {
|
|
messages = append(messages, Message{Role: role, Content: text, RoleID: id, Time: at.UnixMilli()})
|
|
}
|
|
}
|
|
rows.Close()
|
|
if err = rows.Err(); err != nil {
|
|
return err
|
|
}
|
|
if !hasUser || !hasAssistant || len(messages) == 0 {
|
|
return nil
|
|
}
|
|
serialized, _ := json.Marshal(messages)
|
|
sum := sha256.Sum256(serialized)
|
|
digest := hex.EncodeToString(sum[:])
|
|
var already bool
|
|
err = conn.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM audit_event WHERE tenant_id=$1::uuid AND resource_id=$2::uuid AND action='memory.viking.session_sync' AND detail->>'digest'=$3 AND detail->>'collection'=$4)`, tenant, publicID, digest, s.client.options.ResourceID).Scan(&already)
|
|
if err != nil || already {
|
|
return err
|
|
}
|
|
for start := 0; start < len(messages); start += 100 {
|
|
end := min(start+100, len(messages))
|
|
if err = s.client.AddSession(ctx, SessionID(publicID, start/100), user, messages[start:end], when.UnixMilli()); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
ok, err = s.Eligible(ctx, tenant, subject)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !ok {
|
|
return s.client.DeleteUser(ctx, user)
|
|
}
|
|
detail, _ := json.Marshal(map[string]any{"digest": digest, "collection": s.client.options.ResourceID, "message_count": len(messages), "user_id": user})
|
|
_, err = conn.Exec(ctx, `INSERT INTO audit_event(tenant_id,actor_type,action,resource_type,resource_id,result,detail) VALUES($1::uuid,'system','memory.viking.session_sync','conversation',$2::uuid,'success',$3::jsonb)`, tenant, publicID, detail)
|
|
return err
|
|
}
|
|
|
|
func (s *Service) DeleteSubject(ctx context.Context, tenant, subject string) error {
|
|
if tenant == "" || subject == "" {
|
|
return errors.New("Viking erasure requires tenant and subject")
|
|
}
|
|
conn, err := s.pool.Acquire(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer conn.Release()
|
|
user := UserID(tenant, subject)
|
|
if _, err = conn.Exec(ctx, `SELECT pg_advisory_lock(hashtextextended($1,0))`, user); err != nil {
|
|
return err
|
|
}
|
|
defer conn.Exec(context.Background(), `SELECT pg_advisory_unlock(hashtextextended($1,0))`, user)
|
|
return s.client.DeleteUser(ctx, user)
|
|
}
|