1"""Entity extractor — parses subject/predicate/object triples from text."""
2
3from __future__ import annotations
4
5from typing import TYPE_CHECKING
6
7from lexigram.ai.memory.exceptions import FactExtractionError
8from lexigram.logging import (
9 get_logger,
10)
11
12if TYPE_CHECKING:
13 from collections.abc import Awaitable, Callable
14
15 from lexigram.contracts.ai.memory import MemoryEntry
16
17logger = get_logger(__name__)
18
19# Simple pattern-based fallback triples (subject, predicate, object_)
20Triple = tuple[str, str, str]
21
22
23class EntityExtractor:
24 """Extracts structured subject/predicate/object triples from memory entries.
25
26 When an LLM extract callable is injected, delegation occurs there.
27 Otherwise a lightweight heuristic fallback is used, suitable for
28 testing and environments without LLM access.
29 """
30
31 def __init__(
32 self,
33 extract_fn: Callable[[str], Awaitable[list[Triple]]] | None = None,
34 ) -> None:
35 """Initialise the extractor.
36
37 Args:
38 extract_fn: Async callable that returns triples from raw text.
39 When *None*, a heuristic fallback is used.
40 """
41 self._extract_fn = extract_fn
42
43 async def extract(self, entry: MemoryEntry) -> list[Triple]:
44 """Extract triples from *entry.content*.
45
46 Args:
47 entry: Memory entry whose content to parse.
48
49 Returns:
50 List of (subject, predicate, object_) tuples.
51
52 Raises:
53 FactExtractionError: If the extraction callable raises.
54 """
55 if self._extract_fn:
56 try:
57 return await self._extract_fn(entry.content)
58 except (RuntimeError, ValueError, TypeError) as exc:
59 raise FactExtractionError(
60 f"Extraction failed for entry {entry.id}"
61 ) from exc
62 return self._heuristic_extract(entry)
63
64 def _heuristic_extract(self, entry: MemoryEntry) -> list[Triple]:
65 """Naive heuristic — extracts 'X is Y' and 'X has Y' patterns."""
66 triples: list[Triple] = []
67 content = entry.content
68 for phrase in content.split("."):
69 phrase = phrase.strip()
70 for sep in (" is ", " are ", " has ", " have ", " was ", " were "):
71 if sep in phrase:
72 parts = phrase.split(sep, 1)
73 if len(parts) == 2 and parts[0] and parts[1]:
74 subject = parts[0].strip().lower()
75 predicate = sep.strip()
76 object_ = parts[1].strip().lower()
77 if len(subject) <= 60 and len(object_) <= 120:
78 triples.append((subject, predicate, object_))
79 break
80 return triples
81
82
83# Avoid circular import for type hints
84
85__all__ = ["EntityExtractor"]