From 3dc04bbf6279c2898dc229b35d5191d0d2659287 Mon Sep 17 00:00:00 2001 From: whoisnnamdi Date: Sun, 4 Jan 2026 11:27:02 -0800 Subject: [PATCH] fix: use async clients for acompletion methods --- rlm/clients/openai.py | 3 ++- rlm/clients/portkey.py | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/rlm/clients/openai.py b/rlm/clients/openai.py index 268e77ea9..d5fa99966 100644 --- a/rlm/clients/openai.py +++ b/rlm/clients/openai.py @@ -37,6 +37,7 @@ def __init__( # For vLLM, set base_url to local vLLM server address. self.client = openai.OpenAI(api_key=api_key, base_url=base_url) + self.async_client = openai.AsyncOpenAI(api_key=api_key, base_url=base_url) self.model_name = model_name # Per-model usage tracking @@ -75,7 +76,7 @@ async def acompletion( if not model: raise ValueError("Model name is required for OpenAI client.") - response = await self.client.chat.completions.create(model=model, messages=messages) + response = await self.async_client.chat.completions.create(model=model, messages=messages) self._track_cost(response, model) return response.choices[0].message.content diff --git a/rlm/clients/portkey.py b/rlm/clients/portkey.py index fa57d7025..6835c7404 100644 --- a/rlm/clients/portkey.py +++ b/rlm/clients/portkey.py @@ -1,7 +1,7 @@ from collections import defaultdict from typing import Any -from portkey_ai import Portkey +from portkey_ai import AsyncPortkey, Portkey from portkey_ai.api_resources.types.chat_complete_type import ChatCompletions from rlm.clients.base_lm import BaseLM @@ -22,6 +22,7 @@ def __init__( ): super().__init__(model_name=model_name, **kwargs) self.client = Portkey(api_key=api_key, base_url=base_url) + self.async_client = AsyncPortkey(api_key=api_key, base_url=base_url) self.model_name = model_name # Per-model usage tracking @@ -61,7 +62,7 @@ async def acompletion(self, prompt: str | dict[str, Any], model: str | None = No if not model: raise ValueError("Model name is required for Portkey client.") - response = await self.client.chat.completions.create(model=model, messages=messages) + response = await self.async_client.chat.completions.create(model=model, messages=messages) self._track_cost(response, model) return response.choices[0].message.content