"""
Farnsworth Knowledge Graph v2 - Enhanced with Temporal Edges and Entity Resolution
Q1 2025 Feature: Enhanced Knowledge Graph
- Temporal edge tracking (relationships over time)
- Automated entity resolution and merging
- 3D graph visualization data support
- Relationship evolution tracking
"""
import asyncio
import json
import hashlib
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional, Any
from collections import defaultdict
import numpy as np
from loguru import logger
# Import base classes from existing knowledge graph
from farnsworth.memory.knowledge_graph import (
Entity, Relationship, GraphQuery, KnowledgeGraph
)
@dataclass
class TemporalEdge:
"""A relationship with temporal information."""
source_id: str
target_id: str
relation_type: str
# Temporal data
start_time: datetime
end_time: Optional[datetime] = None # None = still active
is_active: bool = True
# Evolution tracking
weight_history: list[tuple[datetime, float]] = field(default_factory=list)
current_weight: float = 1.0
# Metadata
evidence_timeline: list[dict] = field(default_factory=list) # Timestamped evidence
confidence: float = 0.5
stability_score: float = 0.5 # How stable this relationship is
def to_dict(self) -> dict:
return {
"source_id": self.source_id,
"target_id": self.target_id,
"relation_type": self.relation_type,
"start_time": self.start_time.isoformat(),
"end_time": self.end_time.isoformat() if self.end_time else None,
"is_active": self.is_active,
"weight_history": [
(t.isoformat(), w) for t, w in self.weight_history
],
"current_weight": self.current_weight,
"evidence_timeline": self.evidence_timeline,
"confidence": self.confidence,
"stability_score": self.stability_score,
}
@classmethod
def from_dict(cls, data: dict) -> "TemporalEdge":
weight_history = [
(datetime.fromisoformat(t), w) for t, w in data.get("weight_history", [])
]
return cls(
source_id=data["source_id"],
target_id=data["target_id"],
relation_type=data["relation_type"],
start_time=datetime.fromisoformat(data["start_time"]),
end_time=datetime.fromisoformat(data["end_time"]) if data.get("end_time") else None,
is_active=data.get("is_active", True),
weight_history=weight_history,
current_weight=data.get("current_weight", 1.0),
evidence_timeline=data.get("evidence_timeline", []),
confidence=data.get("confidence", 0.5),
stability_score=data.get("stability_score", 0.5),
)
@dataclass
class EntityCluster:
"""A cluster of entities that might be the same."""
id: str
canonical_id: str # The main entity ID
member_ids: list[str]
similarity_scores: dict[str, float] # member_id -> similarity to canonical
merged_at: Optional[datetime] = None
merge_confidence: float = 0.0
def to_dict(self) -> dict:
return {
"id": self.id,
"canonical_id": self.canonical_id,
"member_ids": self.member_ids,
"similarity_scores": self.similarity_scores,
"merged_at": self.merged_at.isoformat() if self.merged_at else None,
"merge_confidence": self.merge_confidence,
}
@dataclass
class EntityResolutionCandidate:
"""A candidate pair for entity resolution."""
entity1_id: str
entity2_id: str
similarity_score: float
match_reasons: list[str]
suggested_action: str # "merge", "link", "ignore"
@dataclass
class TemporalQuery:
"""Query with temporal constraints."""
query: str
time_point: Optional[datetime] = None # State at specific time
time_range: Optional[tuple[datetime, datetime]] = None # State over range
include_inactive: bool = False # Include ended relationships
track_evolution: bool = False # Include weight history
class KnowledgeGraphV2(KnowledgeGraph):
"""
Extended knowledge graph with temporal edges and entity resolution.
New Features:
- Temporal relationship tracking
- Entity resolution and merging
- Relationship evolution analysis
- 3D visualization data
"""
def __init__(
self,
data_dir: str = "./data/graph_v2",
max_nodes: int = 10000,
auto_link_threshold: float = 0.75,
resolution_threshold: float = 0.85,
):
super().__init__(data_dir, max_nodes, auto_link_threshold)
self.resolution_threshold = resolution_threshold
# Temporal edges storage
self.temporal_edges: dict[str, TemporalEdge] = {} # edge_id -> TemporalEdge
# Entity resolution
self.entity_clusters: dict[str, EntityCluster] = {}
self.canonical_mapping: dict[str, str] = {} # entity_id -> canonical_id
# Resolution candidates
self.resolution_queue: list[EntityResolutionCandidate] = []
# Edge indices
self.edges_by_source: dict[str, set[str]] = defaultdict(set)
self.edges_by_target: dict[str, set[str]] = defaultdict(set)
self.edges_by_time: dict[str, set[str]] = defaultdict(set) # "YYYY-MM" -> edge_ids
self._v2_initialized = False
async def initialize(self):
"""Initialize V2 features."""
await super().initialize()
if self._v2_initialized:
return
await self._load_temporal_data()
await self._load_resolution_data()
self._v2_initialized = True
logger.info(f"Knowledge Graph V2 initialized with {len(self.temporal_edges)} temporal edges")
async def add_temporal_relationship(
self,
source: str,
target: str,
relation_type: str,
weight: float = 1.0,
start_time: Optional[datetime] = None,
evidence: Optional[str] = None,
confidence: float = 0.5,
) -> TemporalEdge:
"""
Add a relationship with temporal tracking.
Args:
source: Source entity name or ID
target: Target entity name or ID
relation_type: Type of relationship
weight: Relationship strength
start_time: When relationship started (default: now)
evidence: Text evidence for the relationship
confidence: Confidence in the relationship
Returns:
TemporalEdge object
"""
async with self._lock:
source_id = self._resolve_entity(source)
target_id = self._resolve_entity(target)
if not source_id or not target_id:
raise ValueError(f"Entity not found: {source if not source_id else target}")
# Check for existing edge
edge_key = f"{source_id}|{target_id}|{relation_type}"
existing_edge_id = None
for eid, edge in self.temporal_edges.items():
if (edge.source_id == source_id and
edge.target_id == target_id and
edge.relation_type == relation_type and
edge.is_active):
existing_edge_id = eid
break
now = datetime.now()
start = start_time or now
if existing_edge_id:
# Update existing edge
edge = self.temporal_edges[existing_edge_id]
edge.weight_history.append((now, edge.current_weight))
edge.current_weight = (edge.current_weight + weight) / 2 # Rolling average
edge.confidence = (edge.confidence + confidence) / 2
if evidence:
edge.evidence_timeline.append({
"timestamp": now.isoformat(),
"evidence": evidence[:500],
})
# Update stability score
edge.stability_score = self._calculate_stability(edge)
await self._save_temporal_edge(edge)
return edge
# Create new edge
edge_id = f"te_{hashlib.md5(edge_key.encode()).hexdigest()[:12]}"
edge = TemporalEdge(
source_id=source_id,
target_id=target_id,
relation_type=relation_type,
start_time=start,
current_weight=weight,
weight_history=[(start, weight)],
confidence=confidence,
)
if evidence:
edge.evidence_timeline.append({
"timestamp": start.isoformat(),
"evidence": evidence[:500],
})
self.temporal_edges[edge_id] = edge
# Update indices
self.edges_by_source[source_id].add(edge_id)
self.edges_by_target[target_id].add(edge_id)
self.edges_by_time[start.strftime("%Y-%m")].add(edge_id)
# Also add to base graph
await super().add_relationship(source_id, target_id, relation_type, weight, evidence)
await self._save_temporal_edge(edge)
return edge
async def end_relationship(
self,
source: str,
target: str,
relation_type: str,
end_time: Optional[datetime] = None,
) -> Optional[TemporalEdge]:
"""Mark a relationship as ended."""
async with self._lock:
source_id = self._resolve_entity(source)
target_id = self._resolve_entity(target)
for eid, edge in self.temporal_edges.items():
if (edge.source_id == source_id and
edge.target_id == target_id and
edge.relation_type == relation_type and
edge.is_active):
edge.end_time = end_time or datetime.now()
edge.is_active = False
await self._save_temporal_edge(edge)
return edge
return None
async def query_at_time(
self,
query: str,
time_point: datetime,
max_entities: int = 10,
) -> GraphQuery:
"""Query the graph state at a specific point in time."""
# Get base query results
result = await super().query(query, max_entities)
# Filter relationships to those active at time_point
active_relationships = []
for rel in result.relationships:
# Find temporal edge
for edge in self.temporal_edges.values():
if (edge.source_id == rel.source_id and
edge.target_id == rel.target_id and
edge.relation_type == rel.relation_type):
# Check if active at time_point
if edge.start_time <= time_point:
if edge.end_time is None or edge.end_time >= time_point:
active_relationships.append(rel)
break
result.relationships = active_relationships
return result
async def get_relationship_evolution(
self,
source: str,
target: str,
relation_type: Optional[str] = None,
) -> list[dict]:
"""Get the evolution of a relationship over time."""
source_id = self._resolve_entity(source)
target_id = self._resolve_entity(target)
evolutions = []
for edge in self.temporal_edges.values():
if edge.source_id == source_id and edge.target_id == target_id:
if relation_type is None or edge.relation_type == relation_type:
evolutions.append({
"relation_type": edge.relation_type,
"start_time": edge.start_time.isoformat(),
"end_time": edge.end_time.isoformat() if edge.end_time else None,
"is_active": edge.is_active,
"weight_history": [
{"time": t.isoformat(), "weight": w}
for t, w in edge.weight_history
],
"current_weight": edge.current_weight,
"stability_score": edge.stability_score,
})
return evolutions
def _calculate_stability(self, edge: TemporalEdge) -> float:
"""Calculate how stable a relationship is over time."""
if len(edge.weight_history) < 2:
return 0.5
weights = [w for _, w in edge.weight_history]
# Calculate variance
variance = np.var(weights)
# Lower variance = more stable
stability = 1.0 / (1.0 + variance)
# Bonus for longevity
if edge.start_time:
days_active = (datetime.now() - edge.start_time).days
longevity_bonus = min(0.2, days_active / 365 * 0.2)
stability += longevity_bonus
return min(1.0, stability)
# Entity Resolution Methods
async def find_resolution_candidates(
self,
min_similarity: Optional[float] = None,
) -> list[EntityResolutionCandidate]:
"""Find entities that might be the same."""
threshold = min_similarity or self.resolution_threshold
candidates = []
entities_list = list(self.entities.values())
for i, e1 in enumerate(entities_list):
for e2 in entities_list[i + 1:]:
similarity, reasons = await self._calculate_entity_similarity(e1, e2)
if similarity >= threshold:
candidates.append(EntityResolutionCandidate(
entity1_id=e1.id,
entity2_id=e2.id,
similarity_score=similarity,
match_reasons=reasons,
suggested_action="merge" if similarity >= 0.95 else "link",
))
# Sort by similarity
candidates.sort(key=lambda c: c.similarity_score, reverse=True)
self.resolution_queue = candidates
return candidates
async def _calculate_entity_similarity(
self,
e1: Entity,
e2: Entity,
) -> tuple[float, list[str]]:
"""Calculate similarity between two entities."""
scores = []
reasons = []
# Name similarity
name_sim = self._string_similarity(e1.name.lower(), e2.name.lower())
if name_sim > 0.8:
scores.append(name_sim)
reasons.append(f"Similar names ({name_sim:.2f})")
# Type match
if e1.entity_type == e2.entity_type:
scores.append(0.3)
reasons.append("Same entity type")
# Embedding similarity
if e1.embedding and e2.embedding:
emb_sim = self._cosine_similarity(
np.array(e1.embedding),
np.array(e2.embedding)
)
if emb_sim > 0.7:
scores.append(emb_sim)
reasons.append(f"Similar embeddings ({emb_sim:.2f})")
# Property overlap
if e1.properties and e2.properties:
common_keys = set(e1.properties.keys()) & set(e2.properties.keys())
if common_keys:
prop_matches = sum(
1 for k in common_keys
if e1.properties[k] == e2.properties[k]
)
if prop_matches > 0:
scores.append(prop_matches / len(common_keys) * 0.5)
reasons.append(f"Matching properties ({prop_matches})")
overall = np.mean(scores) if scores else 0.0
return float(overall), reasons
def _string_similarity(self, s1: str, s2: str) -> float:
"""Calculate string similarity using Levenshtein-like approach."""
if s1 == s2:
return 1.0
# Simple character-level Jaccard
set1 = set(s1)
set2 = set(s2)
intersection = len(set1 & set2)
union = len(set1 | set2)
return intersection / union if union > 0 else 0.0
def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float:
"""Compute cosine similarity."""
norm_a = np.linalg.norm(a)
norm_b = np.linalg.norm(b)
if norm_a == 0 or norm_b == 0:
return 0.0
return float(np.dot(a, b) / (norm_a * norm_b))
async def merge_entities(
self,
entity1_id: str,
entity2_id: str,
keep_id: Optional[str] = None,
) -> Entity:
"""
Merge two entities into one.
Args:
entity1_id: First entity ID
entity2_id: Second entity ID
keep_id: Which ID to keep as canonical (default: older one)
Returns:
Merged entity
"""
async with self._lock:
e1 = self.entities.get(entity1_id)
e2 = self.entities.get(entity2_id)
if not e1 or not e2:
raise ValueError("Entity not found")
# Determine canonical
if keep_id:
canonical = e1 if e1.id == keep_id else e2
secondary = e2 if e1.id == keep_id else e1
else:
# Keep the older one
canonical = e1 if e1.created_at <= e2.created_at else e2
secondary = e2 if e1.created_at <= e2.created_at else e1
# Merge properties
for key, value in secondary.properties.items():
if key not in canonical.properties:
canonical.properties[key] = value
# Merge mention count
canonical.mention_count += secondary.mention_count
# Update last mentioned
canonical.last_mentioned = max(
canonical.last_mentioned,
secondary.last_mentioned
)
# Merge embedding (average)
if canonical.embedding and secondary.embedding:
canonical.embedding = (
(np.array(canonical.embedding) + np.array(secondary.embedding)) / 2
).tolist()
# Update relationships to point to canonical
await self._redirect_relationships(secondary.id, canonical.id)
# Create cluster record
cluster_id = f"cluster_{canonical.id}"
cluster = EntityCluster(
id=cluster_id,
canonical_id=canonical.id,
member_ids=[canonical.id, secondary.id],
similarity_scores={secondary.id: 1.0},
merged_at=datetime.now(),
merge_confidence=1.0,
)
self.entity_clusters[cluster_id] = cluster
self.canonical_mapping[secondary.id] = canonical.id
# Remove secondary entity
del self.entities[secondary.id]
self.name_to_id.pop(secondary.name.lower(), None)
await self._save_entity(canonical)
await self._save_resolution_data()
logger.info(f"Merged entity {secondary.id} into {canonical.id}")
return canonical
async def _redirect_relationships(self, old_id: str, new_id: str):
"""Redirect all relationships from old entity to new."""
# Update temporal edges
for edge_id, edge in list(self.temporal_edges.items()):
if edge.source_id == old_id:
edge.source_id = new_id
self.edges_by_source[old_id].discard(edge_id)
self.edges_by_source[new_id].add(edge_id)
if edge.target_id == old_id:
edge.target_id = new_id
self.edges_by_target[old_id].discard(edge_id)
self.edges_by_target[new_id].add(edge_id)
await self._save_temporal_edge(edge)
# Update base graph
if self.graph is not None:
# Get all edges involving old_id
if old_id in self.graph:
# Redirect outgoing edges
for _, target, data in list(self.graph.out_edges(old_id, data=True)):
self.graph.add_edge(new_id, target, **data)
# Redirect incoming edges
for source, _, data in list(self.graph.in_edges(old_id, data=True)):
self.graph.add_edge(source, new_id, **data)
self.graph.remove_node(old_id)
def get_canonical_id(self, entity_id: str) -> str:
"""Get the canonical ID for an entity (handling merges)."""
return self.canonical_mapping.get(entity_id, entity_id)
# 3D Visualization Data
def get_3d_graph_data(
self,
max_nodes: int = 100,
include_inactive: bool = False,
) -> dict:
"""
Get graph data formatted for 3D visualization.
Returns data suitable for force-directed 3D graph libraries.
"""
nodes = []
links = []
# Get most important entities
sorted_entities = sorted(
self.entities.values(),
key=lambda e: e.mention_count,
reverse=True
)[:max_nodes]
entity_ids = {e.id for e in sorted_entities}
# Format nodes
for entity in sorted_entities:
nodes.append({
"id": entity.id,
"name": entity.name,
"type": entity.entity_type,
"size": min(20, 5 + entity.mention_count), # Node size
"color": self._get_type_color(entity.entity_type),
"properties": entity.properties,
})
# Format links from temporal edges
for edge in self.temporal_edges.values():
if edge.source_id in entity_ids and edge.target_id in entity_ids:
if include_inactive or edge.is_active:
links.append({
"source": edge.source_id,
"target": edge.target_id,
"type": edge.relation_type,
"weight": edge.current_weight,
"active": edge.is_active,
"stability": edge.stability_score,
"start_time": edge.start_time.isoformat(),
"end_time": edge.end_time.isoformat() if edge.end_time else None,
})
return {
"nodes": nodes,
"links": links,
"metadata": {
"total_entities": len(self.entities),
"total_edges": len(self.temporal_edges),
"displayed_entities": len(nodes),
"displayed_edges": len(links),
}
}
def _get_type_color(self, entity_type: str) -> str:
"""Get color for entity type in visualization."""
colors = {
"person": "#4CAF50",
"concept": "#2196F3",
"code": "#FF9800",
"file": "#9C27B0",
"url": "#00BCD4",
"tool": "#F44336",
"default": "#757575",
}
return colors.get(entity_type.lower(), colors["default"])
# Persistence
async def _save_temporal_edge(self, edge: TemporalEdge):
"""Save temporal edge to disk."""
edges_dir = self.data_dir / "temporal_edges"
edges_dir.mkdir(parents=True, exist_ok=True)
edge_id = hashlib.md5(
f"{edge.source_id}|{edge.target_id}|{edge.relation_type}".encode()
).hexdigest()[:12]
edge_file = edges_dir / f"{edge_id}.json"
edge_file.write_text(json.dumps(edge.to_dict(), indent=2), encoding='utf-8')
async def _save_entity(self, entity: Entity):
"""Save entity to disk."""
entities_dir = self.data_dir / "entities"
entities_dir.mkdir(parents=True, exist_ok=True)
entity_file = entities_dir / f"{entity.id}.json"
entity_file.write_text(json.dumps(entity.to_dict(), indent=2), encoding='utf-8')
async def _save_resolution_data(self):
"""Save entity resolution data."""
resolution_file = self.data_dir / "resolution.json"
data = {
"clusters": {k: v.to_dict() for k, v in self.entity_clusters.items()},
"canonical_mapping": self.canonical_mapping,
}
resolution_file.write_text(json.dumps(data, indent=2), encoding='utf-8')
async def _load_temporal_data(self):
"""Load temporal edges from disk."""
edges_dir = self.data_dir / "temporal_edges"
if not edges_dir.exists():
return
for edge_file in edges_dir.glob("*.json"):
try:
data = json.loads(edge_file.read_text(encoding='utf-8'))
edge = TemporalEdge.from_dict(data)
edge_id = edge_file.stem
self.temporal_edges[edge_id] = edge
self.edges_by_source[edge.source_id].add(edge_id)
self.edges_by_target[edge.target_id].add(edge_id)
self.edges_by_time[edge.start_time.strftime("%Y-%m")].add(edge_id)
except Exception as e:
logger.error(f"Failed to load temporal edge {edge_file}: {e}")
async def _load_resolution_data(self):
"""Load entity resolution data."""
resolution_file = self.data_dir / "resolution.json"
if not resolution_file.exists():
return
try:
data = json.loads(resolution_file.read_text(encoding='utf-8'))
for k, v in data.get("clusters", {}).items():
self.entity_clusters[k] = EntityCluster(**{
**v,
"merged_at": datetime.fromisoformat(v["merged_at"]) if v.get("merged_at") else None,
})
self.canonical_mapping = data.get("canonical_mapping", {})
except Exception as e:
logger.error(f"Failed to load resolution data: {e}")
def get_stats(self) -> dict:
"""Get extended statistics."""
base_stats = super().get_stats()
base_stats.update({
"temporal_edges": len(self.temporal_edges),
"active_edges": sum(1 for e in self.temporal_edges.values() if e.is_active),
"entity_clusters": len(self.entity_clusters),
"merged_entities": len(self.canonical_mapping),
"resolution_candidates": len(self.resolution_queue),
})
return base_stats
