442 lines
13 KiB
Go
442 lines
13 KiB
Go
package entities
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/andydixon/relata/internal/analysis"
|
|
"github.com/andydixon/relata/internal/config"
|
|
"github.com/andydixon/relata/internal/domain"
|
|
"github.com/andydixon/relata/internal/ids"
|
|
"github.com/andydixon/relata/internal/jobs"
|
|
"github.com/andydixon/relata/internal/observability"
|
|
)
|
|
|
|
// Service extracts entities and claims from an item's chunks and stores them.
|
|
//
|
|
// Extraction runs in two stages, as separate and separately-labelled sources of
|
|
// evidence: deterministic rules first, then model-assisted extraction for the
|
|
// broader concepts rules cannot reach. Every stored record carries the method
|
|
// that produced it, so a user can tell a matched pattern from a model's guess.
|
|
type Service struct {
|
|
entities domain.EntityRepository
|
|
aliases domain.AliasRepository
|
|
claims domain.ClaimRepository
|
|
analysis *analysis.Service
|
|
cfg *config.Config
|
|
}
|
|
|
|
// NewService builds the extraction service.
|
|
func NewService(ents domain.EntityRepository, aliases domain.AliasRepository,
|
|
claims domain.ClaimRepository, an *analysis.Service, cfg *config.Config) *Service {
|
|
return &Service{entities: ents, aliases: aliases, claims: claims, analysis: an, cfg: cfg}
|
|
}
|
|
|
|
// maxModelChunksPerItem bounds how many chunks are sent to the model for
|
|
// extraction, so a large document cannot consume unlimited model time.
|
|
const maxModelChunksPerItem = 12
|
|
|
|
// observation accumulates everything seen about one entity within an item.
|
|
type observation struct {
|
|
names []string
|
|
typ domain.EntityType
|
|
refs []domain.SourceRef
|
|
aliases map[string]bool
|
|
methods map[domain.ExtractionMethod]bool
|
|
count int
|
|
}
|
|
|
|
// ExtractForItem extracts entities from an item's chunks and stores them.
|
|
//
|
|
// Returns the number of distinct entities recorded. Model failures are returned
|
|
// to the caller, which treats them as degraded quality rather than a pipeline
|
|
// failure: the deterministic results are already persisted by then.
|
|
func (s *Service) ExtractForItem(ctx context.Context, item *domain.Item, chunks []*domain.Chunk, pool *jobs.Pool) (int, error) {
|
|
log := observability.LoggerFrom(ctx)
|
|
obs := map[string]*observation{}
|
|
|
|
record := func(name string, typ domain.EntityType, ref domain.SourceRef, method domain.ExtractionMethod, aliases []string) {
|
|
name = strings.TrimSpace(name)
|
|
if name == "" || !typ.Valid() {
|
|
return
|
|
}
|
|
key := Key(typ, name)
|
|
if Normalise(name) == "" {
|
|
return
|
|
}
|
|
o := obs[key]
|
|
if o == nil {
|
|
o = &observation{
|
|
typ: typ,
|
|
aliases: map[string]bool{},
|
|
methods: map[domain.ExtractionMethod]bool{},
|
|
}
|
|
obs[key] = o
|
|
}
|
|
o.names = append(o.names, name)
|
|
o.methods[method] = true
|
|
o.count++
|
|
// Source references are bounded per entity per item: a name repeated a
|
|
// hundred times does not need a hundred stored references.
|
|
if len(o.refs) < 8 {
|
|
o.refs = append(o.refs, ref)
|
|
}
|
|
for _, a := range aliases {
|
|
if a != "" && !strings.EqualFold(a, name) {
|
|
o.aliases[a] = true
|
|
}
|
|
}
|
|
}
|
|
|
|
// Stage 1: deterministic extraction over every chunk.
|
|
for _, c := range chunks {
|
|
for _, m := range ExtractDeterministic(c.Text) {
|
|
record(m.Name, m.Type, domain.SourceRef{
|
|
ItemID: item.ID,
|
|
ChunkID: c.ID,
|
|
// Offsets are translated from chunk-relative to item-relative,
|
|
// so a stored reference addresses the item's extracted text.
|
|
Start: c.StartOffset + m.Start,
|
|
End: c.StartOffset + m.End,
|
|
Quote: m.Quote,
|
|
Page: c.Page,
|
|
}, domain.MethodDeterministic, nil)
|
|
}
|
|
}
|
|
|
|
// Stage 2: model-assisted extraction over a bounded sample of chunks.
|
|
var modelErr error
|
|
if s.analysis.Enabled() {
|
|
limit := len(chunks)
|
|
if limit > maxModelChunksPerItem {
|
|
limit = maxModelChunksPerItem
|
|
}
|
|
for _, c := range chunks[:limit] {
|
|
if err := ctx.Err(); err != nil {
|
|
modelErr = err
|
|
break
|
|
}
|
|
release, err := pool.AcquireAI(ctx)
|
|
if err != nil {
|
|
modelErr = err
|
|
break
|
|
}
|
|
cands, err := s.analysis.ExtractEntities(ctx, c.Text)
|
|
release()
|
|
if err != nil {
|
|
modelErr = err
|
|
log.Debug("model entity extraction failed for a chunk", "error", err.Error())
|
|
continue
|
|
}
|
|
for _, cand := range cands {
|
|
ref := domain.SourceRef{ItemID: item.ID, ChunkID: c.ID, Page: c.Page}
|
|
// Locate the model's quotation in the chunk so the reference
|
|
// points at real text. Unlocatable evidence yields a reference
|
|
// to the chunk rather than a fabricated offset.
|
|
if cand.Evidence != "" {
|
|
if idx := strings.Index(c.Text, cand.Evidence); idx >= 0 {
|
|
ref.Start = c.StartOffset + idx
|
|
ref.End = ref.Start + len(cand.Evidence)
|
|
ref.Quote = cand.Evidence
|
|
}
|
|
}
|
|
if ref.Quote == "" {
|
|
if idx := strings.Index(c.Text, cand.Name); idx >= 0 {
|
|
ref.Start = c.StartOffset + idx
|
|
ref.End = ref.Start + len(cand.Name)
|
|
ref.Quote = cand.Name
|
|
}
|
|
}
|
|
record(cand.Name, cand.Type, ref, domain.MethodModel, cand.Aliases)
|
|
}
|
|
}
|
|
} else {
|
|
modelErr = analysis.ErrDisabled
|
|
}
|
|
|
|
stored := s.persist(ctx, item, obs)
|
|
return stored, modelErr
|
|
}
|
|
|
|
// persist writes the accumulated observations and proposes aliases.
|
|
func (s *Service) persist(ctx context.Context, item *domain.Item, obs map[string]*observation) int {
|
|
log := observability.LoggerFrom(ctx)
|
|
stored := 0
|
|
saved := make([]*domain.Entity, 0, len(obs))
|
|
|
|
for key, o := range obs {
|
|
canonical := Canonicalise(o.names)
|
|
if canonical == "" {
|
|
continue
|
|
}
|
|
method := domain.MethodModel
|
|
if o.methods[domain.MethodDeterministic] {
|
|
// A rule matched, which is the stronger provenance.
|
|
method = domain.MethodDeterministic
|
|
}
|
|
|
|
aliasList := make([]string, 0, len(o.aliases))
|
|
for a := range o.aliases {
|
|
aliasList = append(aliasList, a)
|
|
}
|
|
|
|
e := &domain.Entity{
|
|
ID: ids.New(),
|
|
WorkspaceID: item.WorkspaceID,
|
|
CanonicalName: canonical,
|
|
NormalisedKey: key,
|
|
Type: o.typ,
|
|
Aliases: aliasList,
|
|
SourceRefs: o.refs,
|
|
ItemIDs: []string{item.ID},
|
|
Method: method,
|
|
Mentions: o.count,
|
|
EvidenceScore: evidenceScore(o),
|
|
}
|
|
if method == domain.MethodModel {
|
|
e.ModelVersion = s.analysis.ModelID()
|
|
}
|
|
|
|
got, err := s.entities.UpsertByKey(ctx, e)
|
|
if err != nil {
|
|
log.Warn("could not store an entity", "entity", canonical, "error", err.Error())
|
|
continue
|
|
}
|
|
stored++
|
|
saved = append(saved, got)
|
|
}
|
|
|
|
s.proposeAliases(ctx, item.WorkspaceID, saved)
|
|
return stored
|
|
}
|
|
|
|
// evidenceScore reflects how well-attested an entity is. It is not a
|
|
// probability: it says "several independent signals point at this", nothing
|
|
// about whether the entity is real.
|
|
func evidenceScore(o *observation) float64 {
|
|
score := 0.3
|
|
if o.methods[domain.MethodDeterministic] {
|
|
score += 0.4
|
|
}
|
|
if o.methods[domain.MethodModel] && o.methods[domain.MethodDeterministic] {
|
|
// Both stages agreeing is the strongest evidence available here.
|
|
score += 0.2
|
|
}
|
|
switch {
|
|
case o.count >= 5:
|
|
score += 0.2
|
|
case o.count >= 2:
|
|
score += 0.1
|
|
}
|
|
if score > 1 {
|
|
score = 1
|
|
}
|
|
return score
|
|
}
|
|
|
|
// aliasProposalThreshold is the similarity at which a name pair is worth a
|
|
// human's attention.
|
|
const aliasProposalThreshold = 0.82
|
|
|
|
// maxAliasComparisons bounds the pairwise comparison so a workspace with many
|
|
// entities does not turn alias proposal into an O(n²) sweep.
|
|
const maxAliasComparisons = 400
|
|
|
|
// proposeAliases suggests that two entities may be the same thing.
|
|
//
|
|
// This is the safety valve for conservative normalisation: "BBC" and "British
|
|
// Broadcasting Corporation" get separate keys, and rather than merging them —
|
|
// which would be a guess with real consequences — Relata records a candidate
|
|
// alias and waits for a person to decide.
|
|
func (s *Service) proposeAliases(ctx context.Context, workspaceID string, saved []*domain.Entity) {
|
|
if len(saved) == 0 {
|
|
return
|
|
}
|
|
log := observability.LoggerFrom(ctx)
|
|
|
|
// Compare each newly-seen entity against existing entities that share its
|
|
// loose key — the cheap indexed way to find plausible pairs.
|
|
comparisons := 0
|
|
for _, e := range saved {
|
|
if comparisons >= maxAliasComparisons {
|
|
return
|
|
}
|
|
if e.Type == domain.EntityDate || e.Type == domain.EntityConcept {
|
|
// Dates are already normalised, and concepts overlap by nature.
|
|
continue
|
|
}
|
|
|
|
candidates, _, err := s.entities.List(ctx, workspaceID, domain.EntityFilter{
|
|
Types: []domain.EntityType{e.Type},
|
|
}, domain.Page{Limit: 200})
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
for _, other := range candidates {
|
|
if other.ID == e.ID {
|
|
continue
|
|
}
|
|
comparisons++
|
|
if comparisons >= maxAliasComparisons {
|
|
break
|
|
}
|
|
sim := AliasSimilarity(e.CanonicalName, other.CanonicalName)
|
|
if sim < aliasProposalThreshold {
|
|
continue
|
|
}
|
|
// Record the proposal in a stable direction so the same pair does
|
|
// not generate two mirror-image proposals.
|
|
from, to := e, other
|
|
if from.ID > to.ID {
|
|
from, to = to, from
|
|
}
|
|
err := s.aliases.Upsert(ctx, &domain.EntityAlias{
|
|
ID: ids.New(),
|
|
WorkspaceID: workspaceID,
|
|
EntityID: from.ID,
|
|
Alias: to.CanonicalName,
|
|
Normalised: Normalise(to.CanonicalName),
|
|
CandidateEntityID: to.ID,
|
|
Status: domain.AliasCandidate,
|
|
Method: domain.MethodDeterministic,
|
|
Similarity: sim,
|
|
Reason: "These names are similar. Relata has NOT merged them: confirm only if they refer to the same thing.",
|
|
})
|
|
if err != nil {
|
|
log.Debug("could not record an alias proposal", "error", err.Error())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Claims
|
|
// ---------------------------------------------------------------------------
|
|
|
|
// ExtractClaimsForItem extracts the statements an item makes.
|
|
//
|
|
// Deterministic extraction picks out sentences carrying a figure or a date —
|
|
// the statements most likely to be compared, agreed with or contradicted.
|
|
// Model-assisted extraction adds the rest. Either way, a claim is stored only
|
|
// when its quotation can be located in the source text.
|
|
func (s *Service) ExtractClaimsForItem(ctx context.Context, item *domain.Item, chunks []*domain.Chunk, pool *jobs.Pool) (int, error) {
|
|
log := observability.LoggerFrom(ctx)
|
|
now := time.Now().UTC()
|
|
|
|
var out []*domain.Claim
|
|
seen := map[string]bool{}
|
|
|
|
add := func(c *domain.Claim) {
|
|
if c.DedupeKey == "" || seen[c.DedupeKey] {
|
|
return
|
|
}
|
|
seen[c.DedupeKey] = true
|
|
out = append(out, c)
|
|
}
|
|
|
|
// Stage 1: deterministic claims.
|
|
for _, ch := range chunks {
|
|
for _, dc := range ExtractDeterministicClaims(ch.Text) {
|
|
add(&domain.Claim{
|
|
ID: ids.New(),
|
|
WorkspaceID: item.WorkspaceID,
|
|
ItemID: item.ID,
|
|
ChunkID: ch.ID,
|
|
Statement: dc.Statement,
|
|
SourceRef: domain.SourceRef{
|
|
ItemID: item.ID,
|
|
ChunkID: ch.ID,
|
|
Start: ch.StartOffset + dc.Start,
|
|
End: ch.StartOffset + dc.End,
|
|
Quote: dc.Quote,
|
|
Page: ch.Page,
|
|
},
|
|
Nature: domain.ClaimQuoted,
|
|
Subject: dc.Subject,
|
|
Predicate: dc.Predicate,
|
|
Numbers: dc.Numbers,
|
|
Dates: dc.Dates,
|
|
Polarity: dc.Polarity,
|
|
Method: domain.MethodDeterministic,
|
|
Confidence: 0.7,
|
|
DedupeKey: ids.Deterministic(item.WorkspaceID, item.ID, "det", dc.Quote),
|
|
CreatedAt: now,
|
|
})
|
|
}
|
|
}
|
|
|
|
// Stage 2: model-assisted claims.
|
|
var modelErr error
|
|
if s.analysis.Enabled() {
|
|
limit := len(chunks)
|
|
if limit > maxModelChunksPerItem {
|
|
limit = maxModelChunksPerItem
|
|
}
|
|
for _, ch := range chunks[:limit] {
|
|
if err := ctx.Err(); err != nil {
|
|
modelErr = err
|
|
break
|
|
}
|
|
release, err := pool.AcquireAI(ctx)
|
|
if err != nil {
|
|
modelErr = err
|
|
break
|
|
}
|
|
cands, err := s.analysis.ExtractClaims(ctx, ch.Text)
|
|
release()
|
|
if err != nil {
|
|
modelErr = err
|
|
log.Debug("model claim extraction failed for a chunk", "error", err.Error())
|
|
continue
|
|
}
|
|
for _, cand := range cands {
|
|
// The analysis layer has already verified the quotation exists
|
|
// and returned where it is.
|
|
if cand.QuoteOffset < 0 {
|
|
continue
|
|
}
|
|
subj, pred, nums, dates, pol := AnalyseStatement(cand.Statement)
|
|
if cand.Subject != "" {
|
|
subj = Normalise(cand.Subject)
|
|
}
|
|
add(&domain.Claim{
|
|
ID: ids.New(),
|
|
WorkspaceID: item.WorkspaceID,
|
|
ItemID: item.ID,
|
|
ChunkID: ch.ID,
|
|
Statement: cand.Statement,
|
|
SourceRef: domain.SourceRef{
|
|
ItemID: item.ID,
|
|
ChunkID: ch.ID,
|
|
Start: ch.StartOffset + cand.QuoteOffset,
|
|
End: ch.StartOffset + cand.QuoteOffset + len(cand.Quote),
|
|
Quote: cand.Quote,
|
|
Page: ch.Page,
|
|
},
|
|
Attribution: cand.Attribution,
|
|
Nature: cand.Nature,
|
|
Subject: subj,
|
|
Predicate: pred,
|
|
Numbers: nums,
|
|
Dates: dates,
|
|
Polarity: pol,
|
|
Method: domain.MethodModel,
|
|
ModelVersion: s.analysis.ModelID(),
|
|
Confidence: 0.6,
|
|
DedupeKey: ids.Deterministic(item.WorkspaceID, item.ID, "model", cand.Quote),
|
|
CreatedAt: now,
|
|
})
|
|
}
|
|
}
|
|
} else {
|
|
modelErr = analysis.ErrDisabled
|
|
}
|
|
|
|
if err := s.claims.ReplaceForItem(ctx, item.WorkspaceID, item.ID, out); err != nil {
|
|
return 0, err
|
|
}
|
|
return len(out), modelErr
|
|
}
|