Refactor OpenAI and Portkey clients to support new prompt formats and improve error handling

- Updated `OpenAIClient` and `PortkeyClient` to accept prompts as strings, dictionaries, or lists of dictionaries.
- Enhanced completion methods to return empty strings for empty responses.
- Improved error handling in `send_lm_request_batched` to provide clearer error messages.
- Added assertions in `LMRequestHandler` to ensure prompts are not None.
- Modified token tracking in `PortkeyClient` to handle optional usage attributes.
- Updated various REPL environments to handle cases where no response is returned from the LM.
- Enhanced logging to capture limit hit events and execution times.
- Adjusted tests to reflect changes in prompt handling and response validation.

Author: pairing4good

.github/workflows/docs.yml

@@ -1,5 +1,3 @@
-import os
-
 from dotenv import load_dotenv
 
 from rlm import RLM

.github/workflows/publish.yml

@@ -19,7 +19,7 @@ def __init__(
         max_tokens: int = 32768,
         **kwargs,
     ):
-        super().__init__(model_name=model_name, **kwargs)
+        super().__init__(model_name=model_name or "claude-3-5-sonnet-20241022", **kwargs)
         self.client = anthropic.Anthropic(api_key=api_key)
         self.async_client = anthropic.AsyncAnthropic(api_key=api_key)
         self.model_name = model_name
@@ -31,7 +31,9 @@ def __init__(
         self.model_output_tokens: dict[str, int] = defaultdict(int)
         self.model_total_tokens: dict[str, int] = defaultdict(int)
 
-    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
+    def completion(
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]], model: str | None = None
+    ) -> str:
         messages, system = self._prepare_messages(prompt)
 
         model = model or self.model_name
@@ -47,7 +49,7 @@ def completion(self, prompt: str | list[dict[str, Any]], model: str | None = Non
         return response.content[0].text
 
     async def acompletion(
-        self, prompt: str | list[dict[str, Any]], model: str | None = None
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]], model: str | None = None
     ) -> str:
         messages, system = self._prepare_messages(prompt)
 
@@ -64,21 +66,18 @@ async def acompletion(
         return response.content[0].text
 
     def _prepare_messages(
-        self, prompt: str | list[dict[str, Any]]
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]]
     ) -> tuple[list[dict[str, Any]], str | None]:
         """Prepare messages and extract system prompt for Anthropic API."""
         system = None
 
         if isinstance(prompt, str):
             messages = [{"role": "user", "content": prompt}]
-        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
-            # Extract system message if present (Anthropic handles system separately)
-            messages = []
-            for msg in prompt:
-                if msg.get("role") == "system":
-                    system = msg.get("content")
-                else:
-                    messages.append(msg)
+        elif isinstance(prompt, dict):
+            # Single message dict - wrap in list
+            messages = [prompt]
+        elif isinstance(prompt, list):
+            messages = prompt
         else:
             raise ValueError(f"Invalid prompt type: {type(prompt)}")
 

examples/bedrock_example.py

@@ -28,7 +28,7 @@ def __init__(
         azure_deployment: str | None = None,
         **kwargs,
     ):
-        super().__init__(model_name=model_name, **kwargs)
+        super().__init__(model_name=model_name or "gpt-4o", **kwargs)
 
         if api_key is None:
             api_key = DEFAULT_AZURE_OPENAI_API_KEY
@@ -69,10 +69,14 @@ def __init__(
         self.model_output_tokens: dict[str, int] = defaultdict(int)
         self.model_total_tokens: dict[str, int] = defaultdict(int)
 
-    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
+    def completion(
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]], model: str | None = None
+    ) -> str:
         if isinstance(prompt, str):
             messages = [{"role": "user", "content": prompt}]
-        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+        elif isinstance(prompt, dict):
+            messages = [prompt]
+        elif isinstance(prompt, list):
             messages = prompt
         else:
             raise ValueError(f"Invalid prompt type: {type(prompt)}")
@@ -89,11 +93,13 @@ def completion(self, prompt: str | list[dict[str, Any]], model: str | None = Non
         return response.choices[0].message.content
 
     async def acompletion(
-        self, prompt: str | list[dict[str, Any]], model: str | None = None
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]], model: str | None = None
     ) -> str:
         if isinstance(prompt, str):
             messages = [{"role": "user", "content": prompt}]
-        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+        elif isinstance(prompt, dict):
+            messages = [prompt]
+        elif isinstance(prompt, list):
             messages = prompt
         else:
             raise ValueError(f"Invalid prompt type: {type(prompt)}")
@@ -109,7 +115,7 @@ async def acompletion(
         self._track_cost(response, model)
         return response.choices[0].message.content
 
-    def _track_cost(self, response: openai.ChatCompletion, model: str):
+    def _track_cost(self, response: Any, model: str):
         self.model_call_counts[model] += 1
 
         usage = getattr(response, "usage", None)

rlm/clients/anthropic.py

@@ -15,11 +15,11 @@ def __init__(self, model_name: str, **kwargs):
         self.kwargs = kwargs
 
     @abstractmethod
-    def completion(self, prompt: str | dict[str, Any]) -> str:
+    def completion(self, prompt: str | dict[str, Any] | list[dict[str, Any]]) -> str:
         raise NotImplementedError
 
     @abstractmethod
-    async def acompletion(self, prompt: str | dict[str, Any]) -> str:
+    async def acompletion(self, prompt: str | dict[str, Any] | list[dict[str, Any]]) -> str:
         raise NotImplementedError
 
     @abstractmethod

rlm/clients/azure_openai.py

@@ -79,7 +79,9 @@ def __init__(
         self.last_prompt_tokens = 0
         self.last_completion_tokens = 0
 
-    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
+    def completion(
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]], model: str | None = None
+    ) -> str:
         """Synchronous completion using Bedrock Converse API."""
         messages, system = self._prepare_messages(prompt)
 
