Associate artifacts with extracted models (#298)

Co-authored-by: Todd Roper <toddroper@me.com>

graph_relations.json

@@ -32,7 +32,8 @@
  ],
  "EXTRACTED_FROM":[
  ["Model","Publication"],
- ["Dataset","Publication"]
+ ["Dataset","Publication"],
+ ["Model", "Artifact"]
  ],
  "CONTAINS":[
  ["Project","Publication"],

tds/db/enums.py

@@ -50,6 +50,7 @@ class ProvenanceSearchTypes(str, Enum):
  parent_model_revisions = "parent_model_revisions"
  parent_models = "parent_models"
  parent_nodes = "parent_nodes"
+ models_from_code = "models_from_code"
 
 
 class RelationType(str, Enum):

tds/db/graph/search_provenance.py

@@ -119,7 +119,6 @@ def set_limit_level(limit):
  """
  )
  response = session.run(query)
-
  return nodes_edges(
  response=response,
  nodes=payload.get("nodes", True),
@@ -315,3 +314,25 @@ def concept_counts(self, payload):
  for key in response:
  counts[key] += 1
  return counts
+
+ def models_from_code(self, payload):
+ """
+ Identifies the code source artifact from which a model was extracted
+ """
+ if payload.get("root_type") not in ("Model"):
+ raise HTTPException(
+ status_code=400,
+ detail="Code artifacts used for model extraction can "
+ "only be found by providing a Model",
+ )
+ with self.graph_db.session() as session:
+ model_id = payload["root_id"]
+
+ query = """
+ MATCH (a:Artifact)<-[r:EXTRACTED_FROM]-(m:Model {id: $model_id})
+ RETURN a
+ """
+
+ response = session.run(query, {"model_id": model_id})
+ response_data = [res.data()["a"]["id"] for res in response]
+ return response_data

tds/modules/provenance/controller.py

@@ -93,6 +93,11 @@ def search_provenance(
  * Requirements: “root_type”, “root_id”
  * Allowed root _types are Model *will be expanded.
 
+ **models_from_code** - Returns the artifact `id` for the artifact
+ from which a model was extracted
+ * Requirements: "root_type", “root_id”
+ * Allowed `root_types` are Model
+
  ## Payload format
 
  The payload for searching needs to match the schema below.

tds/modules/provenance/model.py

@@ -59,7 +59,7 @@ class ProvenanceSearch(BaseModel):
  Provenance Data Model.
  """
 
- root_id: Optional[int]
+ root_id: Optional[int | str]
  root_type: Optional[ProvenanceType]
  user_id: Optional[int]
  curie: Optional[str]

tds/schema/provenance.py

@@ -48,4 +48,5 @@ class ProvenancePayload(BaseModel):
  ProvenanceType.Simulation: "Si",
  ProvenanceType.Project: "Pr",
  ProvenanceType.Concept: "Cn",
+ ProvenanceType.Artifact: "Ar",
 }