feat: supported streaming with multiple candidates (#158)

epam · Dec 18, 2024 · 2996c7e · 2996c7e
1 parent 8b0de00
commit 2996c7e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 13 deletions.
diff --git a/aidial_adapter_vertexai/chat_completion.py b/aidial_adapter_vertexai/chat_completion.py
@@ -66,9 +66,6 @@ async def chat_completion(self, request: Request, response: Response):
         n = params.n or 1
         params.n = None
 
-        if n > 1 and params.stream:
-            raise ValidationError("n>1 is not supported in streaming mode")
-
         if params.max_prompt_tokens is None:
             truncated_prompt = TruncatedPrompt(
                 prompt=prompt, discarded_messages=[]

diff --git a/tests/integration_tests/test_chat_completion.py b/tests/integration_tests/test_chat_completion.py
@@ -227,16 +227,8 @@ def test_case(
             name="multiple candidates",
             max_tokens=10,
             n=5,
-            messages=[user("2+3=?")],
-            expected=(
-                ExpectedException(
-                    type=UnprocessableEntityError,
-                    message="n>1 is not supported in streaming mode",
-                    status_code=422,
-                )
-                if streaming
-                else for_all_choices(lambda _: True, 5)
-            ),
+            messages=[user("2+7=? Reply with a single number")],
+            expected=for_all_choices(lambda s: "9" in s, 5),
         )
 
         # Stop sequences do not work for some reason for CHAT_BISON_2_32K and streaming mode