@@ -112,7 +114,7 @@ def completion(self, prompt: str | list[dict[str, Any]], model: str | None = Non
         return response["output"]["message"]["content"][0]["text"]
 
     async def acompletion(
-        self, prompt: str | list[dict[str, Any]], model: str | None = None
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]], model: str | None = None
     ) -> str:
         """
         Async completion using Bedrock Converse API.
@@ -125,7 +127,7 @@ async def acompletion(
         return self.completion(prompt, model)
 
     def _prepare_messages(
-        self, prompt: str | list[dict[str, Any]]
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]]
     ) -> tuple[list[dict[str, Any]], str | None]:
         """
         Prepare messages for Bedrock Converse API.
@@ -138,7 +140,18 @@ def _prepare_messages(
         if isinstance(prompt, str):
             # Simple string prompt
             messages = [{"role": "user", "content": [{"text": prompt}]}]
-        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+        elif isinstance(prompt, dict):
+            # Single message dict
+            role = prompt.get("role")
+            content = prompt.get("content", "")
+            if role == "system":
+                system = content
+                messages = []
+            elif role in ("user", "assistant"):
+                messages = [{"role": role, "content": [{"text": content}]}]
+            else:
+                raise ValueError(f"Unsupported role: {role}")
+        elif isinstance(prompt, list):
             # Message list format
             messages = []
             for msg in prompt:

rlm/clients/base_lm.py

@@ -26,7 +26,7 @@ def __init__(
         model_name: str | None = "gemini-2.5-flash",
         **kwargs,
     ):
-        super().__init__(model_name=model_name, **kwargs)
+        super().__init__(model_name=model_name or "gemini-2.5-flash", **kwargs)
 
         if api_key is None:
             api_key = DEFAULT_GEMINI_API_KEY
@@ -49,7 +49,9 @@ def __init__(
         self.last_prompt_tokens = 0
         self.last_completion_tokens = 0
 
-    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
+    def completion(
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]], model: str | None = None
+    ) -> str:
         contents, system_instruction = self._prepare_contents(prompt)
 
         model = model or self.model_name
@@ -67,10 +69,10 @@ def completion(self, prompt: str | list[dict[str, Any]], model: str | None = Non
         )
 
         self._track_cost(response, model)
-        return response.text
+        return response.text or ""
 
     async def acompletion(
-        self, prompt: str | list[dict[str, Any]], model: str | None = None
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]], model: str | None = None
     ) -> str:
         contents, system_instruction = self._prepare_contents(prompt)
 
@@ -90,36 +92,49 @@ async def acompletion(
         )
 
         self._track_cost(response, model)
-        return response.text
+        return response.text or ""
 
     def _prepare_contents(
-        self, prompt: str | list[dict[str, Any]]
+        self, prompt: str | dict[str, Any] | list[dict[str, Any]]
     ) -> tuple[list[types.Content] | str, str | None]:
         """Prepare contents and extract system instruction for Gemini API."""
-        system_instruction = None
 
         if isinstance(prompt, str):
             return prompt, None
 
-        if isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
-            # Convert OpenAI-style messages to Gemini format
+        if isinstance(prompt, dict):
+            # Single message dict - wrap in list and process
+            role = prompt.get("role")
+            content = prompt.get("content", "")
+
+            if role == "system":
+                # Gemini handles system instruction separately
+                return [], content
+            elif role == "user":
+                return [types.Content(role="user", parts=[types.Part(text=content)])], None
+            elif role == "assistant":
+                # Gemini uses "model" instead of "assistant"
+                return [types.Content(role="model", parts=[types.Part(text=content)])], None
+            else:
+                # Default to user role for unknown roles
+                return [types.Content(role="user", parts=[types.Part(text=content)])], None
+
+        if isinstance(prompt, list):
+            # List of message dicts - process all
             contents = []
+            system_instruction = None
             for msg in prompt:
                 role = msg.get("role")
                 content = msg.get("content", "")
 
                 if role == "system":
-                    # Gemini handles system instruction separately
                     system_instruction = content
                 elif role == "user":
                     contents.append(types.Content(role="user", parts=[types.Part(text=content)]))
                 elif role == "assistant":
-                    # Gemini uses "model" instead of "assistant"
                     contents.append(types.Content(role="model", parts=[types.Part(text=content)]))
                 else:
-                    # Default to user role for unknown roles
                     contents.append(types.Content(role="user", parts=[types.Part(text=content)]))
-
             return contents, system_instruction
 
         raise ValueError(f"Invalid prompt type: {type(prompt)}")