fix: make LLMDescriptionScorer work inside a running event loop

src/autointent/modules/scoring/_description/llm_encoder.py

@@ -5,10 +5,11 @@
 import asyncio
 import json
 import logging
+from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from pathlib import Path
 from textwrap import dedent
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import aiometer
 import numpy as np
@@ -23,12 +24,17 @@
 from .base import BaseDescriptionScorer
 
 if TYPE_CHECKING:
+    from collections.abc import Coroutine
+
     from numpy.typing import NDArray
 
     from autointent.configs import CrossEncoderConfig, EmbedderConfig
 
 logger = logging.getLogger(__name__)
 
+
+_T = TypeVar("_T")
+
 GENERATOR_CONFIG_FILENAME = "generator_config.json"
 
 
@@ -249,7 +255,7 @@ def _compute_similarities(self, utterances: list[str]) -> NDArray[np.float64]:
                 max_at_once=self.max_concurrent,
                 max_per_second=self.max_per_second,
             )
-            categorizations = self._event_loop.run_until_complete(task)  # type: ignore[arg-type]
+            categorizations = self._run_async(task)  # type: ignore[arg-type]
 
         for i, categorization in enumerate(categorizations):
             if isinstance(categorization, IntentCategorization):
@@ -276,19 +282,30 @@ def clear_cache(self) -> None:
         # Generator doesn't have a clear_ram method, so we just set it to None
         if hasattr(self, "_generator"):
             delattr(self, "_generator")
-        if hasattr(self, "_event_loop"):
-            delattr(self, "_event_loop")
+        self._close_event_loop()
 
     def _init_event_loop(self) -> None:
+        # A private loop, never the caller's: the generator's async client keeps its
+        # connections bound to the loop it first ran on, so all batches reuse this one.
+        self._close_event_loop()
         if self.max_concurrent is not None:
-            try:
-                loop = asyncio.get_event_loop()
-            except RuntimeError:
-                loop = asyncio.new_event_loop()
-            else:
-                if loop.is_closed():
-                    loop = asyncio.new_event_loop()
-            self._event_loop = loop
+            self._event_loop = asyncio.new_event_loop()
+
+    def _close_event_loop(self) -> None:
+        if hasattr(self, "_event_loop"):
+            self._event_loop.close()
+            delattr(self, "_event_loop")
+
+    def _run_async(self, coro: Coroutine[Any, Any, _T]) -> _T:
+        """Run ``coro`` to completion on the scorer's loop, whether or not the caller is already in one."""
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return self._event_loop.run_until_complete(coro)
+        # Called from async code (notebook, async web handler): this thread's loop is busy,
+        # so drive the scorer's loop from a worker thread and block until it finishes.
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            return executor.submit(self._event_loop.run_until_complete, coro).result()
 
     def dump(self, path: str) -> None:
         dump_path = Path(path)

tests/modules/scoring/test_description_llm.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import tempfile
 from typing import TYPE_CHECKING, Any, cast
 
@@ -112,3 +113,44 @@ def test_llm_description_in_pipeline(dataset: Dataset, patch_llm_scorer_generato
     pipeline.fit(dataset)
     predictions = pipeline.predict(["test utterance"])
     assert len(predictions) == 1
+
+
+def _fit_llm_scorer(dataset: Dataset) -> LLMDescriptionScorer:
+    data_handler = DataHandler(dataset)
+    scorer = LLMDescriptionScorer(generator_config={"temperature": 0})
+    labels = data_handler.train_labels(0)
+    assert is_strict_labels(labels)
+    descriptions = data_handler.intent_descriptions
+    assert all(d is not None for d in descriptions)
+    scorer.fit(data_handler.train_utterances(0), labels, cast("list[str]", descriptions))
+    return scorer
+
+
+def test_description_scorer_llm_predict_inside_running_loop(
+    dataset: Dataset, patch_llm_scorer_generator: Generator
+) -> None:
+    """predict() is sync, but must also work when called from async code (notebooks, async servers)."""
+    scorer = _fit_llm_scorer(dataset)
+    utterances = ["What is the balance on my account?", "How do I reset my online banking password?"]
+    expected = scorer.predict(utterances)
+
+    async def predict_from_coroutine() -> npt.NDArray[Any]:
+        return scorer.predict(utterances)
+
+    np.testing.assert_array_equal(asyncio.run(predict_from_coroutine()), expected)
+    # the scorer's loop is still usable from sync code afterwards
+    np.testing.assert_array_equal(scorer.predict(utterances), expected)
+
+
+def test_description_scorer_llm_closes_its_event_loop(dataset: Dataset, patch_llm_scorer_generator: Generator) -> None:
+    scorer = _fit_llm_scorer(dataset)
+    first_loop = scorer._event_loop
+
+    scorer.fit([], [], scorer._description_texts)
+    assert first_loop.is_closed()
+    assert scorer._event_loop is not first_loop
+
+    second_loop = scorer._event_loop
+    scorer.clear_cache()
+    assert second_loop.is_closed()
+    assert not hasattr(scorer, "_event_loop")