weavatrix_memory/extraction/linker/
mod.rs

1use super::{
2    ExtractedEntity, ExtractionInput, LinkCandidate, LinkDecision, LinkMethod, LinkPolicy,
3    key::normalized,
4};
5use crate::{
6    domain::Confidence,
7    error::{MemoryError, Result},
8    id::EntityId,
9};
10use std::collections::{BTreeMap, HashMap};
11
12mod catalog;
13mod scoring;
14use catalog::IndexBucket;
15use scoring::{decision, exact_key, generated_id, scope_score};
16
17#[derive(Debug, Clone)]
18pub struct EntityLinker {
19    nodes: Vec<CatalogEntity>,
20    by_id: HashMap<EntityId, usize>,
21    by_label: HashMap<(String, String), IndexBucket>,
22    by_alias: HashMap<(String, String), IndexBucket>,
23    by_external_id: HashMap<(String, String, String), IndexBucket>,
24}
25
26#[derive(Debug, Clone)]
27struct CatalogEntity {
28    id: EntityId,
29    kind: String,
30    repository: Option<String>,
31    branch: Option<String>,
32}
33
34impl EntityLinker {
35    /// Resolves one extracted mention without silently choosing ambiguity.
36    ///
37    /// # Errors
38    ///
39    /// Rejects stable identifiers whose existing entity has a different kind,
40    /// generated identifier collisions, or invalid policy scores.
41    pub fn link(
42        &self,
43        mention: &ExtractedEntity,
44        input: &ExtractionInput,
45        policy: LinkPolicy,
46    ) -> Result<LinkDecision> {
47        if policy.minimum_score > 10_000 || policy.minimum_margin > 10_000 {
48            return Err(MemoryError::InvalidValue {
49                field: "link_policy",
50                reason: "scores must be between 0 and 10,000 basis points",
51            });
52        }
53        if let Some(stable_id) = &mention.stable_id {
54            return self.link_stable_id(mention, stable_id);
55        }
56        let candidates = self.collect_candidates(mention, input)?;
57        self.resolve_candidates(mention, input, policy, candidates)
58    }
59
60    fn link_stable_id(
61        &self,
62        mention: &ExtractedEntity,
63        stable_id: &EntityId,
64    ) -> Result<LinkDecision> {
65        if let Some(index) = self.by_id.get(stable_id) {
66            if self.nodes[*index].kind != normalized(&mention.kind) {
67                return Err(MemoryError::InvalidValue {
68                    field: "extracted_entity.stable_id",
69                    reason: "existing entity kind does not match extracted kind",
70                });
71            }
72            let score = mention.confidence;
73            return Ok(decision(
74                mention,
75                Some(stable_id.clone()),
76                score,
77                LinkMethod::StableId,
78                vec![LinkCandidate {
79                    entity_id: stable_id.clone(),
80                    score,
81                    method: LinkMethod::StableId,
82                }],
83            ));
84        }
85        Ok(decision(
86            mention,
87            Some(stable_id.clone()),
88            mention.confidence,
89            LinkMethod::Created,
90            Vec::new(),
91        ))
92    }
93
94    fn collect_candidates(
95        &self,
96        mention: &ExtractedEntity,
97        input: &ExtractionInput,
98    ) -> Result<Vec<LinkCandidate>> {
99        let kind = normalized(&mention.kind);
100        let mut candidates = BTreeMap::<EntityId, LinkCandidate>::new();
101        self.add_index_matches(
102            self.by_label
103                .get(&(kind.clone(), normalized(&mention.label))),
104            mention,
105            input,
106            9_000,
107            LinkMethod::Label,
108            &mut candidates,
109        )?;
110        self.add_index_matches(
111            self.by_alias
112                .get(&(kind.clone(), normalized(&mention.label))),
113            mention,
114            input,
115            8_500,
116            LinkMethod::Alias,
117            &mut candidates,
118        )?;
119        for alias in &mention.aliases {
120            self.add_index_matches(
121                self.by_alias.get(&(kind.clone(), normalized(alias))),
122                mention,
123                input,
124                8_500,
125                LinkMethod::Alias,
126                &mut candidates,
127            )?;
128        }
129        for (key, value) in &mention.attributes {
130            if key.starts_with("external_id.") {
131                self.add_index_matches(
132                    self.by_external_id
133                        .get(&(kind.clone(), key.clone(), exact_key(value))),
134                    mention,
135                    input,
136                    9_700,
137                    LinkMethod::ExternalId,
138                    &mut candidates,
139                )?;
140            }
141        }
142        for hint in &mention.hints {
143            if let Some(index) = self.by_id.get(&hint.entity_id)
144                && self.nodes[*index].kind == kind
145            {
146                self.add_candidate(
147                    *index,
148                    mention,
149                    input,
150                    hint.confidence.basis_points(),
151                    LinkMethod::ProviderHint,
152                    &mut candidates,
153                )?;
154            }
155        }
156        let mut candidates = candidates.into_values().collect::<Vec<_>>();
157        candidates.sort_by(|left, right| {
158            right
159                .score
160                .cmp(&left.score)
161                .then_with(|| left.entity_id.cmp(&right.entity_id))
162        });
163        Ok(candidates)
164    }
165
166    #[allow(clippy::too_many_arguments)]
167    fn add_index_matches(
168        &self,
169        indexes: Option<&IndexBucket>,
170        mention: &ExtractedEntity,
171        input: &ExtractionInput,
172        base_score: u16,
173        method: LinkMethod,
174        candidates: &mut BTreeMap<EntityId, LinkCandidate>,
175    ) -> Result<()> {
176        if let Some(indexes) = indexes {
177            match indexes {
178                IndexBucket::One(index) => {
179                    self.add_candidate(*index, mention, input, base_score, method, candidates)?;
180                }
181                IndexBucket::Many(indexes) => {
182                    for index in indexes {
183                        self.add_candidate(*index, mention, input, base_score, method, candidates)?;
184                    }
185                }
186            }
187        }
188        Ok(())
189    }
190
191    #[allow(clippy::too_many_arguments)]
192    fn add_candidate(
193        &self,
194        index: usize,
195        mention: &ExtractedEntity,
196        input: &ExtractionInput,
197        base_score: u16,
198        method: LinkMethod,
199        candidates: &mut BTreeMap<EntityId, LinkCandidate>,
200    ) -> Result<()> {
201        let node = &self.nodes[index];
202        let (score, scoped) = scope_score(base_score, node, input);
203        let score = score.min(mention.confidence.basis_points());
204        let method = if method == LinkMethod::Label && scoped {
205            LinkMethod::ScopedLabel
206        } else {
207            method
208        };
209        let candidate = LinkCandidate {
210            entity_id: node.id.clone(),
211            score: Confidence::from_basis_points(score)?,
212            method,
213        };
214        match candidates.get(&node.id) {
215            Some(existing) if existing.score >= candidate.score => {}
216            _ => {
217                candidates.insert(node.id.clone(), candidate);
218            }
219        }
220        Ok(())
221    }
222
223    fn resolve_candidates(
224        &self,
225        mention: &ExtractedEntity,
226        input: &ExtractionInput,
227        policy: LinkPolicy,
228        candidates: Vec<LinkCandidate>,
229    ) -> Result<LinkDecision> {
230        let Some(best) = candidates.first() else {
231            return self.unmatched(mention, input, policy, candidates);
232        };
233        if best.score.basis_points() < policy.minimum_score {
234            return self.unmatched(mention, input, policy, candidates);
235        }
236        if candidates.get(1).is_some_and(|second| {
237            best.score
238                .basis_points()
239                .saturating_sub(second.score.basis_points())
240                < policy.minimum_margin
241        }) {
242            return Ok(decision(
243                mention,
244                None,
245                best.score,
246                LinkMethod::Ambiguous,
247                candidates,
248            ));
249        }
250        Ok(decision(
251            mention,
252            Some(best.entity_id.clone()),
253            best.score,
254            best.method,
255            candidates,
256        ))
257    }
258
259    fn unmatched(
260        &self,
261        mention: &ExtractedEntity,
262        input: &ExtractionInput,
263        policy: LinkPolicy,
264        candidates: Vec<LinkCandidate>,
265    ) -> Result<LinkDecision> {
266        if !policy.create_unmatched {
267            return Ok(decision(
268                mention,
269                None,
270                Confidence::from_basis_points(0)?,
271                LinkMethod::Unresolved,
272                candidates,
273            ));
274        }
275        let id = generated_id(mention, input)?;
276        if self.by_id.contains_key(&id) {
277            return Err(MemoryError::Extraction {
278                provider: "entity-linker".to_owned(),
279                message: format!("generated entity identifier collides with {id}"),
280            });
281        }
282        Ok(decision(
283            mention,
284            Some(id),
285            mention.confidence,
286            LinkMethod::Created,
287            candidates,
288        ))
289    }
290}
weavatrix_memory/extraction/linker/mod.rs

weavatrix_memory/extraction/linker/
mod.rs
