"""Knowledge Graph Builder — builds graph from ContentList.
Inspired by RAG-Anything / LightRAG's multi-modal knowledge graph approach:
1. Structural entities — every content item becomes a graph node
2. LLM semantic extraction — entities + relationships from text chunks
3. Entity deduplication — merge same-named entities across chunks
4. Cross-modal relationship mapping
5. Hierarchical structure preservation ("belongs_to" chains)
6. Weighted relationship scoring
"""
import asyncio
import hashlib
import json as _json
import re
from typing import Callable, Optional
from agentic_rag.services.knowledge.content_list import ContentItem, ContentList, ContentType
from agentic_rag.services.knowledge.graph.index import Entity, KnowledgeGraph, Relation
from agentic_rag.config.prompts import Prompts
def _stable_id(name: str, prefix: str = "ent") -> str:
"""Generate a stable, deterministic entity ID from its name."""
h = hashlib.md5(name.strip().lower().encode("utf-8")).hexdigest()[:12]
return f"{prefix}-{h}"
def _normalize_name(name: str) -> str:
"""Normalize entity name for deduplication."""
return name.strip().strip("'\".,;:!?()[]{}<>《》").strip()
class GraphBuilder:
"""Builds a knowledge graph from a ContentList.
Two-layer construction:
- **Structural layer**: every content item → entity, with belongs_to/nearby/describes edges
- **Semantic layer** (optional, requires LLM): extracts fine-grained entities and
semantic relationships from text chunks, deduplicates by name, and links to
structural entities via "appears_in" edges.
"""
def __init__(
self,
llm_func: Optional[Callable] = None,
extract_entities: bool = True,
):
"""
Args:
llm_func: LLM function (text) -> str for entity/relation extraction.
extract_entities: Whether to perform semantic entity extraction from text.
"""
self.llm_func = llm_func
self.extract_entities = extract_entities
# ═══════════════════════════════════════════════════════════
# Main entry point
# ═══════════════════════════════════════════════════════════
async def build(self, content_list: ContentList, doc_id: str = "",
existing_graph: KnowledgeGraph | None = None) -> KnowledgeGraph:
"""Build (or extend) a knowledge graph from a ContentList.
Args:
content_list: Parsed and processed document content.
doc_id: Document identifier (used as namespace).
existing_graph: If provided, extend this graph instead of creating new.
Returns:
Populated KnowledgeGraph.
"""
graph = existing_graph if existing_graph is not None else KnowledgeGraph()
import sys
# ── Layer 1: Structural entities ──────────────────────
structural_eids = []
for i, item in enumerate(content_list.items):
entity = await self._item_to_entity(item, i)
eid = graph.add_entity(entity)
structural_eids.append((eid, item))
# ── Document entity ───────────────────────────────────
doc_name = content_list.source or "Document"
if hasattr(doc_name, "split"):
doc_name = doc_name.split("/")[-1] # filename only
doc_entity = Entity(
entity_id=_stable_id(doc_name, "doc"),
name=doc_name,
entity_type="document",
properties={"doc_id": doc_id},
)
doc_eid = graph.add_entity(doc_entity)
for eid, _ in structural_eids:
graph.add_relation(eid, doc_eid, "belongs_to", weight=1.0)
# ── Layer 2: Page-level adjacency ─────────────────────
page_entities: dict[int, list[str]] = {}
for eid, item in structural_eids:
page_entities.setdefault(item.page_idx, []).append(eid)
for page, eids in page_entities.items():
for i in range(len(eids) - 1):
graph.add_relation(eids[i], eids[i + 1], "nearby", weight=0.8)
# ── Layer 3: Cross-modal references ───────────────────
for page, eids in page_entities.items():
page_ents = [graph.get_entity(eid) for eid in eids]
page_ents = [e for e in page_ents if e is not None]
text_ents = [e for e in page_ents if e.type == "text_chunk"]
visual_ents = [e for e in page_ents if e.type in ("image", "table", "equation")]
for text_e in text_ents:
for visual_e in visual_ents:
graph.add_relation(
visual_e.id, text_e.id, "describes",
weight=0.6,
metadata={"page_idx": page},
)
# ── Layer 4: Semantic entity extraction (LLM) ────────
if self.extract_entities and self.llm_func:
all_semantic_entities: dict[str, str] = {} # normalized_name -> eid
text_chunks = [
(eid, item)
for eid, item in structural_eids
if item.type == ContentType.TEXT and item.text and len(item.text.strip()) > 30
]
for eid, item in text_chunks:
try:
chunk_entities, chunk_relations = await self._extract_from_chunk(
item.text
)
# Register semantic entities + "appears_in" links
for ent in chunk_entities:
norm = _normalize_name(ent["name"])
if norm in all_semantic_entities:
sem_eid = all_semantic_entities[norm]
else:
sem_entity = Entity(
entity_id=_stable_id(norm, "ent"),
name=ent["name"],
entity_type=ent.get("type", "CONCEPT").lower(),
properties={
"description": ent.get("description", ""),
"source_doc": doc_name,
},
)
sem_eid = graph.add_entity(sem_entity)
all_semantic_entities[norm] = sem_eid
# Link semantic entity → structural chunk
graph.add_relation(
sem_eid, eid, "appears_in",
weight=0.9,
metadata={"chunk_type": "text"},
)
# Register semantic relationships
for rel in chunk_relations:
src_norm = _normalize_name(rel["source"])
tgt_norm = _normalize_name(rel["target"])
if src_norm in all_semantic_entities and tgt_norm in all_semantic_entities:
src_eid = all_semantic_entities[src_norm]
tgt_eid = all_semantic_entities[tgt_norm]
if src_eid != tgt_eid: # skip self-loops
graph.add_relation(
src_eid, tgt_eid,
rel.get("keywords", "related_to").split(",")[0].strip(),
weight=0.85,
metadata={
"keywords": rel.get("keywords", ""),
"description": rel.get("description", ""),
},
)
except Exception as e:
print(f" [GraphBuilder] ⚠ Semantic extraction failed for chunk: {e}",
flush=True)
sys.stdout.flush()
print(f" [GraphBuilder] Semantic: {len(all_semantic_entities)} unique entities "
f"from {len(text_chunks)} chunks", flush=True)
sys.stdout.flush()
return graph
# ═══════════════════════════════════════════════════════════
# Helpers
# ═══════════════════════════════════════════════════════════
async def _extract_from_chunk(self, text: str) -> tuple[list[dict], list[dict]]:
"""Extract semantic entities and relationships from a text chunk via LLM.
Returns:
(entities_list, relationships_list) — each entity dict has
name/type/description; each relation has source/target/keywords/description.
"""
import sys
prompt = Prompts.KG_ENTITY_EXTRACTION.format(text=text[:2000])
try:
raw = await self.llm_func(prompt)
# Extract result from callable (may return str or object)
if hasattr(raw, "content"):
raw = raw.content
elif hasattr(raw, "choices") and raw.choices:
raw = raw.choices[0].message.content
raw = str(raw).strip() if raw else ""
except Exception as e:
print(f" [GraphBuilder] ⚠ LLM call failed: {e}", flush=True)
sys.stdout.flush()
return [], []
if not raw:
print(f" [GraphBuilder] ⚠ LLM returned empty response", flush=True)
sys.stdout.flush()
return [], []
# Strip Qwen ... reasoning blocks (if present)
raw = re.sub(r'[\s\S]*?', '', raw).strip()
if not raw:
print(f" [GraphBuilder] ⚠ LLM response was all block", flush=True)
sys.stdout.flush()
return [], []
# Parse JSON from LLM output.
# Qwen reasoning models prepend a "thinking process" before the actual
# JSON — we must skip that and find the real output block.
json_str = raw
# Strategy 1: Find ```json fenced block (most reliable)
fence_match = re.search(r'```(?:json)?\s*([\s\S]*?)```', raw)
if fence_match:
json_str = fence_match.group(1).strip()
else:
# Strategy 2: Find the outermost JSON object that contains
# "entities" key. Reasoning models put thinking text BEFORE the
# JSON, so we scan from each `{` and pick the first one that
# produces a complete, parseable dict with the right keys.
json_str = ""
for m in re.finditer(r'\{', raw):
start = m.start()
candidate = raw[start:]
depth = 0
end_pos = -1
for i, ch in enumerate(candidate):
if ch == '{':
depth += 1
elif ch == '}':
depth -= 1
if depth == 0:
end_pos = i + 1
break
if end_pos > 0:
candidate = candidate[:end_pos].strip()
# Quick check: does it look like our target JSON?
if '"entities"' in candidate or '"entity"' in candidate:
try:
_json.loads(candidate)
json_str = candidate
break # found it
except _json.JSONDecodeError:
continue # try next `{`
# Strategy 3: Look for known JSON keys as a fallback
if not json_str or '{' not in json_str:
for marker in ('"entities"', '"entity"'):
idx = raw.find(marker)
if idx > 0:
brace_idx = raw.rfind('{', 0, idx)
if brace_idx >= 0:
json_str = raw[brace_idx:]
depth = 0
end_pos = -1
for i, ch in enumerate(json_str):
if ch == '{': depth += 1
elif ch == '}':
depth -= 1
if depth == 0:
end_pos = i + 1
break
if end_pos > 0:
json_str = json_str[:end_pos].strip()
break
try:
data = _json.loads(json_str)
except _json.JSONDecodeError:
# Attempt repair: try to fix common LLM JSON mistakes
try:
import json as _json_module
data = _json.loads(json_str.replace("'", '"'))
except Exception:
print(f" [GraphBuilder] ⚠ JSON parse failed, raw={raw[:120]}", flush=True)
sys.stdout.flush()
return [], []
if not isinstance(data, dict):
print(f" [GraphBuilder] ⚠ Unexpected LLM output type: {type(data).__name__}", flush=True)
sys.stdout.flush()
return [], []
entities = data.get("entities", []) or data.get("entity", []) or []
relationships = data.get("relationships", []) or data.get("relations", []) or data.get("relationship", []) or []
# Validate
valid_entities = []
for e in entities:
if isinstance(e, dict) and e.get("name", "").strip():
valid_entities.append({
"name": _normalize_name(e["name"]),
"type": str(e.get("type", "CONCEPT")).upper(),
"description": str(e.get("description", ""))[:200],
})
valid_relations = []
for r in relationships:
if isinstance(r, dict) and r.get("source", "").strip() and r.get("target", "").strip():
valid_relations.append({
"source": _normalize_name(r["source"]),
"target": _normalize_name(r["target"]),
"keywords": str(r.get("keywords", "related_to"))[:100],
"description": str(r.get("description", ""))[:200],
})
if not valid_entities:
print(f" [GraphBuilder] ⚠ No valid entities extracted from chunk "
f"(text={text[:60]}...)", flush=True)
sys.stdout.flush()
return valid_entities, valid_relations
async def _item_to_entity(self, item: ContentItem, index: int) -> Entity:
"""Convert a ContentItem to a structural graph Entity."""
type_map = {
ContentType.TEXT: "text_chunk",
ContentType.IMAGE: "image",
ContentType.TABLE: "table",
ContentType.EQUATION: "equation",
ContentType.VIDEO: "video_frame",
ContentType.AUDIO: "audio_segment",
ContentType.CODE: "code_block",
}
entity_type = type_map.get(item.type, "unknown")
searchable = item.to_searchable_text()
# Entity name: extract a meaningful title from the chunk text.
# Strategy (in order of preference):
# 1. Section header: 【...】 or \n...\n
# 2. First complete sentence (ends with 。！？)
# 3. First 50 chars, trimmed at word boundary
if entity_type == "text_chunk" and searchable:
text = searchable.strip()
import re as _re
# Try section header pattern: 【XXX】 or 第X页
sec = _re.search(r'【(.+?)】', text)
if sec:
name = sec.group(0) # e.g. "【药品名称】"
else:
# Try first line if it's a short heading
first_line = text.split('\n')[0].strip()
if len(first_line) <= 30 and not first_line.endswith(('。', '！', '？')):
name = first_line
else:
# Find first complete sentence
name = text[:80]
for sep in ("。", "！", "？", "；"):
idx = name.find(sep)
if 15 < idx < 70:
name = name[:idx + 1]
break
# If name is still too long, trim at 50 chars
if len(name) > 50:
name = name[:50] + "…"
elif searchable:
name = searchable[:50]
else:
name = f"{entity_type}_{index}"
# Generate stable ID from content hash
id_src = (item.text or item.img_path or item.video_path or item.audio_path or str(index))
eid = _stable_id(id_src[:200], entity_type[:6])
return Entity(
entity_id=eid,
name=name,
entity_type=entity_type,
content_item=item,
properties={
"page_idx": item.page_idx,
"content_type": item.type.value,
"has_caption": bool(item.img_caption or item.table_caption or item.video_caption),
"index": index,
},
)