[Core] *Prompt* logprobs support in Multi-step (#8199)

2026-08-20 12:50:59 +08:00 · 2024-09-18 08:38:43 -07:00
parent 7c7714d856
commit a8c1d161a7
5 changed files with 299 additions and 58 deletions
@@ -493,6 +493,7 @@ async def completions_with_server_args(
    '''

    outputs = None
+    max_wait_seconds = 240 * 3  # 240 is default
    with RemoteOpenAIServer(model_name,
                            server_cli_args,
                            max_wait_seconds=max_wait_seconds) as server:
@@ -503,7 +504,7 @@ async def completions_with_server_args(
                                                  stream=False,
                                                  max_tokens=5,
                                                  logprobs=num_logprobs)
-    assert outputs is not None
+    assert outputs is not None, "Completion API call failed."

    return outputs