mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-08-06 10:21:11 +00:00
viettranx
4ebd6698b0
5 CRITICAL sites migrated: UpsertEntity (INSERT), GetEntity (SELECT), DeleteEntity (DELETE), ListEntities (SELECT), SearchEntities (SELECT). All now return wrapped parse errors at the caller boundary instead of hitting FK violations (writes) or returning empty results (reads). Phase 4 Step 4b.3 of agent identity hardening (TD-1).
370 lines
11 KiB
Go
370 lines
11 KiB
Go
package pg
|
|
|
|
import (
|
|
"cmp"
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"fmt"
|
|
"slices"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/nextlevelbuilder/goclaw/internal/store"
|
|
)
|
|
|
|
// PGKnowledgeGraphStore implements store.KnowledgeGraphStore backed by Postgres.
|
|
type PGKnowledgeGraphStore struct {
|
|
db *sql.DB
|
|
embProvider store.EmbeddingProvider
|
|
}
|
|
|
|
// NewPGKnowledgeGraphStore creates a new PG-backed knowledge graph store.
|
|
func NewPGKnowledgeGraphStore(db *sql.DB) *PGKnowledgeGraphStore {
|
|
return &PGKnowledgeGraphStore{db: db}
|
|
}
|
|
|
|
// SetEmbeddingProvider configures the embedding provider for semantic search.
|
|
func (s *PGKnowledgeGraphStore) SetEmbeddingProvider(provider store.EmbeddingProvider) {
|
|
s.embProvider = provider
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) UpsertEntity(ctx context.Context, entity *store.Entity) error {
|
|
aid, err := parseUUID(entity.AgentID)
|
|
if err != nil {
|
|
return fmt.Errorf("kg upsert entity: %w", err)
|
|
}
|
|
props, err := json.Marshal(entity.Properties)
|
|
if err != nil {
|
|
props = []byte("{}")
|
|
}
|
|
now := time.Now()
|
|
id := uuid.Must(uuid.NewV7())
|
|
tid := tenantIDForInsert(ctx)
|
|
var actualID uuid.UUID
|
|
if err = s.db.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, entity.UserID, entity.ExternalID, entity.Name, entity.EntityType,
|
|
entity.Description, props, entity.SourceID, entity.Confidence, tid, now,
|
|
).Scan(&actualID); err != nil {
|
|
return err
|
|
}
|
|
|
|
// Generate embedding in background (best-effort, non-blocking)
|
|
go s.EmbedEntity(context.WithoutCancel(ctx), actualID.String(), entity.Name, entity.Description)
|
|
return nil
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) GetEntity(ctx context.Context, agentID, userID, entityID string) (*store.Entity, error) {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("kg get entity: agent: %w", err)
|
|
}
|
|
eid, err := parseUUID(entityID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("kg get entity: id: %w", err)
|
|
}
|
|
|
|
var row entityRow
|
|
if store.IsSharedKG(ctx) {
|
|
tc, tcArgs, _, err := scopeClause(ctx, 3)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
err = pkgSqlxDB.GetContext(ctx, &row, `
|
|
SELECT id, agent_id, user_id, external_id, name, entity_type, description,
|
|
properties, source_id, confidence, created_at, updated_at
|
|
FROM kg_entities WHERE id = $1 AND agent_id = $2`+tc,
|
|
append([]any{eid, aid}, tcArgs...)...,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
} else {
|
|
tc, tcArgs, _, err := scopeClause(ctx, 4)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
err = pkgSqlxDB.GetContext(ctx, &row, `
|
|
SELECT id, agent_id, user_id, external_id, name, entity_type, description,
|
|
properties, source_id, confidence, created_at, updated_at
|
|
FROM kg_entities WHERE id = $1 AND agent_id = $2 AND user_id = $3`+tc,
|
|
append([]any{eid, aid, userID}, tcArgs...)...,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
e := row.toEntity()
|
|
return &e, nil
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) DeleteEntity(ctx context.Context, agentID, userID, entityID string) error {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return fmt.Errorf("kg delete entity: agent: %w", err)
|
|
}
|
|
eid, err := parseUUID(entityID)
|
|
if err != nil {
|
|
return fmt.Errorf("kg delete entity: 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_entities WHERE id = $1 AND agent_id = $2`+tc,
|
|
append([]any{eid, aid}, tcArgs...)...,
|
|
)
|
|
return err
|
|
}
|
|
tc, tcArgs, _, err := scopeClause(ctx, 4)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = s.db.ExecContext(ctx,
|
|
`DELETE FROM kg_entities WHERE id = $1 AND agent_id = $2 AND user_id = $3`+tc,
|
|
append([]any{eid, aid, userID}, tcArgs...)...,
|
|
)
|
|
return err
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) ListEntities(ctx context.Context, agentID, userID string, opts store.EntityListOptions) ([]store.Entity, error) {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("kg list entities: %w", err)
|
|
}
|
|
limit := opts.Limit
|
|
if limit <= 0 {
|
|
limit = 50
|
|
}
|
|
|
|
// Build dynamic WHERE clause: always filter by agent_id, optionally by user_id and entity_type.
|
|
// Default to current facts only (valid_until IS NULL) — expired entities excluded.
|
|
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++
|
|
}
|
|
if opts.EntityType != "" {
|
|
where += fmt.Sprintf(" AND entity_type = $%d", idx)
|
|
args = append(args, opts.EntityType)
|
|
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, opts.Offset)
|
|
query := fmt.Sprintf(`
|
|
SELECT id, agent_id, user_id, external_id, name, entity_type, description,
|
|
properties, source_id, confidence, created_at, updated_at
|
|
FROM kg_entities WHERE %s
|
|
ORDER BY updated_at DESC LIMIT $%d OFFSET $%d`, where, idx, idx+1)
|
|
|
|
var rows []entityRow
|
|
if err = pkgSqlxDB.SelectContext(ctx, &rows, query, args...); err != nil {
|
|
return nil, err
|
|
}
|
|
entities := make([]store.Entity, len(rows))
|
|
for i := range rows {
|
|
entities[i] = rows[i].toEntity()
|
|
}
|
|
return entities, nil
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) SearchEntities(ctx context.Context, agentID, userID, query string, limit int) ([]store.Entity, error) {
|
|
aid, err := parseUUID(agentID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("kg search entities: %w", err)
|
|
}
|
|
if limit <= 0 {
|
|
limit = 20
|
|
}
|
|
|
|
shared := store.IsSharedKG(ctx)
|
|
|
|
// FTS search using tsvector
|
|
ftsResults, err := s.ftsSearchEntities(ctx, aid, userID, query, limit*2, shared)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Vector search if provider available
|
|
var vecResults []scoredEntity
|
|
if s.embProvider != nil {
|
|
embeddings, embErr := s.embProvider.Embed(ctx, []string{query})
|
|
if embErr == nil && len(embeddings) > 0 {
|
|
vecResults, err = s.vectorSearchEntities(ctx, embeddings[0], aid, userID, limit*2, shared)
|
|
if err != nil {
|
|
vecResults = nil
|
|
}
|
|
}
|
|
}
|
|
|
|
// If no vector results, fall back to FTS-only
|
|
if len(vecResults) == 0 {
|
|
if len(ftsResults) > limit {
|
|
ftsResults = ftsResults[:limit]
|
|
}
|
|
entities := make([]store.Entity, len(ftsResults))
|
|
for i, r := range ftsResults {
|
|
entities[i] = r.Entity
|
|
}
|
|
return entities, nil
|
|
}
|
|
|
|
// Hybrid merge with weights: 0.3 FTS, 0.7 vector
|
|
textW, vecW := 0.3, 0.7
|
|
if len(ftsResults) == 0 {
|
|
textW, vecW = 0, 1.0
|
|
}
|
|
merged := hybridMergeEntities(ftsResults, vecResults, textW, vecW)
|
|
|
|
if len(merged) > limit {
|
|
merged = merged[:limit]
|
|
}
|
|
return merged, nil
|
|
}
|
|
|
|
type scoredEntity struct {
|
|
Entity store.Entity
|
|
Score float64
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) ftsSearchEntities(ctx context.Context, agentID uuid.UUID, userID, query string, limit int, shared bool) ([]scoredEntity, error) {
|
|
where := "agent_id = $1 AND valid_until IS NULL AND tsv @@ plainto_tsquery('simple', $2)"
|
|
args := []any{agentID, query}
|
|
idx := 3
|
|
if !shared && 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, query, limit)
|
|
q := fmt.Sprintf(`
|
|
SELECT id, agent_id, user_id, external_id, name, entity_type, description,
|
|
properties, source_id, confidence, created_at, updated_at,
|
|
ts_rank(tsv, plainto_tsquery('simple', $%d)) AS score
|
|
FROM kg_entities
|
|
WHERE %s
|
|
ORDER BY score DESC LIMIT $%d`, idx, where, idx+1)
|
|
|
|
var sRows []scoredEntityRow
|
|
if err = pkgSqlxDB.SelectContext(ctx, &sRows, q, args...); err != nil {
|
|
return nil, err
|
|
}
|
|
results := make([]scoredEntity, len(sRows))
|
|
for i := range sRows {
|
|
results[i] = scoredEntity{Entity: sRows[i].toEntity(), Score: sRows[i].Score}
|
|
}
|
|
return results, nil
|
|
}
|
|
|
|
func (s *PGKnowledgeGraphStore) vectorSearchEntities(ctx context.Context, embedding []float32, agentID uuid.UUID, userID string, limit int, shared bool) ([]scoredEntity, error) {
|
|
vecStr := vectorToString(embedding)
|
|
|
|
where := "agent_id = $1 AND valid_until IS NULL AND embedding IS NOT NULL"
|
|
args := []any{agentID}
|
|
idx := 2
|
|
if !shared && 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, vecStr, limit)
|
|
q := fmt.Sprintf(`
|
|
SELECT id, agent_id, user_id, external_id, name, entity_type, description,
|
|
properties, source_id, confidence, created_at, updated_at,
|
|
1 - (embedding <=> $%d::vector) AS score
|
|
FROM kg_entities
|
|
WHERE %s
|
|
ORDER BY embedding <=> $%d::vector LIMIT $%d`, idx, where, idx, idx+1)
|
|
|
|
var sRows []scoredEntityRow
|
|
if err = pkgSqlxDB.SelectContext(ctx, &sRows, q, args...); err != nil {
|
|
return nil, err
|
|
}
|
|
results := make([]scoredEntity, len(sRows))
|
|
for i := range sRows {
|
|
results[i] = scoredEntity{Entity: sRows[i].toEntity(), Score: sRows[i].Score}
|
|
}
|
|
return results, nil
|
|
}
|
|
|
|
// hybridMergeEntities combines ILIKE and vector results with weighted scoring.
|
|
func hybridMergeEntities(ilike, vec []scoredEntity, textWeight, vectorWeight float64) []store.Entity {
|
|
type mergedEntry struct {
|
|
Entity store.Entity
|
|
Score float64
|
|
}
|
|
seen := make(map[string]*mergedEntry)
|
|
|
|
for _, r := range ilike {
|
|
if existing, ok := seen[r.Entity.ID]; ok {
|
|
existing.Score += r.Score * textWeight
|
|
} else {
|
|
seen[r.Entity.ID] = &mergedEntry{Entity: r.Entity, Score: r.Score * textWeight}
|
|
}
|
|
}
|
|
for _, r := range vec {
|
|
if existing, ok := seen[r.Entity.ID]; ok {
|
|
existing.Score += r.Score * vectorWeight
|
|
} else {
|
|
seen[r.Entity.ID] = &mergedEntry{Entity: r.Entity, Score: r.Score * vectorWeight}
|
|
}
|
|
}
|
|
|
|
results := make([]store.Entity, 0, len(seen))
|
|
scores := make(map[string]float64, len(seen))
|
|
for id, entry := range seen {
|
|
results = append(results, entry.Entity)
|
|
scores[id] = entry.Score
|
|
}
|
|
|
|
slices.SortFunc(results, func(a, b store.Entity) int {
|
|
return cmp.Compare(scores[b.ID], scores[a.ID]) // descending
|
|
})
|
|
|
|
return results
|
|
}
|