update for ruff checks

Author: alexzhang13

examples/lm_in_repl.py

@@ -6,12 +6,13 @@
 """
 
 import os
+
+from dotenv import load_dotenv
+
 from rlm.clients.portkey import PortkeyClient
 from rlm.core.lm_handler import LMHandler
 from rlm.environments.local_repl import LocalREPL
 
-from dotenv import load_dotenv
-
 load_dotenv()
 
 setup_code = """

examples/modal_repl_example.py

@@ -4,10 +4,10 @@
 Run with: python -m examples.modal_repl_example
 """
 
-from rlm.environments.modal_repl import ModalREPL
 from rlm.clients.base_lm import BaseLM
 from rlm.core.lm_handler import LMHandler
-from rlm.core.types import UsageSummary, ModelUsageSummary
+from rlm.core.types import ModelUsageSummary, UsageSummary
+from rlm.environments.modal_repl import ModalREPL
 
 
 class MockLM(BaseLM):

examples/rlm_example.py

@@ -1,9 +1,10 @@
-from rlm.logger import RLMLogger
-from rlm import RLM
-
 import os
+
 from dotenv import load_dotenv
 
+from rlm import RLM
+from rlm.logger import RLMLogger
+
 load_dotenv()
 
 logger = RLMLogger(log_dir="./logs")
@@ -19,7 +20,5 @@
     logger=logger,
 )
 
-result = rlm.completion(
-    "Using your code, solve 2^(2^(2^(2))). Show your work in Python."
-)
+result = rlm.completion("Using your code, solve 2^(2^(2^(2))). Show your work in Python.")
 print(result)

rlm/clients/__init__.py

@@ -1,15 +1,16 @@
-from rlm.clients.base_lm import BaseLM
-from rlm.core.types import ClientBackend
+from typing import Any
 
 from dotenv import load_dotenv
-from typing import Dict, Any
+
+from rlm.clients.base_lm import BaseLM
+from rlm.core.types import ClientBackend
 
 load_dotenv()
 
 
 def get_client(
     backend: ClientBackend,
-    backend_kwargs: Dict[str, Any],
+    backend_kwargs: dict[str, Any],
 ) -> BaseLM:
     """
     Routes a specific backend and the args (as a dict) to the appropriate client if supported.

rlm/clients/anthropic.py

@@ -1,9 +1,11 @@
-from rlm.clients.base_lm import BaseLM
-from rlm.core.types import UsageSummary, ModelUsageSummary
-from typing import Dict, Any, Optional, List
 from collections import defaultdict
+from typing import Any
+
 import anthropic
 
+from rlm.clients.base_lm import BaseLM
+from rlm.core.types import ModelUsageSummary, UsageSummary
+
 
 class AnthropicClient(BaseLM):
     """
@@ -13,7 +15,7 @@ class AnthropicClient(BaseLM):
     def __init__(
         self,
         api_key: str,
-        model_name: Optional[str] = None,
+        model_name: str | None = None,
         max_tokens: int = 32768,
         **kwargs,
     ):
@@ -24,14 +26,12 @@ def __init__(
         self.max_tokens = max_tokens
 
         # Per-model usage tracking
-        self.model_call_counts: Dict[str, int] = defaultdict(int)
-        self.model_input_tokens: Dict[str, int] = defaultdict(int)
-        self.model_output_tokens: Dict[str, int] = defaultdict(int)
-        self.model_total_tokens: Dict[str, int] = defaultdict(int)
+        self.model_call_counts: dict[str, int] = defaultdict(int)
+        self.model_input_tokens: dict[str, int] = defaultdict(int)
+        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: Optional[str] = None
-    ) -> str:
+    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
         messages, system = self._prepare_messages(prompt)
 
         model = model or self.model_name
@@ -47,7 +47,7 @@ def completion(
         return response.content[0].text
 
     async def acompletion(
-        self, prompt: str | List[Dict[str, Any]], model: Optional[str] = None
+        self, prompt: str | list[dict[str, Any]], model: str | None = None
     ) -> str:
         messages, system = self._prepare_messages(prompt)
 
@@ -64,16 +64,14 @@ async def acompletion(
         return response.content[0].text
 
     def _prepare_messages(
-        self, prompt: str | List[Dict[str, Any]]
-    ) -> tuple[List[Dict[str, Any]], Optional[str]]:
+        self, prompt: str | 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
-        ):
+        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:
@@ -90,9 +88,7 @@ def _track_cost(self, response: anthropic.types.Message, model: str):
         self.model_call_counts[model] += 1
         self.model_input_tokens[model] += response.usage.input_tokens
         self.model_output_tokens[model] += response.usage.output_tokens
-        self.model_total_tokens[model] += (
-            response.usage.input_tokens + response.usage.output_tokens
-        )
+        self.model_total_tokens[model] += response.usage.input_tokens + response.usage.output_tokens
 
         # Track last call for handler to read
         self.last_prompt_tokens = response.usage.input_tokens

rlm/clients/base_lm.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Any
+from typing import Any
 
 from rlm.core.types import UsageSummary
 
@@ -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]) -> str:
         raise NotImplementedError
 
     @abstractmethod
-    async def acompletion(self, prompt: str | Dict[str, Any]) -> str:
+    async def acompletion(self, prompt: str | dict[str, Any]) -> str:
         raise NotImplementedError
 
     @abstractmethod

rlm/clients/litellm.py

@@ -1,9 +1,11 @@
-from rlm.clients.base_lm import BaseLM
-from rlm.core.types import UsageSummary, ModelUsageSummary
-from typing import Dict, Any, Optional, List
 from collections import defaultdict
+from typing import Any
+
 import litellm
 
+from rlm.clients.base_lm import BaseLM
+from rlm.core.types import ModelUsageSummary, UsageSummary
+
 
 class LiteLLMClient(BaseLM):
     """
