mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-29 04:26:20 +00:00
698 lines
22 KiB
Python
698 lines
22 KiB
Python
"""
|
|
Context Graph Implementation
|
|
|
|
This module provides a synchronous, in-memory implementation of the GraphStore protocol,
|
|
designed for building and querying context graphs from conversations and entities.
|
|
|
|
It formalizes context as a graph of connections, enabling meaningful connections between
|
|
concepts, entities, and conversations.
|
|
|
|
Key Features:
|
|
- In-memory GraphStore implementation
|
|
- Entity and relationship extraction from conversations
|
|
- BFS-based neighbor discovery
|
|
- Type-based indexing
|
|
- Export to dictionary format
|
|
"""
|
|
|
|
from collections import defaultdict, deque
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List, Optional, Set, Union
|
|
|
|
from ..utils.logging import get_logger
|
|
from ..utils.progress_tracker import get_progress_tracker
|
|
from .entity_linker import EntityLinker
|
|
|
|
|
|
@dataclass
|
|
class ContextNode:
|
|
"""Context graph node (Internal implementation)."""
|
|
|
|
node_id: str
|
|
node_type: str
|
|
content: str
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
properties: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to dictionary format."""
|
|
props = self.properties.copy()
|
|
props.update(self.metadata)
|
|
props["content"] = self.content
|
|
return {"id": self.node_id, "type": self.node_type, "properties": props}
|
|
|
|
|
|
@dataclass
|
|
class ContextEdge:
|
|
"""Context graph edge (Internal implementation)."""
|
|
|
|
source_id: str
|
|
target_id: str
|
|
edge_type: str
|
|
weight: float = 1.0
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to dictionary format."""
|
|
return {
|
|
"source_id": self.source_id,
|
|
"target_id": self.target_id,
|
|
"type": self.edge_type,
|
|
"weight": self.weight,
|
|
"properties": self.metadata,
|
|
}
|
|
|
|
|
|
class ContextGraph:
|
|
"""
|
|
In-memory implementation of context graph.
|
|
|
|
Provides capabilities to build, store, and query a context graph.
|
|
"""
|
|
|
|
def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs):
|
|
"""
|
|
Initialize context graph.
|
|
|
|
Args:
|
|
config: Configuration dictionary
|
|
**kwargs: Additional configuration options:
|
|
- extract_entities: Extract entities from content (default: True)
|
|
- extract_relationships: Extract relationships (default: True)
|
|
- entity_linker: Entity linker instance
|
|
"""
|
|
self.logger = get_logger("context_graph")
|
|
self.config = config or {}
|
|
self.config.update(kwargs)
|
|
|
|
self.extract_entities = self.config.get("extract_entities", True)
|
|
self.extract_relationships = self.config.get("extract_relationships", True)
|
|
|
|
self.entity_linker = self.config.get("entity_linker") or EntityLinker()
|
|
|
|
# Graph structure
|
|
self.nodes: Dict[str, ContextNode] = {}
|
|
self.edges: List[ContextEdge] = []
|
|
|
|
# Adjacency list for efficient traversal: source_id -> list of edges
|
|
self._adjacency: Dict[str, List[ContextEdge]] = defaultdict(list)
|
|
|
|
# Indexes
|
|
self.node_type_index: Dict[str, Set[str]] = defaultdict(set)
|
|
self.edge_type_index: Dict[str, List[ContextEdge]] = defaultdict(list)
|
|
|
|
# Progress tracker
|
|
self.progress_tracker = get_progress_tracker()
|
|
|
|
# --- GraphStore Protocol Implementation ---
|
|
|
|
def add_nodes(self, nodes: List[Dict[str, Any]]) -> int:
|
|
"""
|
|
Add nodes to graph.
|
|
|
|
Args:
|
|
nodes: List of nodes to add (dicts with id, type, properties)
|
|
|
|
Returns:
|
|
Number of nodes added
|
|
"""
|
|
count = 0
|
|
for node in nodes:
|
|
# Extract content from properties if not explicit
|
|
node_props = node.get("properties", {})
|
|
content = node_props.get("content", node.get("id"))
|
|
metadata = {k: v for k, v in node_props.items() if k != "content"}
|
|
|
|
internal_node = ContextNode(
|
|
node_id=node.get("id"),
|
|
node_type=node.get("type", "entity"),
|
|
content=content,
|
|
metadata=metadata,
|
|
properties=node_props,
|
|
)
|
|
|
|
if self._add_internal_node(internal_node):
|
|
count += 1
|
|
return count
|
|
|
|
def add_edges(self, edges: List[Dict[str, Any]]) -> int:
|
|
"""
|
|
Add edges to graph.
|
|
|
|
Args:
|
|
edges: List of edges to add (dicts with source_id, target_id, type,
|
|
weight, properties)
|
|
|
|
Returns:
|
|
Number of edges added
|
|
"""
|
|
count = 0
|
|
for edge in edges:
|
|
internal_edge = ContextEdge(
|
|
source_id=edge.get("source_id"),
|
|
target_id=edge.get("target_id"),
|
|
edge_type=edge.get("type", "related_to"),
|
|
weight=edge.get("weight", 1.0),
|
|
metadata=edge.get("properties", {}),
|
|
)
|
|
|
|
if self._add_internal_edge(internal_edge):
|
|
count += 1
|
|
return count
|
|
|
|
def get_neighbors(self, node_id: str, hops: int = 1) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get neighbors of a node.
|
|
|
|
Returns list of dicts with neighbor info.
|
|
"""
|
|
if node_id not in self.nodes:
|
|
return []
|
|
|
|
neighbors = []
|
|
visited = {node_id}
|
|
queue = deque([(node_id, 0)]) # (current_id, current_hop)
|
|
|
|
while queue:
|
|
current_id, current_hop = queue.popleft()
|
|
|
|
if current_hop >= hops:
|
|
continue
|
|
|
|
# Get outgoing edges
|
|
outgoing_edges = self._adjacency.get(current_id, [])
|
|
for edge in outgoing_edges:
|
|
neighbor_id = edge.target_id
|
|
if neighbor_id not in visited:
|
|
visited.add(neighbor_id)
|
|
queue.append((neighbor_id, current_hop + 1))
|
|
|
|
if neighbor_id in self.nodes:
|
|
node = self.nodes[neighbor_id]
|
|
neighbors.append(
|
|
{
|
|
"id": node.node_id,
|
|
"type": node.node_type,
|
|
"content": node.content,
|
|
"relationship": edge.edge_type,
|
|
"weight": edge.weight,
|
|
"hop": current_hop + 1,
|
|
}
|
|
)
|
|
|
|
return neighbors
|
|
|
|
def query(self, query: str) -> List[Dict[str, Any]]:
|
|
"""
|
|
Execute a simple keyword search query on the graph nodes.
|
|
|
|
Args:
|
|
query: Keyword query string
|
|
|
|
Returns:
|
|
List of matching node dicts
|
|
"""
|
|
results = []
|
|
query_lower = query.lower().split()
|
|
|
|
for node in self.nodes.values():
|
|
content_lower = node.content.lower()
|
|
if any(word in content_lower for word in query_lower):
|
|
# Calculate simple score
|
|
overlap = sum(1 for word in query_lower if word in content_lower)
|
|
score = overlap / len(query_lower) if query_lower else 0.0
|
|
|
|
results.append(
|
|
{
|
|
"node": node.to_dict(),
|
|
"score": score,
|
|
"content": node.content,
|
|
}
|
|
)
|
|
|
|
return sorted(results, key=lambda x: x["score"], reverse=True)
|
|
|
|
def add_node(
|
|
self,
|
|
node_id: str,
|
|
node_type: str,
|
|
content: Optional[str] = None,
|
|
**properties,
|
|
) -> bool:
|
|
"""
|
|
Add a single node to the graph.
|
|
|
|
Args:
|
|
node_id: Unique identifier
|
|
node_type: Node type (e.g., 'entity', 'concept')
|
|
content: Node content/label
|
|
**properties: Additional properties
|
|
"""
|
|
content = content or node_id
|
|
return self._add_internal_node(
|
|
ContextNode(
|
|
node_id=node_id,
|
|
node_type=node_type,
|
|
content=content,
|
|
metadata=properties,
|
|
properties=properties,
|
|
)
|
|
)
|
|
|
|
def add_edge(
|
|
self,
|
|
source_id: str,
|
|
target_id: str,
|
|
edge_type: str = "related_to",
|
|
weight: float = 1.0,
|
|
**properties,
|
|
) -> bool:
|
|
"""
|
|
Add a single edge to the graph.
|
|
|
|
Args:
|
|
source_id: Source node ID
|
|
target_id: Target node ID
|
|
edge_type: Relationship type
|
|
weight: Edge weight
|
|
**properties: Additional properties
|
|
"""
|
|
return self._add_internal_edge(
|
|
ContextEdge(
|
|
source_id=source_id,
|
|
target_id=target_id,
|
|
edge_type=edge_type,
|
|
weight=weight,
|
|
metadata=properties,
|
|
)
|
|
)
|
|
|
|
def save_to_file(self, path: str) -> None:
|
|
"""
|
|
Save context graph to file (JSON format).
|
|
|
|
Args:
|
|
path: File path to save to
|
|
"""
|
|
import json
|
|
|
|
data = {
|
|
"nodes": [node.to_dict() for node in self.nodes.values()],
|
|
"edges": [edge.to_dict() for edge in self.edges],
|
|
}
|
|
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=2, ensure_ascii=False)
|
|
|
|
self.logger.info(f"Saved context graph to {path}")
|
|
|
|
def load_from_file(self, path: str) -> None:
|
|
"""
|
|
Load context graph from file (JSON format).
|
|
|
|
Args:
|
|
path: File path to load from
|
|
"""
|
|
import json
|
|
import os
|
|
|
|
if not os.path.exists(path):
|
|
self.logger.warning(f"File not found: {path}")
|
|
return
|
|
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
# Clear existing
|
|
self.nodes.clear()
|
|
self.edges.clear()
|
|
self._adjacency.clear()
|
|
self.node_type_index.clear()
|
|
self.edge_type_index.clear()
|
|
|
|
# Load nodes
|
|
nodes = data.get("nodes", [])
|
|
self.add_nodes(nodes)
|
|
|
|
# Load edges
|
|
edges = data.get("edges", [])
|
|
self.add_edges(edges)
|
|
|
|
self.logger.info(f"Loaded context graph from {path}")
|
|
|
|
def find_node(self, node_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Find a node by ID."""
|
|
node = self.nodes.get(node_id)
|
|
if node:
|
|
return {
|
|
"id": node.node_id,
|
|
"type": node.node_type,
|
|
"content": node.content,
|
|
"metadata": node.metadata,
|
|
}
|
|
return None
|
|
|
|
def find_nodes(self, node_type: Optional[str] = None) -> List[Dict[str, Any]]:
|
|
"""Find nodes, optionally filtered by type."""
|
|
if node_type:
|
|
node_ids = self.node_type_index.get(node_type, set())
|
|
nodes = [self.nodes[nid] for nid in node_ids]
|
|
else:
|
|
nodes = self.nodes.values()
|
|
|
|
return [
|
|
{
|
|
"id": n.node_id,
|
|
"type": n.node_type,
|
|
"content": n.content,
|
|
"metadata": n.metadata,
|
|
}
|
|
for n in nodes
|
|
]
|
|
|
|
def find_edges(self, edge_type: Optional[str] = None) -> List[Dict[str, Any]]:
|
|
"""Find edges, optionally filtered by type."""
|
|
if edge_type:
|
|
edges = self.edge_type_index.get(edge_type, [])
|
|
else:
|
|
edges = self.edges
|
|
|
|
return [
|
|
{
|
|
"source": e.source_id,
|
|
"target": e.target_id,
|
|
"type": e.edge_type,
|
|
"weight": e.weight,
|
|
"metadata": e.metadata,
|
|
}
|
|
for e in edges
|
|
]
|
|
|
|
def stats(self) -> Dict[str, Any]:
|
|
"""Get graph statistics."""
|
|
return {
|
|
"node_count": len(self.nodes),
|
|
"edge_count": len(self.edges),
|
|
"node_types": {k: len(v) for k, v in self.node_type_index.items()},
|
|
"edge_types": {k: len(v) for k, v in self.edge_type_index.items()},
|
|
"density": self.density(),
|
|
}
|
|
|
|
def density(self) -> float:
|
|
"""Calculate graph density."""
|
|
n = len(self.nodes)
|
|
if n < 2:
|
|
return 0.0
|
|
max_edges = n * (n - 1) # Directed graph
|
|
return len(self.edges) / max_edges
|
|
|
|
# --- Internal Helpers ---
|
|
|
|
def _add_internal_node(self, node: ContextNode) -> bool:
|
|
"""Internal method to add a node."""
|
|
self.nodes[node.node_id] = node
|
|
self.node_type_index[node.node_type].add(node.node_id)
|
|
return True
|
|
|
|
def _add_internal_edge(self, edge: ContextEdge) -> bool:
|
|
"""Internal method to add an edge."""
|
|
# Ensure nodes exist
|
|
if edge.source_id not in self.nodes:
|
|
self._add_internal_node(
|
|
ContextNode(edge.source_id, "entity", edge.source_id)
|
|
)
|
|
if edge.target_id not in self.nodes:
|
|
self._add_internal_node(
|
|
ContextNode(edge.target_id, "entity", edge.target_id)
|
|
)
|
|
|
|
self.edges.append(edge)
|
|
self.edge_type_index[edge.edge_type].append(edge)
|
|
self._adjacency[edge.source_id].append(edge)
|
|
return True
|
|
|
|
# --- Builder Methods (Legacy/Utility) ---
|
|
|
|
def build_from_conversations(
|
|
self,
|
|
conversations: List[Union[str, Dict[str, Any]]],
|
|
link_entities: bool = True,
|
|
extract_intents: bool = False,
|
|
extract_sentiments: bool = False,
|
|
**options,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Build context graph from conversations and return dict representation.
|
|
|
|
Args:
|
|
conversations: List of conversation files or dictionaries
|
|
...
|
|
|
|
Returns:
|
|
Graph dictionary (nodes, edges)
|
|
"""
|
|
tracking_id = self.progress_tracker.start_tracking(
|
|
file=None,
|
|
module="context",
|
|
submodule="ContextGraph",
|
|
message=f"Building graph from {len(conversations)} conversations",
|
|
)
|
|
|
|
try:
|
|
for conv in conversations:
|
|
conv_data = (
|
|
conv if isinstance(conv, dict) else self._load_conversation(conv)
|
|
)
|
|
self._process_conversation(
|
|
conv_data,
|
|
extract_intents=extract_intents,
|
|
extract_sentiments=extract_sentiments,
|
|
)
|
|
|
|
if link_entities:
|
|
self._link_entities()
|
|
|
|
self.progress_tracker.stop_tracking(tracking_id, status="completed")
|
|
return self.to_dict()
|
|
|
|
except Exception as e:
|
|
self.progress_tracker.stop_tracking(
|
|
tracking_id, status="failed", message=str(e)
|
|
)
|
|
raise
|
|
|
|
def build_from_entities_and_relationships(
|
|
self,
|
|
entities: List[Dict[str, Any]],
|
|
relationships: List[Dict[str, Any]],
|
|
**kwargs,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Build graph from entities and relationships.
|
|
|
|
Args:
|
|
entities: List of entity dictionaries
|
|
relationships: List of relationship dictionaries
|
|
**kwargs: Additional options
|
|
|
|
Returns:
|
|
Graph dictionary (nodes, edges)
|
|
"""
|
|
tracking_id = self.progress_tracker.start_tracking(
|
|
file=None,
|
|
module="context",
|
|
submodule="ContextGraph",
|
|
message=(
|
|
f"Building graph from {len(entities)} entities and "
|
|
f"{len(relationships)} relationships"
|
|
),
|
|
)
|
|
|
|
try:
|
|
# Add entities
|
|
for entity in entities:
|
|
entity_id = entity.get("id") or entity.get("entity_id")
|
|
if entity_id:
|
|
self._add_internal_node(
|
|
ContextNode(
|
|
node_id=entity_id,
|
|
node_type=entity.get("type", "entity"),
|
|
content=entity.get("text")
|
|
or entity.get("label")
|
|
or entity_id,
|
|
metadata=entity,
|
|
properties=entity,
|
|
)
|
|
)
|
|
|
|
# Add relationships
|
|
for rel in relationships:
|
|
source = rel.get("source_id")
|
|
target = rel.get("target_id")
|
|
if source and target:
|
|
self._add_internal_edge(
|
|
ContextEdge(
|
|
source_id=source,
|
|
target_id=target,
|
|
edge_type=rel.get("type", "related_to"),
|
|
weight=rel.get("confidence", 1.0),
|
|
metadata=rel,
|
|
)
|
|
)
|
|
|
|
self.progress_tracker.stop_tracking(tracking_id, status="completed")
|
|
return self.to_dict()
|
|
|
|
except Exception as e:
|
|
self.progress_tracker.stop_tracking(
|
|
tracking_id, status="failed", message=str(e)
|
|
)
|
|
raise
|
|
|
|
def _process_conversation(self, conv_data: Dict[str, Any], **kwargs) -> None:
|
|
"""Process a single conversation."""
|
|
conv_id = conv_data.get("id") or f"conv_{hash(str(conv_data)) % 10000}"
|
|
|
|
# Add conversation node
|
|
self._add_internal_node(
|
|
ContextNode(
|
|
node_id=conv_id,
|
|
node_type="conversation",
|
|
content=conv_data.get("content", "") or conv_data.get("summary", ""),
|
|
metadata={"timestamp": conv_data.get("timestamp")},
|
|
)
|
|
)
|
|
|
|
# Track name to ID mapping for relationship resolution
|
|
name_to_id = {}
|
|
|
|
# Extract entities
|
|
if self.extract_entities:
|
|
for entity in conv_data.get("entities", []):
|
|
entity_id = entity.get("id") or entity.get("entity_id")
|
|
entity_text = (
|
|
entity.get("text")
|
|
or entity.get("label")
|
|
or entity.get("name")
|
|
or entity_id
|
|
)
|
|
entity_type = entity.get("type", "entity")
|
|
|
|
# Generate ID if missing
|
|
if not entity_id and entity_text and self.entity_linker:
|
|
# Use EntityLinker to generate ID
|
|
if hasattr(self.entity_linker, "_generate_entity_id"):
|
|
entity_id = self.entity_linker._generate_entity_id(
|
|
entity_text, entity_type
|
|
)
|
|
else:
|
|
# Fallback ID generation
|
|
import hashlib
|
|
|
|
entity_hash = hashlib.md5(
|
|
f"{entity_text}_{entity_type}".encode()
|
|
).hexdigest()[:12]
|
|
entity_id = f"{entity_type.lower()}_{entity_hash}"
|
|
|
|
if entity_id:
|
|
if entity_text:
|
|
name_to_id[entity_text] = entity_id
|
|
|
|
self._add_internal_node(
|
|
ContextNode(
|
|
node_id=entity_id,
|
|
node_type="entity",
|
|
content=entity_text,
|
|
metadata={"type": entity_type, **entity},
|
|
)
|
|
)
|
|
self._add_internal_edge(
|
|
ContextEdge(
|
|
source_id=conv_id,
|
|
target_id=entity_id,
|
|
edge_type="mentions",
|
|
)
|
|
)
|
|
|
|
# Extract relationships
|
|
if self.extract_relationships:
|
|
for rel in conv_data.get("relationships", []):
|
|
source = rel.get("source_id")
|
|
target = rel.get("target_id")
|
|
|
|
# Resolve IDs from names if missing
|
|
if not source and rel.get("source") and rel.get("source") in name_to_id:
|
|
source = name_to_id[rel.get("source")]
|
|
|
|
if not target and rel.get("target") and rel.get("target") in name_to_id:
|
|
target = name_to_id[rel.get("target")]
|
|
|
|
if source and target:
|
|
self._add_internal_edge(
|
|
ContextEdge(
|
|
source_id=source,
|
|
target_id=target,
|
|
edge_type=rel.get("type", "related_to"),
|
|
weight=rel.get("confidence", 1.0),
|
|
)
|
|
)
|
|
|
|
def _link_entities(self) -> None:
|
|
"""Link similar entities using EntityLinker."""
|
|
if not self.entity_linker:
|
|
return
|
|
|
|
entity_nodes = [n for n in self.nodes.values() if n.node_type == "entity"]
|
|
for i, node1 in enumerate(entity_nodes):
|
|
for node2 in entity_nodes[i + 1 :]:
|
|
similarity = self.entity_linker._calculate_text_similarity(
|
|
node1.content.lower(), node2.content.lower()
|
|
)
|
|
if similarity >= self.entity_linker.similarity_threshold:
|
|
self._add_internal_edge(
|
|
ContextEdge(
|
|
source_id=node1.node_id,
|
|
target_id=node2.node_id,
|
|
edge_type="similar_to",
|
|
weight=similarity,
|
|
)
|
|
)
|
|
|
|
def _load_conversation(self, file_path: str) -> Dict[str, Any]:
|
|
"""Load conversation from file."""
|
|
from ..utils.helpers import read_json_file
|
|
from pathlib import Path
|
|
|
|
return read_json_file(Path(file_path))
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Export graph to dictionary format."""
|
|
return {
|
|
"nodes": [
|
|
{
|
|
"id": n.node_id,
|
|
"type": n.node_type,
|
|
"content": n.content,
|
|
"metadata": n.metadata,
|
|
}
|
|
for n in self.nodes.values()
|
|
],
|
|
"edges": [
|
|
{
|
|
"source": e.source_id,
|
|
"target": e.target_id,
|
|
"type": e.edge_type,
|
|
"weight": e.weight,
|
|
}
|
|
for e in self.edges
|
|
],
|
|
"statistics": {
|
|
"node_count": len(self.nodes),
|
|
"edge_count": len(self.edges),
|
|
},
|
|
}
|
|
|
|
|
|
# For backward compatibility
|
|
ContextGraphBuilder = ContextGraph
|