feat: structure RAG context output to respect relationship directionality #807

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:

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']})."
         )