Add a sample function that computes a batch of template results in parallel

effectful/handlers/llm/sampling.py

@@ -1,11 +1,13 @@
+import threading
 from collections import Counter
 from concurrent import futures
 from concurrent.futures.thread import ThreadPoolExecutor
 
 from effectful.handlers.llm import Template
+from effectful.handlers.llm.providers import completion, tool_call
 from effectful.internals.runtime import get_interpretation, interpreter
-from effectful.ops.semantics import fwd
-from effectful.ops.syntax import ObjectInterpretation, implements
+from effectful.ops.semantics import fwd, handler
+from effectful.ops.syntax import ObjectInterpretation, defop, implements
 
 
 class KAheadSampler[**P, T](ObjectInterpretation):
@@ -45,3 +47,55 @@ def n_votes_ahead():
                 tasks.append(executor.submit(interpreter(intp)(fwd), *args, **kwargs))
         executor.shutdown()
         return self.votes.most_common(1)[0][0]
+
+
+def sample(template, n):
+    @defop
+    def in_nested_call() -> bool:
+        return False
+
+    lock = threading.Lock()
+
+    def _template_call(template, *args, **kwargs):
+        if in_nested_call():
+            return fwd()
+
+        with handler({in_nested_call: lambda: True}):
+            with ThreadPoolExecutor() as executor:
+
+                @interpreter(get_interpretation())
+                def do_work():
+                    lock.acquire()
+                    try:
+                        result = fwd()
+                    finally:
+                        assert lock.locked()
+                        lock.release()
+                    return result
+
+                tasks = [executor.submit(do_work) for _ in range(n)]
+                completed = futures.wait(tasks, return_when=futures.ALL_COMPLETED)
+                return [t.result() for t in completed.done]
+
+    def _completion(*args, **kwargs):
+        lock.release()
+        result = fwd()
+        lock.acquire()
+        return result
+
+    def _tool_call(*args, **kwargs):
+        lock.acquire()
+        try:
+            result = fwd()
+        except Exception as e:
+            lock.release()
+            raise e
+        return result
+
+    return handler(
+        {
+            Template.__call__: _template_call,
+            completion: _completion,
+            tool_call: _tool_call,
+        }
+    )(template)