fix: keep agent streams alive during long tasks

SovereignAI-Starter/scripts/eval_knowledge_retrieval.py

@@ -4,10 +4,12 @@ from __future__ import annotations
 
 import argparse
 import atexit
+import hashlib
 import json
 import shutil
 import sys
 import tempfile
+from contextlib import nullcontext
 from datetime import datetime, timezone
 from pathlib import Path
 from urllib.error import HTTPError
@@ -68,6 +70,15 @@ def main() -> int:
     parser.add_argument("--base-url", default="http://127.0.0.1:8000")
     parser.add_argument("--dataset", type=Path, default=DATASET)
     parser.add_argument("--output", type=Path, default=None)
+    parser.add_argument(
+        "--workspace-dir",
+        type=Path,
+        default=None,
+        help=(
+            "Use an existing empty directory beside or inside the project for fixtures. "
+            "Only fixture files whose contents remain unchanged are removed afterward."
+        ),
+    )
     parser.add_argument("--request-timeout", type=float, default=180.0)
     parser.add_argument("--case", action="append", dest="case_ids")
     parser.add_argument(
@@ -90,6 +101,23 @@ def main() -> int:
     if not cases or len({case["id"] for case in cases}) != len(cases):
         raise ValueError("Retrieval dataset must be non-empty with unique IDs.")
 
+    if args.workspace_dir is None:
+        workspace_context = tempfile.TemporaryDirectory(
+            prefix="sovereignai-rag-eval-"
+        )
+    else:
+        workspace_path = args.workspace_dir.resolve(strict=True)
+        if not (
+            workspace_path.is_relative_to(ROOT)
+            or workspace_path.parent == ROOT.parent
+        ):
+            parser.error(
+                "--workspace-dir must be inside the project or its immediate parent."
+            )
+        if not workspace_path.is_dir() or any(workspace_path.iterdir()):
+            parser.error("--workspace-dir must be an existing empty directory.")
+        workspace_context = nullcontext(str(workspace_path))
+
     results: list[dict] = []
     index_summary: dict[str, object] = {}
     files = sorted({case["document"] for case in cases})
@@ -107,10 +135,18 @@ def main() -> int:
             pass
 
     atexit.register(logout)
-    with tempfile.TemporaryDirectory(prefix="sovereignai-rag-eval-") as temporary:
+    with workspace_context as temporary:
         workspace_path = Path(temporary)
+        fixture_digests: dict[Path, str] = {}
         for case in cases:
-            destination = workspace_path / case["document"]
+            relative_document = Path(case["document"])
+            if (
+                relative_document.is_absolute()
+                or ".." in relative_document.parts
+                or not relative_document.parts
+            ):
+                parser.error(f"Unsafe fixture document path: {case['document']!r}")
+            destination = workspace_path / relative_document
             destination.parent.mkdir(parents=True, exist_ok=True)
             source_path = case.get("source_path")
             if source_path:
@@ -120,6 +156,9 @@ def main() -> int:
                 shutil.copyfile(source, destination)
             elif not destination.exists():
                 destination.write_text(case["text"], encoding="utf-8")
+            fixture_digests[relative_document] = hashlib.sha256(
+                destination.read_bytes()
+            ).hexdigest()
         index_payload = {"workspace_path": str(workspace_path), "files": files}
         index_attempted = False
         project_registration_attempted = False
@@ -138,12 +177,24 @@ def main() -> int:
             if Path(registration.get("path", "")).resolve() != workspace_path.resolve():
                 raise RuntimeError("The API registered a different fixture workspace.")
             index_attempted = True
-            index_result = post_json(
-                index_url, index_payload, timeout=args.request_timeout, token=token
-            )
-            if len(index_result.get("indexed", [])) != len(files):
-                raise RuntimeError("The API did not confirm every fixture document.")
-            indexed_documents = index_result["indexed"]
+            indexed_documents = []
+            for file_number, relative_path in enumerate(files, start=1):
+                print(
+                    f"Indexing fixture {file_number}/{len(files)}: {relative_path}",
+                    flush=True,
+                )
+                index_result = post_json(
+                    index_url,
+                    {"workspace_path": str(workspace_path), "files": [relative_path]},
+                    timeout=args.request_timeout,
+                    token=token,
+                )
+                indexed = index_result.get("indexed", [])
+                if len(indexed) != 1:
+                    raise RuntimeError(
+                        f"The API did not confirm fixture document {relative_path}."
+                    )
+                indexed_documents.extend(indexed)
             index_summary = {
                 "semantic_indexed_documents": sum(
                     bool(item.get("semantic_indexed")) for item in indexed_documents
@@ -240,6 +291,42 @@ def main() -> int:
                         )
                 except Exception as cleanup_error:
                     cleanup_errors.append(cleanup_error)
+            if args.workspace_dir is not None:
+                for relative_document, expected_digest in fixture_digests.items():
+                    destination = workspace_path / relative_document
+                    try:
+                        if not destination.is_file():
+                            raise RuntimeError(
+                                f"Evaluation fixture disappeared: {relative_document}"
+                            )
+                        actual_digest = hashlib.sha256(destination.read_bytes()).hexdigest()
+                        if actual_digest != expected_digest:
+                            raise RuntimeError(
+                                f"Evaluation fixture changed; leaving it in place: {relative_document}"
+                            )
+                        destination.unlink()
+                    except Exception as cleanup_error:
+                        cleanup_errors.append(cleanup_error)
+                parents = sorted(
+                    {
+                        (workspace_path / relative_document).parent
+                        for relative_document in fixture_digests
+                    },
+                    key=lambda path: len(path.parts),
+                    reverse=True,
+                )
+                for parent in parents:
+                    if parent == workspace_path:
+                        continue
+                    try:
+                        parent.rmdir()
+                    except OSError:
+                        if parent.exists():
+                            cleanup_errors.append(
+                                RuntimeError(
+                                    f"Evaluation directory is not empty; leaving it in place: {parent}"
+                                )
+                            )
             if cleanup_errors:
                 if original_error:
                     for cleanup_error in cleanup_errors: