Fixing openrouter pricing rate limit

.github/workflows/unit_tests.yml

@@ -38,6 +38,9 @@ jobs:
       - name: Install Playwright
         run: playwright install chromium --with-deps
 
+      - name: Download WebArena / VisualWebArena ressource files
+        run: python -c 'import nltk; nltk.download("punkt_tab")'
+
       - name: Fetch MiniWob
         uses: actions/checkout@v4
         with:
@@ -58,4 +61,4 @@ jobs:
       - name: Run AgentLab Unit Tests
         env:
           MINIWOB_URL: "http://localhost:8080/miniwob/"
-        run: pytest -n 5 --durations=10 -m 'not pricy' -v tests/
+        run: pytest -n 5 --durations=10 -m 'not pricy' -v tests/

src/agentlab/llm/chat_api.py

@@ -208,6 +208,10 @@ def handle_error(error, itr, min_retry_wait_time, max_retry):
     return error_type
 
 
+class OpenRouterError(openai.OpenAIError):
+    pass
+
+
 class ChatModel(AbstractChatModel):
     def __init__(
         self,
@@ -274,6 +278,12 @@ def __call__(self, messages: list[dict]) -> dict:
                     temperature=self.temperature,
                     max_tokens=self.max_tokens,
                 )
+
+                if completion.usage is None:
+                    raise OpenRouterError(
+                        "The completion object does not contain usage information. This is likely a bug in the OpenRouter API."
+                    )
+
                 self.success = True
                 break
             except openai.OpenAIError as e:

src/agentlab/llm/tracking.py

@@ -1,3 +1,4 @@
+from functools import cache
 import os
 import threading
 from contextlib import contextmanager
@@ -61,6 +62,7 @@ def wrapper(self, obs):
     return wrapper
 
 
+@cache
 def get_pricing_openrouter():
     api_key = os.getenv("OPENROUTER_API_KEY")
     assert api_key, "OpenRouter API key is required"