From ded0701504de1c687761f588969cd8a5965a06be Mon Sep 17 00:00:00 2001
From: Andy Lee <andylizf@outlook.com>
Date: Thu, 14 Aug 2025 14:29:57 -0700
Subject: [PATCH] core: auto-cleanup for LeannSearcher/LeannChat
 (__enter__/__exit__/__del__); ensure server terminate/kill robustness;
 benchmarks: use searcher.cleanup(); docs: suggest uv run

---
 .../diskann_vs_hnsw_speed_comparison.py       |  6 ++--
 docs/configuration-guide.md                   |  2 +-
 packages/leann-core/src/leann/api.py          | 33 +++++++++++++++++++
 .../src/leann/embedding_server_manager.py     | 13 ++++++--
 4 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/benchmarks/diskann_vs_hnsw_speed_comparison.py b/benchmarks/diskann_vs_hnsw_speed_comparison.py
index bb89692..c287afa 100644
--- a/benchmarks/diskann_vs_hnsw_speed_comparison.py
+++ b/benchmarks/diskann_vs_hnsw_speed_comparison.py
@@ -113,10 +113,10 @@ def benchmark_backend(
         ]
         score_validity_rate = len(valid_scores) / len(all_scores) if all_scores else 0
 
-        # Clean up
+        # Clean up (ensure embedding server shutdown and object GC)
         try:
-            if hasattr(searcher, "__del__"):
-                searcher.__del__()
+            if hasattr(searcher, "cleanup"):
+                searcher.cleanup()
             del searcher
             del builder
             gc.collect()
diff --git a/docs/configuration-guide.md b/docs/configuration-guide.md
index 036c1c2..c1402a1 100644
--- a/docs/configuration-guide.md
+++ b/docs/configuration-guide.md
@@ -113,7 +113,7 @@ ollama pull nomic-embed-text
 --backend-name diskann --graph-degree 32 --build-complexity 64
 ```
 
-**Performance Benchmark**: Run `python benchmarks/diskann_vs_hnsw_speed_comparison.py` to compare DiskANN and HNSW on your system.
+**Performance Benchmark**: Run `uv run benchmarks/diskann_vs_hnsw_speed_comparison.py` to compare DiskANN and HNSW on your system.
 
 ## LLM Selection: Engine and Model Comparison
 
diff --git a/packages/leann-core/src/leann/api.py b/packages/leann-core/src/leann/api.py
index ba84014..144e858 100644
--- a/packages/leann-core/src/leann/api.py
+++ b/packages/leann-core/src/leann/api.py
@@ -665,6 +665,23 @@ class LeannSearcher:
         if hasattr(self.backend_impl, "embedding_server_manager"):
             self.backend_impl.embedding_server_manager.stop_server()
 
+    # Enable automatic cleanup patterns
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        try:
+            self.cleanup()
+        except Exception:
+            pass
+
+    def __del__(self):
+        try:
+            self.cleanup()
+        except Exception:
+            # Avoid noisy errors during interpreter shutdown
+            pass
+
 
 class LeannChat:
     def __init__(
@@ -743,3 +760,19 @@ class LeannChat:
         """
         if hasattr(self.searcher, "cleanup"):
             self.searcher.cleanup()
+
+    # Enable automatic cleanup patterns
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        try:
+            self.cleanup()
+        except Exception:
+            pass
+
+    def __del__(self):
+        try:
+            self.cleanup()
+        except Exception:
+            pass
diff --git a/packages/leann-core/src/leann/embedding_server_manager.py b/packages/leann-core/src/leann/embedding_server_manager.py
index 3ed223f..05c8639 100644
--- a/packages/leann-core/src/leann/embedding_server_manager.py
+++ b/packages/leann-core/src/leann/embedding_server_manager.py
@@ -268,8 +268,12 @@ class EmbeddingServerManager:
             f"Terminating server process (PID: {self.server_process.pid}) for backend {self.backend_module_name}..."
         )
 
-        # Use simple termination - our improved server shutdown should handle this properly
-        self.server_process.terminate()
+        # Use simple termination first; if the server installed signal handlers,
+        # it will exit cleanly. Otherwise escalate to kill after a short wait.
+        try:
+            self.server_process.terminate()
+        except Exception:
+            pass
 
         try:
             self.server_process.wait(timeout=5)  # Give more time for graceful shutdown
@@ -278,7 +282,10 @@ class EmbeddingServerManager:
             logger.warning(
                 f"Server process {self.server_process.pid} did not terminate within 5 seconds, force killing..."
             )
-            self.server_process.kill()
+            try:
+                self.server_process.kill()
+            except Exception:
+                pass
             try:
                 self.server_process.wait(timeout=2)
                 logger.info(f"Server process {self.server_process.pid} killed successfully.")