Merge pull request #58 from AryaXAI/bugfix_test

Dtree prediction path - API split fix
AryaXAI · Nov 4, 2024 · b90d888 · b90d888
2 parents 1d049d1 + 26d6ae0
commit b90d888
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/aryaxai/common/xai_uris.py b/aryaxai/common/xai_uris.py
@@ -76,6 +76,7 @@
 GET_CASES_URI = f"{API_VERSION_V2}/ai-models/get_cases"
 SEARCH_CASE_URI = f"{API_VERSION_V2}/ai-models/search_case"
 CASE_INFO_URI = f"{API_VERSION_V2}/ai-models/get_case_info"
+CASE_DTREE_URI = f"{API_VERSION_V2}/ai-models/get_case_dtree"
 DELETE_CASE_URI = f"{API_VERSION_V2}/project/delete_data_with_filter"
 CASE_LOGS_URI = f"{API_VERSION_V2}/ai-models/explainability_logs"
 GET_VIEWED_CASE_URI = f"{API_VERSION_V2}/ai-models/get_viewed_case"

diff --git a/aryaxai/core/project.py b/aryaxai/core/project.py
@@ -44,6 +44,7 @@
     AVAILABLE_SYNTHETIC_CUSTOM_SERVERS_URI,
     AVAILABLE_TAGS_URI,
     CASE_INFO_URI,
+    CASE_DTREE_URI,
     CASE_LOGS_URI,
     CLEAR_NOTIFICATIONS_URI,
     CREATE_OBSERVATION_URI,
@@ -2256,6 +2257,18 @@ def case_info(
         if not res["success"]:
             raise Exception(res["details"])
 
+        prediction_path_payload = {
+            "project_name": self.project_name,
+            "unique_identifier": unique_identifer,
+            "case_id": case_id,
+            "model_name": res["details"]["model_name"],
+            "data_id": res["details"]["data_id"],
+        }
+        dtree_res = self.api_client.post(CASE_DTREE_URI, prediction_path_payload)
+        if dtree_res["success"]:
+            res["details"]["case_prediction_svg"] = dtree_res["details"]["case_prediction_svg"]
+            res["details"]["case_prediction_path"] = dtree_res["details"]["case_prediction_path"]
+
         case = Case(**res["details"])
 
         return case