xinmi/internal/application/repository/retriever/legacygraph/repository.go

package legacygraph

import (
	"context"
	"encoding/json"
	"fmt"
	"slices"
	"strconv"
	"strings"

	neo4jrepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/neo4j"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/neo4j/neo4j-go-driver/v6/neo4j"
	"gorm.io/gorm"
)

// Repository preserves source records while exposing verified native chunks to RAG.
type Repository struct {
	db       *gorm.DB
	fallback interfaces.RetrieveGraphRepository
}

func NewRepository(db *gorm.DB, driver neo4j.Driver) *Repository {
	r := &Repository{db: db}
	if driver != nil {
		r.fallback = neo4jrepo.NewNeo4jRepository(driver)
	}
	return r
}
func NewRetrieveRepository(db *gorm.DB, driver neo4j.Driver) interfaces.RetrieveGraphRepository {
	return NewRepository(db, driver)
}
func (r *Repository) AddGraph(ctx context.Context, ns types.NameSpace, graphs []*types.GraphData) error {
	return r.addNativeGraphs(ctx, ns, graphs, nil)
}
func (r *Repository) DelGraph(ctx context.Context, ns []types.NameSpace) error {
	return r.deleteNativeGraphs(ctx, ns)
}

type Record struct {
	ID      int64           `json:"id" gorm:"column:legacy_id"`
	Payload json.RawMessage `json:"payload" gorm:"column:payload"`
}
type Evidence struct {
	Record
	NativeKnowledgeID string `json:"native_knowledge_id"`
	NativeChunkID     string `json:"native_chunk_id"`
	MappingState      string `json:"mapping_state"`
}
type Overview struct {
	Counts        map[string]int64 `json:"counts"`
	Sources       int64            `json:"sources"`
	MappedSources int64            `json:"mapped_sources"`
	Nodes         []Record         `json:"nodes"`
	Relations     []Record         `json:"relations"`
	Evidence      []Evidence       `json:"evidence"`
	NextCursor    int64            `json:"next_cursor"`
}

