From da24cffe73a340cc8f0950a5ef158bbe46446b4a Mon Sep 17 00:00:00 2001
From: nanglo123 <tnanglo1998@gmail.com>
Date: Tue, 9 Jul 2024 13:43:35 -0400
Subject: [PATCH] Add add_relation endpoint processing Relationship object
 models

---
 mira/dkg/api.py    |  9 ++++-----
 mira/dkg/client.py | 39 ++++++++++++++++++++++++++++++++++-----
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/mira/dkg/api.py b/mira/dkg/api.py
index 7ac8b5f0..c109beb1 100644
--- a/mira/dkg/api.py
+++ b/mira/dkg/api.py
@@ -11,7 +11,7 @@
 from scipy.spatial import distance
 from typing_extensions import Literal
 
-from mira.dkg.client import AskemEntity, Entity
+from mira.dkg.client import AskemEntity, Entity, Relation
 from mira.dkg.utils import DKG_REFINER_RELS
 
 __all__ = [
@@ -335,7 +335,7 @@ def get_relations(
 if active_add_relation_endpoint:
     @api_blueprint.post(
         "/add_nodes",
-        response_model=Union[AskemEntity,Entity],
+        response_model=None,
         tags=["relations"],
     )
     def add_nodes(
@@ -345,7 +345,6 @@ def add_nodes(
         """Add a list of nodes to the DKG"""
         for entity in node_list:
             request.app.state.client.add_node(entity)
-        return node_list[0]
 
     @api_blueprint.post(
         "/add_relations",
@@ -354,11 +353,11 @@ def add_nodes(
     )
     def add_relations(
         request: Request,
-        relation_list: List[Dict[str, Any]]
+        relation_list: List[Relation]
     ):
         """Add a list of relations to the DKG"""
         for relation in relation_list:
-            request.app.state.client.relation(relation)
+            request.app.state.client.add_relation(relation)
 
 
 class IsOntChildResult(BaseModel):
diff --git a/mira/dkg/client.py b/mira/dkg/client.py
index 37cf00dd..7411b5a3 100644
--- a/mira/dkg/client.py
+++ b/mira/dkg/client.py
@@ -41,6 +41,17 @@
 TxResult: TypeAlias = Optional[List[List[Any]]]
 
 
+class Relation(BaseModel):
+    """A relationship between two entities in the DKG"""
+    source_curie: str
+    target_curie: str
+    type: str
+    pred: str
+    source: str
+    graph: str
+    version: str
+
+
 class Entity(BaseModel):
     """An entity in the domain knowledge graph."""
 
@@ -313,7 +324,7 @@ def add_node(self, entity):
         Parameters
         ----------
         entity:
-            The node that will be added to the DKG
+            The node object that will be added to the DKG
         """
         curie = entity.id
         name = entity.name
@@ -339,16 +350,34 @@ def add_node(self, entity):
 
         self.create_tx(create_source_node_query)
 
-    def add_relation(self, relation_dict):
+    def add_relation(self, relation):
         """Add a relation to the DKG
 
         Parameters
         ----------
-        relation_dict:
-            The dictionary containing the relationship information
+        relation:
+            The relation object that will be added to the DKG
         """
+        source_curie = relation.source_curie
+        target_curie = relation.target_curie
+        type = relation.type
+        pred = relation.pred
+        source = relation.source
+        version = relation.version
+        graph = relation.graph
+
+        create_relation_query = (
+            f"MATCH (source_node {{curie: '{source_curie}'}}), "
+            f"(target_node {{curie: '{target_curie}'}}) "
+            f"MERGE (source_node)-[rel:{type}]->(target_node)"
+            f"SET rel.pred = '{pred}'"
+            f"SET rel.source = '{source}'"
+            f"SET rel.version = '{version}'"
+            f"SET rel.graph = '{graph}'"
+        )
+        
+        self.create_tx(create_relation_query)
 
-        pass
 
     def create_single_property_node_index(
         self,