Improve retrieval of deep workspace evidence

SovereignAI-Starter/scripts/eval_knowledge_retrieval.py

@@ -6,9 +6,11 @@ import argparse
 import atexit
 import json
 import shutil
+import sys
 import tempfile
 from datetime import datetime, timezone
 from pathlib import Path
+from urllib.error import HTTPError
 from urllib.request import Request, urlopen
 
 
@@ -28,8 +30,12 @@ def post_json(
     if token:
         headers["Authorization"] = f"Bearer {token}"
     request = Request(url, data=body, headers=headers, method="POST")
-    with urlopen(request, timeout=timeout) as response:
-        return json.loads(response.read().decode("utf-8"))
+    try:
+        with urlopen(request, timeout=timeout) as response:
+            return json.loads(response.read().decode("utf-8"))
+    except HTTPError as error:
+        detail = error.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"HTTP {error.code} from {url}: {detail[:1000]}") from error
 
 
 def delete_json(
@@ -49,8 +55,12 @@ def delete_json(
         headers=headers,
         method="DELETE",
     )
-    with urlopen(request, timeout=timeout) as response:
-        return json.loads(response.read().decode("utf-8"))
+    try:
+        with urlopen(request, timeout=timeout) as response:
+            return json.loads(response.read().decode("utf-8"))
+    except HTTPError as error:
+        detail = error.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"HTTP {error.code} from {url}: {detail[:1000]}") from error
 
 
 def main() -> int:
@@ -59,6 +69,7 @@ def main() -> int:
     parser.add_argument("--dataset", type=Path, default=DATASET)
     parser.add_argument("--output", type=Path, default=None)
     parser.add_argument("--request-timeout", type=float, default=180.0)
+    parser.add_argument("--case", action="append", dest="case_ids")
     args = parser.parse_args()
 
     cases = [
@@ -66,6 +77,11 @@ def main() -> int:
         for line in args.dataset.read_text(encoding="utf-8").splitlines()
         if line.strip()
     ]
+    if args.case_ids:
+        cases = [case for case in cases if case["id"] in args.case_ids]
+        missing = set(args.case_ids) - {case["id"] for case in cases}
+        if missing:
+            parser.error(f"Unknown case ID(s): {', '.join(sorted(missing))}")
     if not cases or len({case["id"] for case in cases}) != len(cases):
         raise ValueError("Retrieval dataset must be non-empty with unique IDs.")
 
@@ -99,7 +115,9 @@ def main() -> int:
             elif not destination.exists():
                 destination.write_text(case["text"], encoding="utf-8")
         index_payload = {"workspace_path": str(workspace_path), "files": files}
+        index_attempted = False
         try:
+            index_attempted = True
             index_result = post_json(
                 index_url, index_payload, timeout=args.request_timeout, token=token
             )
@@ -142,19 +160,34 @@ def main() -> int:
                         "expected_evidence": case["expected_evidence"],
                         "rank": rank,
                         "top_paths": ranked_paths,
+                        "expected_document_chunks": [
+                            item.get("chunk")
+                            for item in retrieved
+                            if item.get("path") == case["document"]
+                        ],
                         "evidence_found": evidence_found,
                         "search_mode": response.get("search_mode", "unknown"),
                     }
                 )
         finally:
-            deletion = delete_json(
-                index_url,
-                index_payload,
-                timeout=args.request_timeout,
-                token=token,
-            )
-            if len(deletion.get("deleted", [])) != len(files):
-                raise RuntimeError("The API did not confirm cleanup for every fixture document.")
+            if index_attempted:
+                original_error = sys.exc_info()[0] is not None
+                try:
+                    deletion = delete_json(
+                        index_url,
+                        index_payload,
+                        timeout=args.request_timeout,
+                        token=token,
+                    )
+                    if len(deletion.get("deleted", [])) != len(files):
+                        raise RuntimeError(
+                            "The API did not confirm cleanup for every fixture document."
+                        )
+                except Exception as cleanup_error:
+                    if original_error:
+                        print(f"Cleanup also failed: {cleanup_error}", file=sys.stderr)
+                    else:
+                        raise
 
     logout()