diff --git a/backend/app/rag/graph_builder.py b/backend/app/rag/graph_builder.py index 8e2a0f0e..e883c9ce 100644 --- a/backend/app/rag/graph_builder.py +++ b/backend/app/rag/graph_builder.py @@ -88,6 +88,29 @@ def extract_entities(text: str) -> List[Entity]: return list(entities.values()) +def extract_relationship_predicate(text_chunk: str, ent1_text: str, ent2_text: str) -> str: + """Extract a clean relationship predicate between two entities using local context.""" + if not text_chunk or not ent1_text or not ent2_text: + return "co-occurs with" + + idx1 = text_chunk.find(ent1_text) + idx2 = text_chunk.find(ent2_text) + + if idx1 == -1 or idx2 == -1: + return "co-occurs with" + + start = min(idx1 + len(ent1_text), idx2 + len(ent2_text)) + end = max(idx1, idx2) + + between_text = text_chunk[start:end].strip(" ,.-_()\t\n") + clean_predicate = " ".join(between_text.split()) + + if not clean_predicate: + return "co-occurs with" + if len(clean_predicate.split()) > 8: + return "appears in context with" + + return clean_predicate def build_graph(chunks: List[Dict[str, Any]]) -> nx.Graph: """Build an entity co-occurrence graph from document chunks.""" @@ -115,16 +138,27 @@ def build_graph(chunks: List[Dict[str, Any]]) -> nx.Graph: ) for left_index, left in enumerate(entities): + left_id = left.id + adj_dict = graph[left_id] + for right in entities[left_index + 1:]: - if graph.has_edge(left.id, right.id): - graph[left.id][right.id]["weight"] += 1 - graph[left.id][right.id]["pages"].add(page) - graph[left.id][right.id]["chunks"].add(chunk_index) + right_id = right.id + + # Extract the dynamic semantic relationship context string + predicate = extract_relationship_predicate(text, left.text, right.text) + + if right_id in adj_dict: + edge = adj_dict[right_id] + edge["weight"] += 1 + edge["pages"].add(page) + edge["chunks"].add(chunk_index) + edge["relationships"].add(predicate) else: graph.add_edge( - left.id, - right.id, + left_id, + right_id, weight=1, + relationships={predicate}, pages={page}, chunks={chunk_index}, ) @@ -135,12 +169,13 @@ def build_graph(chunks: List[Dict[str, Any]]) -> nx.Graph: def _convert_sets_for_json(graph: nx.Graph) -> None: for _, data in graph.nodes(data=True): - data["pages"] = sorted(item for item in data.get("pages", []) if item is not None) - data["chunks"] = sorted(item for item in data.get("chunks", []) if item is not None) + data["pages"] = sorted((item for item in data.get("pages", []) if item is not None), key=str) + data["chunks"] = sorted((item for item in data.get("chunks", []) if item is not None), key=str) for _, _, data in graph.edges(data=True): - data["pages"] = sorted(item for item in data.get("pages", []) if item is not None) - data["chunks"] = sorted(item for item in data.get("chunks", []) if item is not None) + data["pages"] = sorted((item for item in data.get("pages", []) if item is not None), key=str) + data["chunks"] = sorted((item for item in data.get("chunks", []) if item is not None), key=str) + data["relationships"] = sorted(list(data.get("relationships", []))) def save_graph(graph: nx.Graph, user_id: str, document_id: str) -> Path: diff --git a/backend/app/rag/graph_retriever.py b/backend/app/rag/graph_retriever.py index 39841aa1..71177ac4 100644 --- a/backend/app/rag/graph_retriever.py +++ b/backend/app/rag/graph_retriever.py @@ -84,18 +84,21 @@ def get_entity_context( ) for neighbor_id in neighbors: edge = graph[node_id][neighbor_id] - left = _node_name(graph, node_id) - right = _node_name(graph, neighbor_id) - key = _relationship_key(left.casefold(), right.casefold()) - existing = relationships.setdefault( - key, - { - "left": left, - "right": right, + source_name = _node_name(graph, node_id) + target_name = _node_name(graph, neighbor_id) + + # Deduplicate safely using alphabetical keys behind the scenes + dedup_key = tuple(sorted((source_name.casefold(), target_name.casefold()))) + + if dedup_key not in relationships: + relationships[dedup_key] = { + "subject": source_name, # The matched anchor entity + "object": target_name, # The connected neighbor entity "weight": 0, "pages": set(), - }, - ) + } + + existing = relationships[dedup_key] existing["weight"] = int(existing["weight"]) + int(edge.get("weight", 1)) existing["pages"].update(edge.get("pages", [])) except Exception as exc: @@ -115,7 +118,7 @@ def get_entity_context( for item in ranked: pages = sorted(item["pages"]) lines.append( - f"- {item['left']} is related to {item['right']} " + f"- {item['subject']} is related to {item['object']} " f"through document co-occurrence on {_format_pages(pages)} " f"(strength: {item['weight']})." )