Model added to similarity client

materialsproject · May 21, 2021 · 6dd8f80 · 6dd8f80
1 parent 8841c26
commit 6dd8f80
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/mp_api/routes/similarity/client.py b/src/mp_api/routes/similarity/client.py
@@ -1,7 +1,9 @@
 from mp_api.core.client import BaseRester
+from mp_api.routes.similarity.models import SimilarityDoc
 
 
 class SimilarityRester(BaseRester):
 
     suffix = "similarity"
+    document_model = SimilarityDoc  # type: ignore
     primary_key = "task_id"