mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-08-06 16:24:22 +00:00
viettranx
e6896d2673
All 11 sites migrated (counted higher than scout's 7 CRITICAL because plan treats the file as a single commit unit): - UpsertRelation (agent, source, target — 3 sites in one INSERT) - DeleteRelation (agent, id) - ListRelations, ListAllRelations (agent + entity) - IngestExtraction (agent) - PruneByConfidence, Stats (agent) Phase 4 Step 4b.4 of agent identity hardening (TD-1).
383 lines
11 KiB
Go
383 lines
11 KiB
Go
package pg
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log/slog"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/nextlevelbuilder/goclaw/internal/store"
|
|
)
|
|
|
|
func (s *PGKnowledgeGraphStore) UpsertRelation(ctx context.Context, relation *store.Relation) error {
|
|
aid, err := parseUUID(relation.AgentID)
|
|
if err != nil {
|
|
return fmt.Errorf("kg upsert relation: agent: %w", err)
|
|
}
|
|
src, err := parseUUID(relation.SourceEntityID)
|
|
if err != nil {
|
|
return fmt.Errorf("kg upsert relation: source: %w", err)
|
|
}
|
|
tgt, err := parseUUID(relation.TargetEntityID)
|
|
if err != nil {
|
|
return fmt.Errorf("kg upsert relation: target: %w", err)
|
|
}
|
|
props, err := json.Marshal(relation.Properties)
|
|
if err != nil {
|
|
props = []byte("{}")
|
|
}
|
|
id := uuid.Must(uuid.NewV7())
|
|
now := time.Now()
|
|
tid := tenantIDForInsert(ctx)
|
|
_, err = s.db.ExecContext(ctx, `
|
|
INSERT INTO kg_relations
|
|
(id, agent_id, user_id, source_entity_id, relation_type, target_entity_id, confidence, properties, tenant_id, created_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
|
ON CONFLICT (agent_id, user_id, source_entity_id, relation_type, target_entity_id) DO UPDATE SET
|
|
confidence = EXCLUDED.confidence,
|
|
properties = EXCLUDED.properties,
|
|
tenant_id = EXCLUDED.tenant_id`,
|
|
id, aid, relation.UserID, src, relation.RelationType, tgt, relation.Confidence, props, tid, now,
|
|
)
|
|
return err
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) DeleteRelation(ctx context.Context, agentID, userID, relationID string) error {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return fmt.Errorf("kg delete relation: agent: %w", err)
|
|
}
|
|
rid, err := parseUUID(relationID)
|
|
if err != nil {
|
|
return fmt.Errorf("kg delete relation: id: %w", err)
|
|
}
|
|
if store.IsSharedKG(ctx) {
|
|
tc, tcArgs, _, err := scopeClause(ctx, 3)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = s.db.ExecContext(ctx,
|
|
`DELETE FROM kg_relations WHERE id = $1 AND agent_id = $2`+tc,
|
|
append([]any{rid, aid}, tcArgs...)...,
|
|
)
|
|
return err
|
|
}
|
|
tc, tcArgs, _, err := scopeClause(ctx, 4)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = s.db.ExecContext(ctx,
|
|
`DELETE FROM kg_relations WHERE id = $1 AND agent_id = $2 AND user_id = $3`+tc,
|
|
append([]any{rid, aid, userID}, tcArgs...)...,
|
|
)
|
|
return err
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) ListRelations(ctx context.Context, agentID, userID, entityID string) ([]store.Relation, error) {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("kg list relations: agent: %w", err)
|
|
}
|
|
eid, err := parseUUID(entityID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("kg list relations: entity: %w", err)
|
|
}
|
|
|
|
var q string
|
|
var args []any
|
|
if store.IsSharedKG(ctx) {
|
|
tc, tcArgs, _, err := scopeClause(ctx, 3)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
q = `SELECT id, agent_id, user_id, source_entity_id, relation_type, target_entity_id,
|
|
confidence, properties, created_at
|
|
FROM kg_relations
|
|
WHERE agent_id = $1 AND valid_until IS NULL
|
|
AND (source_entity_id = $2 OR target_entity_id = $2)` + tc + `
|
|
ORDER BY created_at DESC`
|
|
args = append([]any{aid, eid}, tcArgs...)
|
|
} else {
|
|
tc, tcArgs, _, err := scopeClause(ctx, 4)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
q = `SELECT id, agent_id, user_id, source_entity_id, relation_type, target_entity_id,
|
|
confidence, properties, created_at
|
|
FROM kg_relations
|
|
WHERE agent_id = $1 AND user_id = $2 AND valid_until IS NULL
|
|
AND (source_entity_id = $3 OR target_entity_id = $3)` + tc + `
|
|
ORDER BY created_at DESC`
|
|
args = append([]any{aid, userID, eid}, tcArgs...)
|
|
}
|
|
|
|
var rRows []relationRow
|
|
if err := pkgSqlxDB.SelectContext(ctx, &rRows, q, args...); err != nil {
|
|
return nil, err
|
|
}
|
|
result := make([]store.Relation, len(rRows))
|
|
for i := range rRows {
|
|
result[i] = rRows[i].toRelation()
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) ListAllRelations(ctx context.Context, agentID, userID string, limit int) ([]store.Relation, error) {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("kg list all relations: %w", err)
|
|
}
|
|
if limit <= 0 {
|
|
limit = 200
|
|
}
|
|
where := "agent_id = $1 AND valid_until IS NULL"
|
|
args := []any{aid}
|
|
idx := 2
|
|
if !store.IsSharedKG(ctx) && userID != "" {
|
|
where += fmt.Sprintf(" AND user_id = $%d", idx)
|
|
args = append(args, userID)
|
|
idx++
|
|
}
|
|
tc, tcArgs, _, err := scopeClause(ctx, idx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if tc != "" {
|
|
where += tc
|
|
args = append(args, tcArgs...)
|
|
idx++
|
|
}
|
|
args = append(args, limit)
|
|
q := fmt.Sprintf(`
|
|
SELECT id, agent_id, user_id, source_entity_id, relation_type, target_entity_id,
|
|
confidence, properties, created_at
|
|
FROM kg_relations WHERE %s
|
|
ORDER BY created_at DESC LIMIT $%d`, where, idx)
|
|
var rRows []relationRow
|
|
if err = pkgSqlxDB.SelectContext(ctx, &rRows, q, args...); err != nil {
|
|
return nil, err
|
|
}
|
|
result := make([]store.Relation, len(rRows))
|
|
for i := range rRows {
|
|
result[i] = rRows[i].toRelation()
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) IngestExtraction(ctx context.Context, agentID, userID string, entities []store.Entity, relations []store.Relation) ([]string, error) {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("kg ingest extraction: agent: %w", err)
|
|
}
|
|
tx, err := s.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer tx.Rollback() //nolint:errcheck
|
|
|
|
now := time.Now()
|
|
tid := tenantIDForInsert(ctx)
|
|
|
|
// Upsert entities and build external_id → DB UUID lookup for relations
|
|
extIDToUUID := make(map[string]uuid.UUID, len(entities))
|
|
for i := range entities {
|
|
e := &entities[i]
|
|
e.AgentID = agentID
|
|
e.UserID = userID
|
|
props, _ := json.Marshal(e.Properties)
|
|
id := uuid.Must(uuid.NewV7())
|
|
// Use RETURNING to get the actual ID (could be existing row on conflict)
|
|
var actualID uuid.UUID
|
|
if err := tx.QueryRowContext(ctx, `
|
|
INSERT INTO kg_entities
|
|
(id, agent_id, user_id, external_id, name, entity_type, description, properties, source_id, confidence, tenant_id, created_at, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $12)
|
|
ON CONFLICT (agent_id, user_id, external_id) DO UPDATE SET
|
|
name = EXCLUDED.name,
|
|
entity_type = EXCLUDED.entity_type,
|
|
description = EXCLUDED.description,
|
|
properties = EXCLUDED.properties,
|
|
source_id = EXCLUDED.source_id,
|
|
confidence = EXCLUDED.confidence,
|
|
tenant_id = EXCLUDED.tenant_id,
|
|
updated_at = EXCLUDED.updated_at
|
|
RETURNING id`,
|
|
id, aid, userID, e.ExternalID, e.Name, e.EntityType,
|
|
e.Description, props, e.SourceID, e.Confidence, tid, now,
|
|
).Scan(&actualID); err != nil {
|
|
return nil, err
|
|
}
|
|
extIDToUUID[e.ExternalID] = actualID
|
|
}
|
|
|
|
// Batch-generate embeddings for all upserted entities (fire-and-forget on error).
|
|
if s.embProvider != nil && len(extIDToUUID) > 0 {
|
|
texts := make([]string, 0, len(entities))
|
|
ids := make([]uuid.UUID, 0, len(entities))
|
|
for _, e := range entities {
|
|
texts = append(texts, e.Name+" "+e.Description)
|
|
ids = append(ids, extIDToUUID[e.ExternalID])
|
|
}
|
|
embeddings, embErr := s.embProvider.Embed(ctx, texts)
|
|
if embErr != nil {
|
|
slog.Warn("kg entity embedding batch failed", "error", embErr)
|
|
} else {
|
|
for i, emb := range embeddings {
|
|
if len(emb) == 0 {
|
|
continue
|
|
}
|
|
vecStr := vectorToString(emb)
|
|
if _, err := tx.ExecContext(ctx,
|
|
`UPDATE kg_entities SET embedding = $1::vector WHERE id = $2`,
|
|
vecStr, ids[i],
|
|
); err != nil {
|
|
slog.Warn("kg entity embedding update failed", "entity_id", ids[i], "error", err)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
for i := range relations {
|
|
r := &relations[i]
|
|
r.AgentID = agentID
|
|
r.UserID = userID
|
|
// Resolve external_id references to actual DB UUIDs
|
|
src, ok1 := extIDToUUID[r.SourceEntityID]
|
|
tgt, ok2 := extIDToUUID[r.TargetEntityID]
|
|
if !ok1 || !ok2 {
|
|
continue // skip relations referencing unknown entities
|
|
}
|
|
props, _ := json.Marshal(r.Properties)
|
|
id := uuid.Must(uuid.NewV7())
|
|
if _, err := tx.ExecContext(ctx, `
|
|
INSERT INTO kg_relations
|
|
(id, agent_id, user_id, source_entity_id, relation_type, target_entity_id, confidence, properties, tenant_id, created_at)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
|
ON CONFLICT (agent_id, user_id, source_entity_id, relation_type, target_entity_id) DO UPDATE SET
|
|
confidence = EXCLUDED.confidence,
|
|
properties = EXCLUDED.properties,
|
|
tenant_id = EXCLUDED.tenant_id`,
|
|
id, aid, userID, src, r.RelationType, tgt, r.Confidence, props, tid, now,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
if err := tx.Commit(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Collect upserted entity IDs for downstream processing (e.g. dedup)
|
|
entityIDs := make([]string, 0, len(extIDToUUID))
|
|
for _, uid := range extIDToUUID {
|
|
entityIDs = append(entityIDs, uid.String())
|
|
}
|
|
return entityIDs, nil
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) PruneByConfidence(ctx context.Context, agentID, userID string, minConfidence float64) (int, error) {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("kg prune: %w", err)
|
|
}
|
|
var res sql.Result
|
|
if store.IsSharedKG(ctx) {
|
|
tc, tcArgs, _, tcErr := scopeClause(ctx, 3)
|
|
if tcErr != nil {
|
|
return 0, tcErr
|
|
}
|
|
res, err = s.db.ExecContext(ctx,
|
|
`DELETE FROM kg_entities WHERE agent_id = $1 AND confidence < $2`+tc,
|
|
append([]any{aid, minConfidence}, tcArgs...)...,
|
|
)
|
|
} else {
|
|
tc, tcArgs, _, tcErr := scopeClause(ctx, 4)
|
|
if tcErr != nil {
|
|
return 0, tcErr
|
|
}
|
|
res, err = s.db.ExecContext(ctx,
|
|
`DELETE FROM kg_entities WHERE agent_id = $1 AND user_id = $2 AND confidence < $3`+tc,
|
|
append([]any{aid, userID, minConfidence}, tcArgs...)...,
|
|
)
|
|
}
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
n, _ := res.RowsAffected()
|
|
return int(n), nil
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) Stats(ctx context.Context, agentID, userID string) (*store.GraphStats, error) {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("kg stats: %w", err)
|
|
}
|
|
stats := &store.GraphStats{EntityTypes: make(map[string]int)}
|
|
|
|
userFilter := ""
|
|
args := []any{aid}
|
|
idx := 2
|
|
if userID != "" {
|
|
userFilter = fmt.Sprintf(" AND user_id = $%d", idx)
|
|
args = append(args, userID)
|
|
idx++
|
|
}
|
|
tc, tcArgs, _, err := scopeClause(ctx, idx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
tenantFilter := tc
|
|
args = append(args, tcArgs...)
|
|
|
|
if err := s.db.QueryRowContext(ctx,
|
|
`SELECT COUNT(*) FROM kg_entities WHERE agent_id = $1 AND valid_until IS NULL`+userFilter+tenantFilter, args...,
|
|
).Scan(&stats.EntityCount); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := s.db.QueryRowContext(ctx,
|
|
`SELECT COUNT(*) FROM kg_relations WHERE agent_id = $1 AND valid_until IS NULL`+userFilter+tenantFilter, args...,
|
|
).Scan(&stats.RelationCount); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
rows, err := s.db.QueryContext(ctx,
|
|
`SELECT entity_type, COUNT(*) FROM kg_entities WHERE agent_id = $1 AND valid_until IS NULL`+userFilter+tenantFilter+` GROUP BY entity_type`, args...,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
for rows.Next() {
|
|
var t string
|
|
var c int
|
|
if err := rows.Scan(&t, &c); err != nil {
|
|
continue
|
|
}
|
|
stats.EntityTypes[t] = c
|
|
}
|
|
|
|
// Fetch distinct user IDs (only when not filtering by specific user)
|
|
if userID == "" {
|
|
uidRows, uidErr := s.db.QueryContext(ctx,
|
|
`SELECT DISTINCT user_id FROM kg_entities WHERE agent_id = $1`+tenantFilter+` AND user_id != '' ORDER BY user_id`,
|
|
append([]any{aid}, tcArgs...)...,
|
|
)
|
|
if uidErr == nil {
|
|
defer uidRows.Close()
|
|
for uidRows.Next() {
|
|
var uid string
|
|
if uidRows.Scan(&uid) == nil && uid != "" {
|
|
stats.UserIDs = append(stats.UserIDs, uid)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return stats, nil
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) Close() error { return nil }
|