This commit is contained in:
Codex Local
@@ -0,0 +1,161 @@
|
||||
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]) + "…"
|
||||
}
|
||||
Reference in New Issue
Block a user