[Tests] Disable retries and use context manager for openai client (#7565)

2024-08-26 21:33:17 -07:00
parent 2eedede875
commit 39178c7fbc
15 changed files with 130 additions and 93 deletions
--- a/tests/multi_step/test_correctness_async_llm.py
+++ b/tests/multi_step/test_correctness_async_llm.py
@@ -28,12 +28,12 @@ async def completions_with_server_args(prompts: List[str], model_name: str,

    outputs = None
    with RemoteOpenAIServer(model_name, server_cli_args) as server:
-        client = server.get_async_client()
-        outputs = await client.completions.create(model=model_name,
-                                                  prompt=prompts,
-                                                  temperature=0,
-                                                  stream=False,
-                                                  max_tokens=5)
+        async with server.get_async_client() as client:
+            outputs = await client.completions.create(model=model_name,
+                                                      prompt=prompts,
+                                                      temperature=0,
+                                                      stream=False,
+                                                      max_tokens=5)
    assert outputs is not None

    return outputs