Added Embeddings Option on custom Triplets

llama-index/lib64

@@ -0,0 +1 @@
+lib

llama_index/indices/knowledge_graph/base.py

@@ -217,17 +217,25 @@ def _insert(self, nodes: Sequence[BaseNode], **insert_kwargs: Any) -> None:
                     )
                     self._index_struct.add_to_embedding_dict(triplet_str, rel_embedding)
 
-    def upsert_triplet(self, triplet: Tuple[str, str, str]) -> None:
-        """Insert triplets.
+    def upsert_triplet(
+        self, triplet: Tuple[str, str, str], include_embeddings: bool = False
+    ) -> None:
+        """Insert triplets and optionally embeddings.
 
         Used for manual insertion of KG triplets (in the form
         of (subject, relationship, object)).
 
         Args:
-            triplet (str): Knowledge triplet
-
+            triplet (tuple): Knowledge triplet
+            embedding (Any, optional): Embedding option for the triplet. Defaults to None.
         """
         self._graph_store.upsert_triplet(*triplet)
+        triplet_str = str(triplet)
+        if include_embeddings:
+            set_embedding = self._service_context.embed_model.get_text_embedding(
+                triplet_str
+            )
+            self._index_struct.add_to_embedding_dict(str(triplet), set_embedding)
 
     def add_node(self, keywords: List[str], node: BaseNode) -> None:
         """Add node.
@@ -243,7 +251,10 @@ def add_node(self, keywords: List[str], node: BaseNode) -> None:
         self._docstore.add_documents([node], allow_update=True)
 
     def upsert_triplet_and_node(
-        self, triplet: Tuple[str, str, str], node: BaseNode
+        self,
+        triplet: Tuple[str, str, str],
+        node: BaseNode,
+        include_embeddings: bool = False,
     ) -> None:
         """Upsert KG triplet and node.
 
@@ -254,11 +265,18 @@ def upsert_triplet_and_node(
         Args:
             keywords (List[str]): Keywords to index the node.
             node (Node): Node to be indexed.
+            include_embeddings (bool): Option to add embeddings for triplets. Defaults to False
 
         """
         subj, _, obj = triplet
         self.upsert_triplet(triplet)
         self.add_node([subj, obj], node)
+        triplet_str = str(triplet)
+        if include_embeddings:
+            set_embedding = self._service_context.embed_model.get_text_embedding(
+                triplet_str
+            )
+            self._index_struct.add_to_embedding_dict(str(triplet), set_embedding)
 
     def _delete_node(self, node_id: str, **delete_kwargs: Any) -> None:
         """Delete a node."""

storage/default__vector_store.json

@@ -0,0 +1 @@
+{ "embedding_dict": {}, "text_id_to_ref_doc_id": {}, "metadata_dict": {} }

storage/docstore.json

@@ -0,0 +1,69 @@
+{
+  "docstore/data": {
+    "8e2fa999-a703-445c-b678-d67cdc72647a": {
+      "__data__": {
+        "id_": "8e2fa999-a703-445c-b678-d67cdc72647a",
+        "embedding": null,
+        "metadata": {},
+        "excluded_embed_metadata_keys": [],
+        "excluded_llm_metadata_keys": [],
+        "relationships": {},
+        "text": "a",
+        "start_char_idx": null,
+        "end_char_idx": null,
+        "text_template": "{metadata_str}\n\n{content}",
+        "metadata_template": "{key}: {value}",
+        "metadata_seperator": "\n",
+        "class_name": "TextNode"
+      },
+      "__type__": "1"
+    },
+    "02ba5550-91e9-4854-99f6-14b8f8f9d83c": {
+      "__data__": {
+        "id_": "02ba5550-91e9-4854-99f6-14b8f8f9d83c",
+        "embedding": null,
+        "metadata": {},
+        "excluded_embed_metadata_keys": [],
+        "excluded_llm_metadata_keys": [],
+        "relationships": {},
+        "text": "b",
+        "start_char_idx": null,
+        "end_char_idx": null,
+        "text_template": "{metadata_str}\n\n{content}",
+        "metadata_template": "{key}: {value}",
+        "metadata_seperator": "\n",
+        "class_name": "TextNode"
+      },
+      "__type__": "1"
+    },
+    "2bee59d9-c054-4fc4-b81a-b3c0b4a0f46d": {
+      "__data__": {
+        "id_": "2bee59d9-c054-4fc4-b81a-b3c0b4a0f46d",
+        "embedding": null,
+        "metadata": {},
+        "excluded_embed_metadata_keys": [],
+        "excluded_llm_metadata_keys": [],
+        "relationships": {},
+        "text": "c",
+        "start_char_idx": null,
+        "end_char_idx": null,
+        "text_template": "{metadata_str}\n\n{content}",
+        "metadata_template": "{key}: {value}",
+        "metadata_seperator": "\n",
+        "class_name": "TextNode"
+      },
+      "__type__": "1"
+    }
+  },
+  "docstore/metadata": {
+    "8e2fa999-a703-445c-b678-d67cdc72647a": {
+      "doc_hash": "5f546eb4606b5c2b7d2a449a5cc2bbb477ed5a246c7051ce871b12f2dbfc8419"
+    },
+    "02ba5550-91e9-4854-99f6-14b8f8f9d83c": {
+      "doc_hash": "ef2a875407aca2dd75a7a1dccc3b202918a9bc8e543e0b3def3f34ccd5e56a7b"
+    },
+    "2bee59d9-c054-4fc4-b81a-b3c0b4a0f46d": {
+      "doc_hash": "7ed0bc022e61394074858c49e5e2f8844146ae76fd02cc3b74ed8fb450e33b42"
+    }
+  }
+}

storage/graph_store.json

@@ -0,0 +1 @@
+{ "graph_dict": {} }

storage/image__vector_store.json

@@ -0,0 +1 @@
+{ "embedding_dict": {}, "text_id_to_ref_doc_id": {}, "metadata_dict": {} }

storage/index_store.json

@@ -0,0 +1,8 @@
+{
+  "index_store/data": {
+    "1790cecc-e4c7-48ae-850a-dac2811bf147": {
+      "__type__": "list",
+      "__data__": "{\"index_id\": \"1790cecc-e4c7-48ae-850a-dac2811bf147\", \"summary\": null, \"nodes\": [\"8e2fa999-a703-445c-b678-d67cdc72647a\", \"02ba5550-91e9-4854-99f6-14b8f8f9d83c\", \"2bee59d9-c054-4fc4-b81a-b3c0b4a0f46d\"]}"
+    }
+  }
+}