162 lines
6.7 KiB
Go
162 lines
6.7 KiB
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"
|
|
|
|
"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 {
|
|
pool *pgxpool.Pool
|
|
embedder embeddingprovider.Provider
|
|
audit *audit.Repository
|
|
now func() time.Time
|
|
}
|
|
|
|
type conversationScope struct {
|
|
internalID, publicID, tenantID, subjectID, taskRef, engineMode, 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) 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[:])
|
|
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,rt.provider_task_ref,c.engine_mode,c.status
|
|
FROM conversation c JOIN rtc_task rt ON rt.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.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" {
|
|
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]) + "…"
|
|
}
|