Feat/completions n

python/freetoken/server/api_server.py

@@ -375,16 +375,23 @@ async def stream_generate(self, uid: int):
         yield b"data: [DONE]\n\n"
         logger.debug("Finished streaming response for user %s", uid)
 
-    async def stream_with_cancellation(self, generator, request: Request, uid: int):
+    async def stream_with_cancellation(
+        self, generator, request: Request, uid: int | Callable[[], list[int]]
+    ):
+        """``uid`` may be a callable returning the uids to abort, for a stream (OpenAI ``n``)
+        whose requests are submitted after it starts."""
         try:
             async for chunk in generator:
                 # detect if the client has disconnected
                 if await request.is_disconnected():
-                    logger.info("Client disconnected for user %s", uid)
+                    logger.info(
+                        "Client disconnected for user %s", uid if isinstance(uid, int) else "(n > 1)"
+                    )
                     raise asyncio.CancelledError
                 yield chunk
         except asyncio.CancelledError:
-            asyncio.create_task(self.abort_user(uid))
+            for live_uid in uid() if callable(uid) else [uid]:
+                asyncio.create_task(self.abort_user(live_uid))
             raise
 
     async def abort_user(self, uid: int):

python/freetoken/server/generation.py

@@ -16,7 +16,7 @@
 import asyncio
 import json
 import time
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Awaitable, Callable
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -507,6 +507,109 @@ async def with_keepalive(events: AsyncIterator[GenEvent], interval: float):
             task.cancel()
 
 
+_CHOICE_DONE = object()
+
+
+@dataclass
+class _ChoiceFailed:
+    exc: Exception
+
+
+class ChoiceFanOut:
+    """Runs ``n`` generations of one prompt (OpenAI ``n``) and merges their outputs as
+    ``(choice_index, item)`` pairs. Choice 0 is submitted first and the rest follow its first
+    output for which ``releases_siblings`` holds: by then its prompt is in the prefix cache,
+    so they skip a second prefill.
+
+    ``submit(uid)`` sends one request; ``consume(uid)`` yields its outputs, ends at its
+    terminal one and raises on failure. Closing ``events()`` early (error, disconnect)
+    aborts every choice not yet finished, including ones never submitted."""
+
+    def __init__(
+        self,
+        n: int,
+        state: Any,
+        submit: Callable[[int], Awaitable[None]],
+        consume: Callable[[int], AsyncIterator[Any]],
+        releases_siblings: Callable[[Any], bool],
+    ) -> None:
+        self._n = n
+        self._state = state
+        self._submit = submit
+        self._consume = consume
+        self._releases_siblings = releases_siblings
+        self._uids: list[int] = []
+        self._live: set[int] = set()
+        self._closed = False
+
+    async def start(self) -> None:
+        """Admit every choice and submit choice 0. Admission is decided once for the whole
+        request, so it fails here, before a stream opens, rather than mid-response."""
+        try:
+            for _ in range(self._n):
+                uid = self._state.new_user()
+                self._uids.append(uid)
+                self._live.add(uid)
+            await self._submit(self._uids[0])
+        except BaseException:
+            self.abort_live()
+            raise
+
+    def take_live_uids(self) -> list[int]:
+        """Stop launching choices and hand over the uids not yet finished, to abort."""
+        self._closed = True
+        uids = list(self._live)
+        self._live.clear()
+        return uids
+
+    def abort_live(self) -> None:
+        """Abort every choice not yet finished; a no-op once ``events()`` has cleaned up."""
+        for uid in self.take_live_uids():
+            asyncio.create_task(self._state.abort_user(uid))
+
+    async def events(self) -> AsyncIterator[tuple[int, Any]]:
+        if not self._uids:
+            await self.start()
+        queue: asyncio.Queue[tuple[int, Any]] = asyncio.Queue()
+        released = asyncio.Event()
+
+        async def run(index: int) -> None:
+            uid = self._uids[index]
+            try:
+                if index:
+                    await released.wait()
+                    if self._closed:
+                        return
+                    await self._submit(uid)
+                async for item in self._consume(uid):
+                    if self._releases_siblings(item):
+                        released.set()
+                    queue.put_nowait((index, item))
+                self._live.discard(uid)
+            except Exception as exc:  # noqa: BLE001 -- re-raised by the consumer below
+                self._closed = True
+                queue.put_nowait((index, _ChoiceFailed(exc)))
+            finally:
+                released.set()
+                queue.put_nowait((index, _CHOICE_DONE))
+
+        tasks = [asyncio.create_task(run(index)) for index in range(self._n)]
+        try:
+            pending = self._n
+            while pending:
+                index, item = await queue.get()
+                if item is _CHOICE_DONE:
+                    pending -= 1
+                elif isinstance(item, _ChoiceFailed):
+                    raise item.exc
+                else:
+                    yield index, item
+        finally:
+            for task in tasks:
+                task.cancel()
+            self.abort_live()
+
+
 def _record_generation(
     *,
     source: str | None,