In-memory entity store with heuristic extraction + JSON persistence.
| 33 | |
| 34 | |
| 35 | class KnowledgeGraph: |
| 36 | """In-memory entity store with heuristic extraction + JSON persistence.""" |
| 37 | |
| 38 | def __init__(self) -> None: |
| 39 | self.entities: dict[str, Entity] = {} # key: f"{type}:{name}" |
| 40 | |
| 41 | @staticmethod |
| 42 | def _key(name: str, etype: str) -> str: |
| 43 | return f"{etype}:{name}" |
| 44 | |
| 45 | def add(self, name: str, etype: str, now: float = 0.0) -> None: |
| 46 | name = name.strip() |
| 47 | if not name: |
| 48 | return |
| 49 | key = self._key(name, etype) |
| 50 | ent = self.entities.get(key) |
| 51 | if ent is None: |
| 52 | self.entities[key] = Entity(name=name, type=etype, count=1, last_seen=now) |
| 53 | else: |
| 54 | ent.count += 1 |
| 55 | if now: |
| 56 | ent.last_seen = now |
| 57 | |
| 58 | def record_from_text(self, text: str, now: float = 0.0) -> int: |
| 59 | """Extract entities from ``text``; returns the number of mentions recorded.""" |
| 60 | if not text: |
| 61 | return 0 |
| 62 | n = 0 |
| 63 | for m in _FILE_RE.finditer(text): |
| 64 | self.add(m.group(1), "file", now) |
| 65 | n += 1 |
| 66 | for m in _SYMBOL_RE.finditer(text): |
| 67 | self.add(m.group(1), "symbol", now) |
| 68 | n += 1 |
| 69 | for m in _URL_RE.finditer(text): |
| 70 | self.add(m.group(0), "url", now) |
| 71 | n += 1 |
| 72 | return n |
| 73 | |
| 74 | def stats(self) -> dict[str, int]: |
| 75 | by_type: dict[str, int] = {} |
| 76 | for ent in self.entities.values(): |
| 77 | by_type[ent.type] = by_type.get(ent.type, 0) + 1 |
| 78 | return {"total": len(self.entities), **by_type} |
| 79 | |
| 80 | def top(self, limit: int = 20) -> list[Entity]: |
| 81 | return sorted(self.entities.values(), key=lambda e: (-e.count, -e.last_seen, e.name))[:limit] |
| 82 | |
| 83 | def clear(self) -> None: |
| 84 | self.entities.clear() |
| 85 | |
| 86 | # — persistence — |
| 87 | def to_dict(self) -> dict: |
| 88 | return {"entities": [asdict(e) for e in self.entities.values()]} |
| 89 | |
| 90 | @classmethod |
| 91 | def from_dict(cls, data: dict) -> "KnowledgeGraph": |
| 92 | g = cls() |
no outgoing calls