package memorytool import ( "context" "crypto/sha256" "encoding/hex" "encoding/json" "errors" "fmt" "strings" "time" "unicode/utf8" "data-terminal-cloud/internal/audit" "data-terminal-cloud/internal/embeddingprovider" "data-terminal-cloud/internal/retrieval" "data-terminal-cloud/internal/vikingmemory" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" ) var ( ErrConversationUnavailable = errors.New("memory tool conversation is unavailable") ErrEngineModeUnsupported = errors.New("memory tool is unavailable for this engine mode") ) type Result struct { MemoryID string `json:"memory_id"` MemoryVersionID string `json:"memory_version_id"` VersionNo int64 `json:"version_no"` Content string `json:"content"` SourceRefs json.RawMessage `json:"source_refs"` Relevance float64 `json:"relevance"` } type Response struct { Items []Result `json:"items"` Degraded bool `json:"degraded"` } type Service struct { viking *vikingmemory.Service pool *pgxpool.Pool embedder embeddingprovider.Provider audit *audit.Repository now func() time.Time } type conversationScope struct { internalID, publicID, tenantID, subjectID, taskRef, engineMode, runtimeType, status string } func NewService(pool *pgxpool.Pool, embedder embeddingprovider.Provider, auditRepository *audit.Repository) (*Service, error) { if pool == nil || embedder == nil || auditRepository == nil { return nil, errors.New("memory tool database, embedding provider and audit repository are required") } return &Service{pool: pool, embedder: embedder, audit: auditRepository, now: time.Now}, nil } func (s *Service) ConfigureVikingMemory(memory *vikingmemory.Service) { s.viking = memory } func (s *Service) Search(ctx context.Context, claims Claims, query string, limit int, requestID string) (Response, error) { started := s.now().UTC() query = strings.TrimSpace(query) if query == "" || utf8.RuneCountInString(query) > 500 || limit < 1 || limit > 5 { return Response{}, errors.New("memory query and limit are invalid") } scope, err := s.resolveScope(ctx, claims) if err != nil { return Response{}, err } digest := sha256.Sum256([]byte(query)) queryDigest := hex.EncodeToString(digest[:]) if s.viking != nil { searchContext, cancel := context.WithTimeout(ctx, 3*time.Second) memories, searchErr := s.viking.Search(searchContext, scope.tenantID, scope.subjectID, query, limit) cancel() if searchErr == nil && len(memories) > 0 { items := make([]Result, 0, len(memories)) for _, item := range memories { refs, _ := json.Marshal([]map[string]string{{"provider": "viking", "session_id": item.SessionID}}) items = append(items, Result{MemoryID: item.ID, Content: item.Content, SourceRefs: refs, Relevance: item.Score}) } s.recordCall(ctx, scope, queryDigest, len(items), "SUCCEEDED", "", started) return Response{Items: items}, nil } } embedding, err := s.embedder.Embed(ctx, query) if err != nil { s.recordCall(ctx, scope, queryDigest, 0, "FAILED", "EMBEDDING_UNAVAILABLE", started) return Response{Items: []Result{}, Degraded: true}, nil } vector, err := embeddingprovider.Literal(embedding.Vector) if err != nil { s.recordCall(ctx, scope, queryDigest, 0, "FAILED", "EMBEDDING_INVALID", started) return Response{Items: []Result{}, Degraded: true}, nil } tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{AccessMode: pgx.ReadWrite}) if err != nil { return Response{}, err } defer tx.Rollback(context.Background()) // This legacy tool has no authenticated current USER utterance binding. // It cannot widen access to RESPOND_ONLY just because the model invoked it. memories, err := retrieval.SelectMemoryForContextVector(ctx, tx, scope.tenantID, scope.subjectID, vector, limit) if err != nil { return Response{}, err } items := make([]Result, 0, len(memories)) totalRunes := 0 for _, memory := range memories { content := strings.TrimSpace(memory.Content) if totalRunes+utf8.RuneCountInString(content) > 800 { continue } totalRunes += utf8.RuneCountInString(content) items = append(items, Result{MemoryID: memory.MemoryID, MemoryVersionID: memory.MemoryVersionID, VersionNo: memory.VersionNo, Content: content, SourceRefs: memory.SourceRefs, Relevance: memory.Relevance}) if _, err := tx.Exec(ctx, `INSERT INTO conversation_memory_use(tenant_id,subject_id,conversation_id,memory_id,memory_version_id,use_type,relevance,source_refs) SELECT $1::uuid,$2::uuid,$3::uuid,cm.id,mv.id,'in_call_tool',$6,$7::jsonb FROM curated_memory cm JOIN memory_version mv ON mv.memory_id=cm.id WHERE cm.public_id=$4::uuid AND mv.public_id=$5::uuid ON CONFLICT(conversation_id,memory_version_id,use_type) DO UPDATE SET relevance=GREATEST(conversation_memory_use.relevance,EXCLUDED.relevance),used_at=now()`, scope.tenantID, scope.subjectID, scope.internalID, memory.MemoryID, memory.MemoryVersionID, memory.Relevance, memory.SourceRefs); err != nil { return Response{}, err } } status := "SUCCEEDED" if len(items) == 0 { status = "EMPTY" } latency := int(s.now().UTC().Sub(started).Milliseconds()) if latency < 0 { latency = 0 } if _, err := tx.Exec(ctx, `INSERT INTO memory_tool_call(tenant_id,subject_id,conversation_id,query_digest,result_count,status,latency_ms,engine_mode) VALUES($1::uuid,$2::uuid,$3::uuid,$4,$5,$6,$7,$8)`, scope.tenantID, scope.subjectID, scope.internalID, queryDigest, len(items), status, latency, scope.engineMode); err != nil { return Response{}, err } if err := s.audit.Record(ctx, tx, audit.Event{TenantID: &scope.tenantID, ActorType: "service", Action: "memory.tool.search", ResourceType: "conversation", ResourceID: &scope.publicID, RequestID: requestID, Result: "success", Detail: map[string]any{"result_count": len(items), "engine_mode": scope.engineMode}}); err != nil { return Response{}, err } if err := tx.Commit(ctx); err != nil { return Response{}, err } return Response{Items: items, Degraded: false}, nil } func (s *Service) resolveScope(ctx context.Context, claims Claims) (conversationScope, error) { var scope conversationScope err := s.pool.QueryRow(ctx, `SELECT c.id::text,c.public_id::text,c.tenant_id::text,c.subject_id::text, COALESCE(rt.provider_task_ref,hs.provider_session_ref),c.engine_mode,c.runtime_type,c.status FROM conversation c LEFT JOIN rtc_task rt ON rt.conversation_id=c.id LEFT JOIN hardware_agent_session hs ON hs.conversation_id=c.id WHERE c.public_id=$1::uuid`, claims.ConversationID). Scan(&scope.internalID, &scope.publicID, &scope.tenantID, &scope.subjectID, &scope.taskRef, &scope.engineMode, &scope.runtimeType, &scope.status) if errors.Is(err, pgx.ErrNoRows) || scope.taskRef != claims.TaskRef { return conversationScope{}, ErrConversationUnavailable } if err != nil { return conversationScope{}, fmt.Errorf("resolve memory tool conversation: %w", err) } if scope.engineMode != "s2s_hybrid" && scope.runtimeType != "volc_hardware_agent" { return conversationScope{}, ErrEngineModeUnsupported } if scope.status != "ACTIVE" && scope.status != "STARTING_AGENT" { return conversationScope{}, ErrConversationUnavailable } return scope, nil } func (s *Service) recordCall(ctx context.Context, scope conversationScope, queryDigest string, resultCount int, status, errorCode string, started time.Time) { latency := int(s.now().UTC().Sub(started).Milliseconds()) if latency < 0 { latency = 0 } _, _ = s.pool.Exec(ctx, `INSERT INTO memory_tool_call(tenant_id,subject_id,conversation_id,query_digest,result_count,status,error_code,latency_ms,engine_mode) VALUES($1::uuid,$2::uuid,$3::uuid,$4,$5,$6,$7,$8,$9)`, scope.tenantID, scope.subjectID, scope.internalID, queryDigest, resultCount, status, errorCode, latency, scope.engineMode) } func trimRunes(value string, limit int) string { runes := []rune(strings.TrimSpace(value)) if len(runes) <= limit { return string(runes) } return string(runes[:limit]) + "…" }