add debug flag to yql builder

src/cpr_sdk/models/search.py

@@ -375,6 +375,7 @@ class Hit(BaseModel):
     corpus_import_id: Optional[str] = None
     metadata: Optional[Sequence[dict[str, str]]] = None
     concepts: Optional[Sequence[Concept]] = None
+    debug_info: Optional[dict] = None
 
     @classmethod
     def from_vespa_response(cls, response_hit: dict) -> "Hit":
@@ -419,6 +420,8 @@ def from_vespa_response(cls, response_hit: dict) -> "Document":
             if family_publication_ts
             else None
         )
+        debug_info = {k: v for k, v in fields.items() if "token" in k}
+
         return cls(
             family_name=fields.get("family_name"),
             family_description=fields.get("family_description"),
@@ -439,6 +442,7 @@ def from_vespa_response(cls, response_hit: dict) -> "Document":
             corpus_import_id=fields.get("corpus_import_id"),
             metadata=fields.get("metadata"),
             concepts=fields.get("concepts"),
+            debug_info=debug_info,
         )
 
 
@@ -467,6 +471,8 @@ def from_vespa_response(cls, response_hit: dict) -> "Passage":
             else None
         )
 
+        debug_info = {k: v for k, v in fields.items() if "token" in k}
+
         return cls(
             family_name=fields.get("family_name"),
             family_description=fields.get("family_description"),
@@ -492,6 +498,7 @@ def from_vespa_response(cls, response_hit: dict) -> "Passage":
             text_block_coords=fields.get("text_block_coords"),
             metadata=fields.get("metadata"),
             concepts=fields.get("concepts"),
+            debug_info=debug_info,
         )
 
 

src/cpr_sdk/search_adaptors.py

@@ -69,15 +69,17 @@ def __init__(
 
         self.client = Vespa(url=instance_url, cert=cert_path, key=key_path)
 
-    def search(self, parameters: SearchParameters) -> SearchResponse:
+    def search(
+        self, parameters: SearchParameters, debug: bool = False
+    ) -> SearchResponse:
         """
         Search a vespa instance
 
         :param SearchParameters parameters: a search request object
         :return SearchResponse: a list of families, with response metadata
         """
         total_time_start = time.time()
-        vespa_request_body = build_vespa_request_body(parameters)
+        vespa_request_body = build_vespa_request_body(parameters, debug=debug)
         query_time_start = time.time()
         try:
             vespa_response = self.client.query(body=vespa_request_body)

src/cpr_sdk/vespa.py

@@ -80,11 +80,13 @@ def find_vespa_cert_paths() -> tuple[Optional[str], Optional[str]]:
     return cert_path, key_path
 
 
-def build_vespa_request_body(parameters: SearchParameters) -> dict[str, str]:
+def build_vespa_request_body(
+    parameters: SearchParameters, debug: bool = False
+) -> dict[str, str]:
     """Constructs the payload for a vespa query"""
     sensitive = is_sensitive_query(parameters.query_string, SENSITIVE_QUERY_TERMS)
 
-    yql = YQLBuilder(params=parameters, sensitive=sensitive).to_str()
+    yql = YQLBuilder(params=parameters, sensitive=sensitive, debug=debug).to_str()
     vespa_request_body: dict[str, Any] = {
         "yql": yql,
         "timeout": "20",

src/cpr_sdk/yql_builder.py

@@ -24,17 +24,25 @@ class YQLBuilder:
                 max($MAX_HITS_PER_FAMILY)
                 each(
                     output(
-                        summary(search_summary)
+                        summary($SUMMARY)
                     )
                 )
             )
         )
     """
     )
 
-    def __init__(self, params: SearchParameters, sensitive: bool = False) -> None:
+    def __init__(
+        self, params: SearchParameters, sensitive: bool = False, debug: bool = False
+    ) -> None:
         self.params = params
         self.sensitive = sensitive
+        self.debug = debug
+
+    def build_summary(self) -> str:
+        """Creates the part of the query that determines which fields to return"""
+
+        return "search_summary_with_tokens" if self.debug else "search_summary"
 
     def build_sources(self) -> str:
         """Creates the part of the query that determines which sources to search"""
@@ -220,6 +228,7 @@ def to_str(self) -> str:
             LIMIT=self.build_limit(),
             SORT=self.build_sort(),
             MAX_HITS_PER_FAMILY=self.build_max_hits_per_family(),
+            SUMMARY=self.build_summary(),
         )
         return " ".join(yql.split())