Fix nvidia_llm httpx connection leaks and improve crew.py retrieval logging

agents/crew.py

@@ -185,7 +185,7 @@ def run_query_crew(query: str, top_k: int = None, max_tokens: int = None, use_ve
     if status_callback:
         status_callback("inference")
 
-    print(f"\n[Retrieval Phase] Performing vector+graph search for: '{query}'")
+    log.info(f"[Retrieval Phase] Performing vector+graph search for: '{query}'")
     t0 = time.time()
 
     # Import pipeline modules directly for fast retrieval (bypasses CrewAI overhead)
@@ -211,7 +211,7 @@ def run_query_crew(query: str, top_k: int = None, max_tokens: int = None, use_ve
                         r_chunks = vector_store.query_bm25(clean_name, top_k=top_k if top_k is not None else cfg.TOP_K_VECTOR, session_token=session_token)
                         graph_results.extend(r_chunks)
         except Exception as e:
-            print(f"[Retrieval] Graph search failed (non-fatal): {e}")
+            log.warning(f"[Retrieval] Graph search failed (non-fatal): {e}")
     if status_callback:
         status_callback({"status": "graph", "chunks": len(graph_results)})
 
@@ -228,7 +228,7 @@ def run_query_crew(query: str, top_k: int = None, max_tokens: int = None, use_ve
                 session_token=session_token,
             )
         except Exception as e:
-            print(f"[Retrieval] Vector dense search failed: {e}")
+            log.warning(f"[Retrieval] Vector dense search failed: {e}")
     if status_callback:
         status_callback({"status": "vector", "chunks": len(vec_results)})
         
@@ -244,7 +244,7 @@ def run_query_crew(query: str, top_k: int = None, max_tokens: int = None, use_ve
                 session_token=session_token,
             )
         except Exception as e:
-            print(f"[Retrieval] Vector BM25 search failed: {e}")
+            log.warning(f"[Retrieval] Vector BM25 search failed: {e}")
     if status_callback:
         status_callback({"status": "bm25", "chunks": len(bm25_results)})
 
@@ -258,8 +258,8 @@ def run_query_crew(query: str, top_k: int = None, max_tokens: int = None, use_ve
     unique_chunks = list(all_chunks.values())
     
     t_retrieval = time.time() - t0
-    print(f"[Retrieval] Vector chunks: {len(vec_results)}, Graph chunks: {len(graph_results)}, BM25 chunks: {len(bm25_results)}")
-    print(f"[Retrieval Phase] Done in {t_retrieval:.2f}s — "
+    log.info(f"[Retrieval] Vector chunks: {len(vec_results)}, Graph chunks: {len(graph_results)}, BM25 chunks: {len(bm25_results)}")
+    log.info(f"[Retrieval Phase] Done in {t_retrieval:.2f}s — "
           f"{len(unique_chunks)} unique chunks retrieved from Vector DB, Graph DB, and BM25.")
 
     # ── Phase 2: Reranking Agent ───────────────────────────────────────────────
@@ -268,7 +268,7 @@ def run_query_crew(query: str, top_k: int = None, max_tokens: int = None, use_ve
         
     retrieval_is_empty = len(unique_chunks) == 0
 
-    print(f"[Gatekeeper] empty={retrieval_is_empty}")
+    log.info(f"[Gatekeeper] empty={retrieval_is_empty}")
 
     if retrieval_is_empty:
         end_time = time.time()
@@ -300,9 +300,9 @@ def run_query_crew(query: str, top_k: int = None, max_tokens: int = None, use_ve
                 
             # Sort by score descending
             unique_chunks.sort(key=lambda x: x.get("agent_score", -9999.0), reverse=True)
-            print("[Reranking] Successfully reranked chunks using CrossEncoder.")
+            log.info("[Reranking] Successfully reranked chunks using CrossEncoder.")
         except Exception as e:
-            print(f"[Reranking] CrossEncoder Failed: {e}. Proceeding without reranking.")
+            log.warning(f"[Reranking] CrossEncoder Failed: {e}. Proceeding without reranking.")
     
     # Take Final Top 10
     final_top_k = top_k if top_k is not None else cfg.TOP_K_VECTOR

agents/nvidia_llm.py

@@ -39,6 +39,14 @@ def after_request(response):
 # ── Inference Queue (multi-user serialization) ─────────────────────────────────
 _inference_queue: _queue_module.Queue = _queue_module.Queue(maxsize=_QUEUE_MAX_SIZE)
 
+if not NVIDIA_API_KEY:
+    log.warning("NVIDIA_API_KEY is not set. API calls might fail if the token is required.")
+
+_nvidia_client = OpenAI(
+    base_url="https://nim.api.nvidia.com/v1",
+    api_key=NVIDIA_API_KEY or "dummy-key-if-not-required-locally"
+)
+
 def _run_inference(data: dict) -> dict:
     raw_prompt = data.get("prompt", "")
     if not raw_prompt:
@@ -51,14 +59,6 @@ def _run_inference(data: dict) -> dict:
     top_p = float(data.get("top_p", 1.0))
     model_name = data.get("model", DEFAULT_MODEL)
 
-    if not NVIDIA_API_KEY:
-        log.warning("NVIDIA_API_KEY is not set. API calls might fail if the token is required.")
-
-    client = OpenAI(
-        base_url="https://nim.api.nvidia.com/v1",
-        api_key=NVIDIA_API_KEY or "dummy-key-if-not-required-locally"
-    )
-
     choices = []
     
     for i, prompt in enumerate(prompts):
@@ -70,7 +70,7 @@ def _run_inference(data: dict) -> dict:
             ]
 
         try:
-            completion = client.chat.completions.create(
+            completion = _nvidia_client.chat.completions.create(
                 model=model_name,
                 messages=messages,
                 temperature=temperature,
@@ -83,11 +83,15 @@ def _run_inference(data: dict) -> dict:
             print(f"\n[CONSOLE STREAM] Generating via NVIDIA for: {model_name}")
             print("-" * 30)
 
-            for chunk in completion:
-                if chunk.choices and chunk.choices[0].delta.content is not None:
-                    content = chunk.choices[0].delta.content
-                    print(content, end="", flush=True)
-                    full_output += content
+            try:
+                for chunk in completion:
+                    if chunk.choices and chunk.choices[0].delta.content is not None:
+                        content = chunk.choices[0].delta.content
+                        print(content, end="", flush=True)
+                        full_output += content
+            finally:
+                if hasattr(completion, "close"):
+                    completion.close()
 
             print("\n" + "-" * 30)