@@ -13,9 +15,9 @@ class LiteLLMClient(BaseLM):
 
     def __init__(
         self,
-        model_name: Optional[str] = None,
-        api_key: Optional[str] = None,
-        api_base: Optional[str] = None,
+        model_name: str | None = None,
+        api_key: str | None = None,
+        api_base: str | None = None,
         **kwargs,
     ):
         super().__init__(model_name=model_name, **kwargs)
@@ -24,19 +26,15 @@ def __init__(
         self.api_base = api_base
 
         # Per-model usage tracking
-        self.model_call_counts: Dict[str, int] = defaultdict(int)
-        self.model_input_tokens: Dict[str, int] = defaultdict(int)
-        self.model_output_tokens: Dict[str, int] = defaultdict(int)
-        self.model_total_tokens: Dict[str, int] = defaultdict(int)
+        self.model_call_counts: dict[str, int] = defaultdict(int)
+        self.model_input_tokens: dict[str, int] = defaultdict(int)
+        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: Optional[str] = None
-    ) -> str:
+    def completion(self, prompt: str | 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, list) and all(isinstance(item, dict) for item in prompt):
             messages = prompt
         else:
             raise ValueError(f"Invalid prompt type: {type(prompt)}")
@@ -56,13 +54,11 @@ def completion(
         return response.choices[0].message.content
 
     async def acompletion(
-        self, prompt: str | List[Dict[str, Any]], model: Optional[str] = None
+        self, prompt: str | 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, list) and all(isinstance(item, dict) for item in prompt):
             messages = prompt
         else:
             raise ValueError(f"Invalid prompt type: {type(prompt)}")

rlm/clients/openai.py

@@ -1,11 +1,12 @@
-from rlm.clients.base_lm import BaseLM, UsageSummary, ModelUsageSummary
-from typing import Dict, Any, Optional, List
-from collections import defaultdict
-import openai
 import os
+from collections import defaultdict
+from typing import Any
 
+import openai
 from dotenv import load_dotenv
 
+from rlm.clients.base_lm import BaseLM, ModelUsageSummary, UsageSummary
+
 load_dotenv()
 
 # Load API keys from environment variables
@@ -20,9 +21,9 @@ class OpenAIClient(BaseLM):
 
     def __init__(
         self,
-        api_key: Optional[str] = None,
-        model_name: Optional[str] = None,
-        base_url: Optional[str] = None,
+        api_key: str | None = None,
+        model_name: str | None = None,
+        base_url: str | None = None,
         **kwargs,
     ):
         super().__init__(model_name=model_name, **kwargs)
@@ -38,19 +39,15 @@ def __init__(
         self.model_name = model_name
 
         # Per-model usage tracking
-        self.model_call_counts: Dict[str, int] = defaultdict(int)
-        self.model_input_tokens: Dict[str, int] = defaultdict(int)
-        self.model_output_tokens: Dict[str, int] = defaultdict(int)
-        self.model_total_tokens: Dict[str, int] = defaultdict(int)
+        self.model_call_counts: dict[str, int] = defaultdict(int)
+        self.model_input_tokens: dict[str, int] = defaultdict(int)
+        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: Optional[str] = None
-    ) -> str:
+    def completion(self, prompt: str | 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, list) and all(isinstance(item, dict) for item in prompt):
             messages = prompt
         else:
             raise ValueError(f"Invalid prompt type: {type(prompt)}")
@@ -64,13 +61,11 @@ def completion(
         return response.choices[0].message.content
 
     async def acompletion(
-        self, prompt: str | List[Dict[str, Any]], model: Optional[str] = None
+        self, prompt: str | 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, list) and all(isinstance(item, dict) for item in prompt):
             messages = prompt
         else:
             raise ValueError(f"Invalid prompt type: {type(prompt)}")
@@ -79,9 +74,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.client.chat.completions.create(model=model, messages=messages)
         self._track_cost(response, model)
         return response.choices[0].message.content