func (r *Repository) scope(ctx context.Context, kb string) (*gorm.DB, error) {
	tenant, ok := types.TenantIDFromContext(ctx)
	if !ok || tenant == 0 || kb == "" {
		return nil, fmt.Errorf("graph tenant and knowledge base are required")
	}
	return r.db.WithContext(ctx).Table("native_graph_records").Where("tenant_id = ? AND knowledge_base_id = ?", tenant, kb), nil
}
func (r *Repository) Browse(ctx context.Context, kb, query string, cursor, entity int64) (*Overview, error) {
	scope, err := r.scope(ctx, kb)
	if err != nil {
		return nil, err
	}
	out := &Overview{Counts: map[string]int64{}, Nodes: []Record{}, Relations: []Record{}, Evidence: []Evidence{}}
	var counts []struct {
		SourceTable string
		N           int64
	}
	if err = scope.Session(&gorm.Session{}).Select("source_table, count(*) AS n").Group("source_table").Scan(&counts).Error; err != nil {
		return nil, err
	}
	for _, c := range counts {
		out.Counts[c.SourceTable] = c.N
	}
	tenant := types.MustTenantIDFromContext(ctx)
	var coverage struct {
		Sources int64
		Mapped  int64
	}
	// Overview coverage records migration-time matches; each returned link is revalidated below.
	err = r.db.WithContext(ctx).Raw(`SELECT count(*) AS sources, count(native_chunk_id) AS mapped FROM native_graph_sources WHERE tenant_id=? AND knowledge_base_id=?`, tenant, kb).Scan(&coverage).Error
	if err != nil {
		return nil, err
	}
	out.Sources = coverage.Sources
	out.MappedSources = coverage.Mapped
	q := scope.Session(&gorm.Session{}).Select("legacy_id, payload").Where("source_table = ? AND legacy_id > ?", "knowledgeentity", cursor)
	if entity > 0 {
		q = q.Where("legacy_id = ?", entity)
	} else if strings.TrimSpace(query) != "" {
		q = q.Where("strpos(lower(payload->>'canonical_name'), lower(?)) > 0 OR strpos(lower(?), lower(payload->>'canonical_name')) > 0", query, query)
	}
	if err = q.Order("legacy_id").Limit(51).Find(&out.Nodes).Error; err != nil {
		return nil, err
	}
	if len(out.Nodes) > 50 {
		out.NextCursor = out.Nodes[49].ID
		out.Nodes = out.Nodes[:50]
	}
	if err = r.projectNativeEntityAttributes(ctx, types.NameSpace{KnowledgeBase: kb}, out.Nodes, nil); err != nil {
		return nil, err
	}
	if entity <= 0 {
		return out, nil
	}
	err = r.db.WithContext(ctx).Raw(`SELECT r.legacy_id, r.payload || jsonb_build_object('from_name',a.payload->>'canonical_name','to_name',b.payload->>'canonical_name') AS payload FROM native_graph_records r LEFT JOIN native_graph_records a ON a.tenant_id=r.tenant_id AND a.knowledge_base_id=r.knowledge_base_id AND a.source_table='knowledgeentity' AND a.legacy_id=(r.payload->>'from_entity_id')::bigint LEFT JOIN native_graph_records b ON b.tenant_id=r.tenant_id AND b.knowledge_base_id=r.knowledge_base_id AND b.source_table='knowledgeentity' AND b.legacy_id=(r.payload->>'to_entity_id')::bigint WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.source_table='knowledgerelation' AND (r.payload->>'from_entity_id'=? OR r.payload->>'to_entity_id'=?) ORDER BY r.legacy_id LIMIT 100`, tenant, kb, strconv.FormatInt(entity, 10), strconv.FormatInt(entity, 10)).Scan(&out.Relations).Error
	if err != nil {
		return nil, err
	}
	err = r.db.WithContext(ctx).Raw(`SELECT r.legacy_id,r.payload, COALESCE(c.knowledge_id,'') AS native_knowledge_id, COALESCE(c.id,'') AS native_chunk_id, CASE WHEN c.id IS NOT NULL THEN s.mapping_state WHEN s.native_chunk_id IS NOT NULL THEN 'stale' ELSE COALESCE(s.mapping_state,'unmapped') END AS mapping_state FROM native_graph_records r LEFT JOIN native_graph_sources s ON s.tenant_id=r.tenant_id AND s.knowledge_base_id=r.knowledge_base_id AND s.source_chunk_id=r.payload->>'source_chunk_id' LEFT JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id AND c.deleted_at IS NULL AND c.is_enabled=true AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256 AND EXISTS (SELECT 1 FROM knowledges k WHERE k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL) AND EXISTS (SELECT 1 FROM knowledge_bases kb WHERE kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL) WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.source_table='knowledgegraphevidence' AND r.payload->>'subject_type'='entity' AND r.payload->>'subject_id'=? ORDER BY r.legacy_id LIMIT 100`, tenant, kb, strconv.FormatInt(entity, 10)).Scan(&out.Evidence).Error
	return out, err
}

