update oai client to pass all args

Author: alexzhang13

rlm/clients/openai.py

@@ -21,6 +21,10 @@
 class OpenAIClient(BaseLM):
     """
     LM Client for running models with the OpenAI API. Works with vLLM as well.
+
+    Any additional keyword arguments (e.g. default_headers, default_query, max_retries)
+    are passed through to the underlying openai.OpenAI and openai.AsyncOpenAI constructors.
+    Only model_name is excluded, since it is not a client constructor argument.
     """
 
     def __init__(
@@ -42,11 +46,16 @@ def __init__(
             elif base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
                 api_key = DEFAULT_PRIME_API_KEY
 
-        # For vLLM, set base_url to local vLLM server address.
-        self.client = openai.OpenAI(api_key=api_key, base_url=base_url, timeout=self.timeout)
-        self.async_client = openai.AsyncOpenAI(
-            api_key=api_key, base_url=base_url, timeout=self.timeout
-        )
+        # Pass through arbitrary kwargs to the OpenAI client (e.g. default_headers, default_query, max_retries).
+        # Exclude model_name since it is not an OpenAI client constructor argument.
+        client_kwargs = {
+            "api_key": api_key,
+            "base_url": base_url,
+            "timeout": self.timeout,
+            **{k: v for k, v in self.kwargs.items() if k != "model_name"},
+        }
+        self.client = openai.OpenAI(**client_kwargs)
+        self.async_client = openai.AsyncOpenAI(**client_kwargs)
         self.model_name = model_name
         self.base_url = base_url  # Track for cost extraction