adapt collect_training_code to use VespaResult

1 files changed, 2 insertions, 10 deletions

diff --git a/python/vespa/vespa/application.py b/python/vespa/vespa/application.py
index 62839de92de..3646cf87bf2 100644
--- a/python/vespa/vespa/application.py
+++ b/python/vespa/vespa/application.py
@@ -114,7 +114,7 @@ class Vespa(object):
             recall=(id_field, [relevant_id]),
             **kwargs
         )
-        hits = get_hits(vespa_result=relevant_id_result)
+        hits = relevant_id_result.hits
         features = []
         if len(hits) == 1 and hits[0]["fields"][id_field] == relevant_id:
             random_hits_result = self.query(
@@ -123,7 +123,7 @@ class Vespa(object):
                 hits=number_additional_docs,
                 **kwargs
             )
-            hits.extend(get_hits(random_hits_result))
+            hits.extend(random_hits_result.hits)
 
             features = annotate_data(
                 hits=hits,
@@ -248,14 +248,6 @@ class Vespa(object):
         return evaluation
 
 
-# todo: create a VespaResult class to store vespa results
-def get_hits(vespa_result):
-    hits = []
-    if "children" in vespa_result["root"]:
-        hits = vespa_result["root"]["children"]
-    return hits
-
-
 # todo: a better pattern for labelled data would be (query_id, query, doc_id, score) with the possibility od
 #  assigning a specific default value for those docs not mentioned
 def annotate_data(hits, query_id, id_field, relevant_id, relevant_score, default_score):