story/internal/memorytool/service.go

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())
	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 := trimRunes(memory.Content, 240)
		if totalRunes+utf8.RuneCountInString(content) > 800 {
			break
		}
		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]) + "…"
}