func (r *Repository) SearchNode(ctx context.Context, ns types.NameSpace, terms []string) (*types.GraphData, error) {
	out := &types.GraphData{Node: []*types.GraphNode{}, Relation: []*types.GraphRelation{}}
	scope, err := r.scope(ctx, ns.KnowledgeBase)
	if err != nil {
		return nil, err
	}
	terms = boundedTerms(terms)
	if len(terms) == 0 {
		return out, nil
	}
	clauses := make([]string, 0, len(terms))
	args := make([]any, 0, len(terms)*2)
	for _, term := range terms {
		clauses = append(clauses, "(strpos(lower(payload->>'canonical_name'),lower(?))>0 OR strpos(lower(?),lower(payload->>'canonical_name'))>0)")
		args = append(args, term, term)
	}
	var nodes []Record
	err = scope.Session(&gorm.Session{}).Select("legacy_id,payload").Where("source_table='knowledgeentity' AND payload->>'status'='active'").Where(strings.Join(clauses, " OR "), args...).Order("legacy_id").Limit(16).Find(&nodes).Error
	if err != nil {
		return nil, err
	}
	tenant := types.MustTenantIDFromContext(ctx)
	seedIDs := []string{}
	seenIDs := map[int64]bool{}
	for _, node := range nodes {
		seedIDs = append(seedIDs, strconv.FormatInt(node.ID, 10))
		seenIDs[node.ID] = true
	}
	var relations []Record
	if len(seedIDs) > 0 {
		q, _ := r.scope(ctx, ns.KnowledgeBase)
		err = q.Select("legacy_id,payload").Where("source_table='knowledgerelation' AND payload->>'status'='active' AND (payload->>'from_entity_id' IN ? OR payload->>'to_entity_id' IN ?)", seedIDs, seedIDs).Order("legacy_id").Limit(32).Find(&relations).Error
		if err != nil {
			return nil, err
		}
		neighborIDs := []int64{}
		for _, rel := range relations {
			var edge struct {
				From int64 `json:"from_entity_id"`
				To   int64 `json:"to_entity_id"`
			}
			if err = json.Unmarshal(rel.Payload, &edge); err != nil {
				return nil, err
			}
			for _, id := range []int64{edge.From, edge.To} {
				if !seenIDs[id] && len(neighborIDs) < 16 {
					seenIDs[id] = true
					neighborIDs = append(neighborIDs, id)
				}
			}
		}
		if len(neighborIDs) > 0 {
			var neighbors []Record
			q, _ = r.scope(ctx, ns.KnowledgeBase)
			if err = q.Select("legacy_id,payload").Where("source_table='knowledgeentity' AND payload->>'status'='active' AND legacy_id IN ?", neighborIDs).Order("legacy_id").Find(&neighbors).Error; err != nil {
				return nil, err
			}
			nodes = append(nodes, neighbors...)
		}
	}
	chunksByEntity := map[int64][]string{}
	chunksByRelation := map[int64][]string{}
	if len(nodes) > 0 {
		type subject struct {
			ID   int64  `json:"subject_id"`
			Type string `json:"subject_type"`
		}
		ids := make([]subject, 0, len(nodes)+len(relations))
		for _, node := range nodes {
			ids = append(ids, subject{node.ID, "entity"})
		}
		for _, relation := range relations {
			ids = append(ids, subject{relation.ID, "relation"})
		}
		encodedIDs, err := json.Marshal(ids)
		if err != nil {
			return nil, err
		}
		var matches []struct {
			SubjectID   int64
			SubjectType string
			ChunkID     string
		}
		// One bounded database request, with the evidence index and four-chunk cap per entity.
		query := `SELECT seed.subject_id,seed.subject_type,matched.chunk_id FROM jsonb_to_recordset(?::jsonb) AS seed(subject_id bigint,subject_type text) CROSS JOIN LATERAL (
			SELECT DISTINCT c.id AS chunk_id FROM native_graph_records e
			JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id'
			JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id AND c.deleted_at IS NULL AND c.is_enabled=true AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256
			JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
			JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
			WHERE e.tenant_id=? AND e.knowledge_base_id=? AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_type'=seed.subject_type AND e.payload->>'subject_id'=seed.subject_id::text AND e.payload->>'status'='active'`
		args := []any{string(encodedIDs), tenant, ns.KnowledgeBase}
		if ns.Knowledge != "" {
			query += " AND c.knowledge_id=?"
			args = append(args, ns.Knowledge)
		}
		query += " ORDER BY c.id LIMIT 4) matched"
		if err = r.db.WithContext(ctx).Raw(query, args...).Scan(&matches).Error; err != nil {
			return nil, err
		}
		for _, match := range matches {
			if match.SubjectType == "entity" {
				chunksByEntity[match.SubjectID] = append(chunksByEntity[match.SubjectID], match.ChunkID)
			} else {
				chunksByRelation[match.SubjectID] = append(chunksByRelation[match.SubjectID], match.ChunkID)
			}
		}
	}
	if err = r.projectNativeEntityAttributes(ctx, ns, nodes, chunksByEntity); err != nil {
		return nil, err
	}
	names := map[int64]string{}
	for _, node := range nodes {
		var entity struct {
			Name       string   `json:"canonical_name"`
			Type       string   `json:"entity_type"`
			Attributes []string `json:"attributes"`
		}
		if err = json.Unmarshal(node.Payload, &entity); err != nil {
			return nil, err
		}
		chunks := chunksByEntity[node.ID]
		if len(chunks) > 0 {
			names[node.ID] = entity.Name
			attributes := []string{entity.Type}
			identity := fmt.Sprintf("legacy_entity:%d", node.ID)
			for _, attribute := range entity.Attributes {
				if attribute != entity.Type && attribute != identity {
					attributes = append(attributes, attribute)
				}
			}
			attributes = append(attributes, identity)
			out.Node = append(out.Node, &types.GraphNode{Name: entity.Name, Chunks: chunks, Attributes: attributes})
		}
	}
	for _, rel := range relations {
		var edge struct {
			From   int64  `json:"from_entity_id"`
			To     int64  `json:"to_entity_id"`
			Type   string `json:"relation_type"`
			Origin string `json:"origin"`
		}
		if err = json.Unmarshal(rel.Payload, &edge); err != nil {
			return nil, err
		}
		if edge.Origin == nativeOrigin && len(chunksByRelation[rel.ID]) == 0 {
			continue
		}
		if names[edge.From] != "" && names[edge.To] != "" {
			out.Relation = append(out.Relation, &types.GraphRelation{Node1: names[edge.From], Node2: names[edge.To], Type: edge.Type})
		}
	}
	if r.fallback != nil {
		fallback, err := r.fallback.SearchNode(ctx, ns, terms)
		if err != nil {
			return nil, err
		}
		if fallback != nil {
			out.Node = append(out.Node, fallback.Node...)
			out.Relation = append(out.Relation, fallback.Relation...)
		}
	}
	return out, nil
}

