fix(performance): avoid repeated SentenceTransformer loading in milvus_search

Implemented thread-safe lazy-loading for SentenceTransformer to eliminate
redundant loading within milvus_search.

Signed-off-by: Antigravity <antigravity@google.com>

Author: Ayush-kathil

server-https/app.py

@@ -9,15 +9,23 @@
 from typing import Dict, Any, List, Optional, AsyncGenerator
 from sentence_transformers import SentenceTransformer
 from pymilvus import connections, Collection
+import threading
+import time
 
-
-embedding_model = None
-# Load embedding model once to avoid repeated initialization overhead
+_model_lock = threading.Lock()
+_embedding_model = None
 def get_embedding_model():
-    global embedding_model
-    if embedding_model is None:
-        embedding_model = SentenceTransformer(EMBEDDING_MODEL)
-    return embedding_model
+    """Thread-safe lazy initialization of SentenceTransformer"""
+    global _embedding_model
+    if _embedding_model is None:
+        with _model_lock:
+            # Double-checked locking
+            if _embedding_model is None:
+                start_t = time.perf_counter()
+                print(f"[INFO] Lazy loading SentenceTransformer model '{EMBEDDING_MODEL}'...")
+                _embedding_model = SentenceTransformer(EMBEDDING_MODEL)
+                print(f"[INFO] Model loaded in {time.perf_counter() - start_t:.3f} seconds.")
+    return _embedding_model
 
 # Config
 KSERVE_URL = os.getenv("KSERVE_URL", "http://llama.docs-agent.svc.cluster.local/openai/v1/chat/completions")

server/app.py

@@ -9,6 +9,8 @@
 from typing import Dict, Any, List
 from sentence_transformers import SentenceTransformer
 from pymilvus import connections, Collection
+import threading
+import time
 
 # Config
 KSERVE_URL = os.getenv("KSERVE_URL", "http://llama.docs-agent.svc.cluster.local/openai/v1/chat/completions")
@@ -66,6 +68,66 @@
 
 
 
+_model_lock = threading.Lock()
+_embedding_model = None
+
+def get_embedding_model():
+    """Thread-safe lazy initialization of SentenceTransformer"""
+    global _embedding_model
+    if _embedding_model is None:
+        with _model_lock:
+            # Double-checked locking
+            if _embedding_model is None:
+                start_t = time.perf_counter()
+                print(f"[INFO] Lazy loading SentenceTransformer model '{EMBEDDING_MODEL}'...")
+                _embedding_model = SentenceTransformer(EMBEDDING_MODEL)
+                print(f"[INFO] Model loaded in {time.perf_counter() - start_t:.3f} seconds.")
+    return _embedding_model
+
+def milvus_search(query: str, top_k: int = 5) -> Dict[str, Any]:
+    """Execute a semantic search in Milvus and return structured JSON serializable results."""
+    try:
+        # Connect to Milvus
+        connections.connect(alias="default", host=MILVUS_HOST, port=MILVUS_PORT)
+        collection = Collection(MILVUS_COLLECTION)
+        collection.load()
+
+        # Thread-safe cached encoder
+        model = get_embedding_model()
+        query_vec = model.encode(query).tolist()
+
+        search_params = {"metric_type": "COSINE", "params": {"nprobe": 32}}
+        results = collection.search(
+            data=[query_vec],
+            anns_field=MILVUS_VECTOR_FIELD,
+            param=search_params,
+            limit=int(top_k),
+            output_fields=["file_path", "content_text", "citation_url"],
+        )
+
+        hits = []
+        for hit in results[0]:
+            similarity = 1.0 - float(hit.distance)
+            entity = hit.entity
+            content_text = entity.get("content_text") or ""
+            if isinstance(content_text, str) and len(content_text) > 400:
+                content_text = content_text[:400] + "..."
+            hits.append({
+                "similarity": similarity,
+                "file_path": entity.get("file_path"),
+                "citation_url": entity.get("citation_url"),
+                "content_text": content_text,
+            })
+        return {"results": hits}
+    except Exception as e:
+        print(f"[ERROR] Milvus search failed: {e}")
+        return {"results": []}
+    finally:
+        try:
+            connections.disconnect(alias="default")
+        except Exception:
+            pass
+
 TOOLS = [
     {
         "type": "function",