Improve entity extraction prompts for relationship array alignment

Author: caseyclements

libs/langchain-mongodb/langchain_mongodb/graphrag/example_templates.py

@@ -25,7 +25,7 @@
         "startDate": ["2018-01-01"]
       }},
       "relationships": {{
-        "targets": ["Jasbinder Kaur", "Jarnail Singh"],
+        "target_ids": ["Jasbinder Kaur", "Jarnail Singh"],
         "types": ["Friend", "Friend"],
         "attributes": [
           {{ "since": ["2019-05-01"] }},
@@ -37,7 +37,7 @@
       "_id": "Jarnail Singh",
       "type": "Person",
       "relationships": {{
-        "targets": ["Alice Palace"],
+        "target_ids": ["Alice Palace"],
         "types": ["Friend"],
         "attributes": [{{ "since": ["2019-05-01"] }}]
       }}
@@ -46,7 +46,7 @@
       "_id": "Jasbinder Kaur",
       "type": "Person",
       "relationships": {{
-        "targets": ["Alice Palace"],
+        "target_ids": ["Alice Palace"],
         "types": ["Friend"],
         "attributes": [{{ "since": ["2015-05-01"], "frequency": ["weekly"] }}]
       }}
@@ -71,8 +71,9 @@
         "location": ["San Francisco"]
       }},
       "relationships": {{
-        "targets": ["Elon Musk", "Sam Altman"],
-        "types": ["Speaker", "Speaker"]
+        "target_ids": ["Elon Musk", "Sam Altman"],
+        "types": ["Speaker", "Speaker"],
+        "attributes": [{{}}, {{}}]
       }}
     }},
     {{ "_id": "Elon Musk", "type": "Person" }},
@@ -92,8 +93,9 @@
       "_id": "Quantum Computing",
       "type": "Concept",
       "relationships": {{
-        "targets": ["Quantum Mechanics"],
-        "types": ["Based On"]
+        "target_ids": ["Quantum Mechanics"],
+        "types": ["Based On"],
+        "attributes": [{{}}]
       }}
     }},
     {{ "_id": "Quantum Mechanics", "type": "Concept" }}
@@ -114,8 +116,9 @@
       "type": "Event",
       "attributes": {{ "date": ["2023-03-01"] }},
       "relationships": {{
-        "targets": ["NASA"],
-        "types": ["Managed By"]
+        "target_ids": ["NASA"],
+        "types": ["Managed By"],
+        "attributes": [{{}}]
       }}
     }},
     {{
@@ -126,8 +129,9 @@
       "_id": "Bill Nelson",
       "type": "Person",
       "relationships": {{
-        "targets": ["Artemis II Mission"],
-        "types": ["Praised By"]
+        "target_ids": ["Artemis II Mission"],
+        "types": ["Praised By"],
+        "attributes": [{{}}]
       }}
     }}
   ]
@@ -146,16 +150,18 @@
       "_id": "Rust",
       "type": "Programming Language",
       "relationships": {{
-        "targets": ["Memory Safety"],
-        "types": ["Ensures"]
+        "target_ids": ["Memory Safety"],
+        "types": ["Ensures"],
+        "attributes": [{{}}]
       }}
     }},
     {{
       "_id": "Memory Safety",
       "type": "Concept",
       "relationships": {{
-        "targets": ["Ownership Model"],
-        "types": ["Uses"]
+        "target_ids": ["Ownership Model"],
+        "types": ["Uses"],
+        "attributes": [{{}}]
       }}
     }},
     {{ "_id": "Ownership Model", "type": "Concept" }}

libs/langchain-mongodb/langchain_mongodb/graphrag/prompts.py

@@ -46,6 +46,37 @@
 Instead of using specific and momentary types such as 'worked_at', use more general and timeless relationship types
 like 'employee'. Add details as attributes. Make sure to use general and timeless relationship types!
 
+### CRITICAL: Array Length Alignment
+The relationships object contains three arrays: `target_ids`, `types`, and `attributes`.
+**These three arrays MUST have EXACTLY the same length.**
+- Each position (index) in these arrays describes ONE complete relationship.
+- Position 0 in `target_ids`, `types`, and `attributes` together describe the first relationship.
+- Position 1 in `target_ids`, `types`, and `attributes` together describe the second relationship.
+- And so on...
+
+If a relationship has no attributes, you MUST still include an empty object `{{}}` in the `attributes` array at that position.
+
+Example of CORRECT alignment:
+```json
+"relationships": {{
+  "target_ids": ["Entity A", "Entity B"],
+  "types": ["partners", "supplier"],
+  "attributes": [
+    {{"since": ["2020"]}},
+    {{}}
+  ]
+}}
+```
+
+Example of INCORRECT (DO NOT DO THIS):
+```json
+"relationships": {{
+  "target_ids": ["Entity A", "Entity B"],
+  "types": ["partners"],
+  "attributes": [{{"since": ["2020"]}}]
+}}
+```
+
 **Allowed Relationship Types**:
 - Extract ONLY relationships whose `type` matches one of the following: {allowed_relationship_types}.
 - If this list is empty, ANY relationship type is permitted.
@@ -64,7 +95,8 @@
 1. Validate that all extracted entities have an `_id` and `type`.
 2. Validate that all `type` values are in {allowed_entity_types}.
 3. Validate that all relationships use keys in {allowed_relationship_types}.
-4. Exclude any entities or relationships failing validation.
+4. **CRITICAL**: For each entity with relationships, verify that `target_ids`, `types`, and `attributes` arrays have EXACTLY the same length.
+5. Exclude any entities or relationships failing validation.
 
 ## Output Schema
 Output a valid JSON document with a single top-level key, `entities`, as an array of objects.