diff --git a/llama_index/indices/knowledge_graph/base.py b/llama_index/indices/knowledge_graph/base.py index 468019202f5b5839fd352ab371b349a8a8574a2a..e38a178feef9f0f969a2bef2b170a60aa6f724fc 100644 --- a/llama_index/indices/knowledge_graph/base.py +++ b/llama_index/indices/knowledge_graph/base.py @@ -156,6 +156,12 @@ class KnowledgeGraphIndex(BaseIndex[KG]): if not subj or not pred or not obj: # skip partial triplets continue + + # Strip double quotes and Capitalize triplets for disambiguation + subj, pred, obj = ( + entity.strip('"').capitalize() for entity in [subj, pred, obj] + ) + results.append((subj, pred, obj)) return results diff --git a/tests/indices/knowledge_graph/test_base.py b/tests/indices/knowledge_graph/test_base.py index 72f322ffbf363bb512facc30d0ed44d714d09d22..c0a0ea365e2e3ae657b0c51685e7f18ba4c658fc 100644 --- a/tests/indices/knowledge_graph/test_base.py +++ b/tests/indices/knowledge_graph/test_base.py @@ -232,6 +232,7 @@ def test__parse_triplet_response( ) assert len(parsed_triplets) == 1 assert len(parsed_triplets[0]) == 3 - assert ("foo", "is", "bar") in parsed_triplets[0] - assert ("hello", "is not", "world") in parsed_triplets[0] - assert ("Jane", "is mother of", "Bob") in parsed_triplets[0] + # Expecting Capitalized triplet Outputs + assert ("Foo", "Is", "Bar") in parsed_triplets[0] + assert ("Hello", "Is not", "World") in parsed_triplets[0] + assert ("Jane", "Is mother of", "Bob") in parsed_triplets[0]