// Attribute ownership is checked independently of the four returned chunks:
// an entity may still have a bare reference after its attribute source expires.
func (r *Repository) projectNativeEntityAttributes(ctx context.Context, ns types.NameSpace, nodes []Record, chunks map[int64][]string) error {
	type subject struct {
		ID         int64           `json:"subject_id"`
		Knowledge  string          `json:"knowledge_id"`
		Type       string          `json:"entity_type"`
		Attributes json.RawMessage `json:"attributes"`
	}
	seeds := []subject{}
	payloads := map[int64]map[string]json.RawMessage{}
	for _, node := range nodes {
		var entity struct {
			Origin     string          `json:"origin"`
			Knowledge  string          `json:"native_knowledge_id"`
			Type       string          `json:"entity_type"`
			Attributes json.RawMessage `json:"attributes"`
		}
		if err := json.Unmarshal(node.Payload, &entity); err != nil {
			return err
		}
		if entity.Origin != nativeOrigin {
			continue
		}
		var payload map[string]json.RawMessage
		if err := json.Unmarshal(node.Payload, &payload); err != nil {
			return err
		}
		payloads[node.ID] = payload
		seeds = append(seeds, subject{node.ID, entity.Knowledge, entity.Type, entity.Attributes})
	}
	if len(seeds) == 0 {
		return nil
	}
	encoded, err := json.Marshal(seeds)
	if err != nil {
		return err
	}
	var matches []struct {
		SubjectID      int64
		Snapshots      json.RawMessage
		Chunks         json.RawMessage
		AttributeChunk string
	}
	query := `SELECT seed.subject_id,jsonb_agg(DISTINCT e.payload->'entity_attributes') AS snapshots,
		to_jsonb((array_agg(DISTINCT c.id ORDER BY c.id))[1:4]) AS chunks,
		COALESCE(min(c.id) FILTER (WHERE e.payload->'entity_attributes'->>0 IS NOT NULL
			AND (e.payload->'entity_attributes'=seed.attributes OR e.payload->'entity_attributes'->>0=seed.entity_type)),'') AS attribute_chunk
		FROM jsonb_to_recordset(?::jsonb) AS seed(subject_id bigint,knowledge_id text,entity_type text,attributes jsonb)
		JOIN native_graph_records e ON e.payload->>'subject_id'=seed.subject_id::text AND e.payload->>'native_knowledge_id'=seed.knowledge_id
		JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id'
		JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id
			AND s.native_knowledge_id=c.knowledge_id AND c.deleted_at IS NULL AND c.is_enabled=true
			AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256
		JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
		JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
		WHERE e.tenant_id=? AND e.knowledge_base_id=? AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_type'='entity'
			AND e.payload->>'origin'=? AND e.payload->>'status'='active' AND s.mapping_state='native_extraction'
			AND e.payload->>'source_content_hash'=s.native_content_sha256 AND e.source_snapshot_sha256=s.native_content_sha256
			AND s.source_content_hash=s.native_content_sha256`
	args := []any{string(encoded), types.MustTenantIDFromContext(ctx), ns.KnowledgeBase, nativeOrigin}
	if ns.Knowledge != "" {
		query += " AND c.knowledge_id=?"
		args = append(args, ns.Knowledge)
	}
	query += " GROUP BY seed.subject_id"
	if err = r.db.WithContext(ctx).Raw(query, args...).Scan(&matches).Error; err != nil {
		return err
	}
	supported := map[int64][]string{}
	attributeChunks := map[int64]string{}
	verifiedChunks := map[int64][]string{}
	for _, match := range matches {
		supported[match.SubjectID] = compatibleNativeAttributeSnapshots(match.Snapshots)
		attributeChunks[match.SubjectID] = match.AttributeChunk
		var current []string
		if err = json.Unmarshal(match.Chunks, &current); err != nil {
			return err
		}
		verifiedChunks[match.SubjectID] = current
	}
	for i, node := range nodes {
		payload, native := payloads[node.ID]
		if !native {
			continue
		}
		if chunks != nil {
			chunks[node.ID] = verifiedChunks[node.ID]
		}
		attributes := []string{}
		entityType := "concept"
		if snapshot := supported[node.ID]; len(snapshot) > 0 {
			var stored []string
			if err := json.Unmarshal(payload["attributes"], &stored); err == nil && slices.Equal(stored, snapshot) {
				attributes = stored
				entityType = snapshot[0]
			}
			var storedType string
			if err := json.Unmarshal(payload["entity_type"], &storedType); err == nil && storedType == snapshot[0] {
				entityType = storedType
			}
		}
		if chunk := attributeChunks[node.ID]; chunks != nil && chunk != "" && (len(attributes) > 0 || entityType != "concept") {
			verified := []string{chunk}
			for _, prior := range chunks[node.ID] {
				if prior != chunk && len(verified) < 4 {
					verified = append(verified, prior)
				}
			}
			chunks[node.ID] = verified
		}
		payload["attributes"], _ = json.Marshal(attributes)
		payload["entity_type"], _ = json.Marshal(entityType)
		if nodes[i].Payload, err = json.Marshal(payload); err != nil {
			return err
		}
	}
	return nil
}

func compatibleNativeAttributeSnapshots(raw json.RawMessage) []string {
	var snapshots []json.RawMessage
	if err := json.Unmarshal(raw, &snapshots); err != nil {
		return nil
	}
	var attributes []string
	for _, snapshot := range snapshots {
		var values []*string
		if err := json.Unmarshal(snapshot, &values); err != nil {
			return nil
		}
		if len(values) == 0 {
			continue
		}
		current := make([]string, len(values))
		for i, value := range values {
			if value == nil {
				return nil
			}
			current[i] = *value
		}
		if len(attributes) > 0 && !slices.Equal(attributes, current) {
			return nil
		}
		attributes = current
	}
	return attributes
}

func boundedTerms(terms []string) []string {
	out := []string{}
	seen := map[string]bool{}
	for _, term := range terms {
		term = strings.TrimSpace(term)
		if term == "" || seen[term] {
			continue
		}
		if len([]rune(term)) > 200 {
			term = string([]rune(term)[:200])
		}
		seen[term] = true
		out = append(out, term)
		if len(out) == 8 {
			break
		}
	}
	return out
}