From d4e663f7bc60135749da0381a245420d4d49c991 Mon Sep 17 00:00:00 2001
From: Joshua Vial <joshua@enspiral.com>
Date: Wed, 20 Dec 2023 10:13:28 +1300
Subject: [PATCH] benchmark work with openrouter

---
 benchmark/benchmark.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/benchmark/benchmark.py b/benchmark/benchmark.py
index 1a22c0f62..e61cf038a 100755
--- a/benchmark/benchmark.py
+++ b/benchmark/benchmark.py
@@ -782,7 +782,22 @@ def run_test(
         chat_history_file=history_fname,
     )
 
-    main_model = models.Model.create(model_name)
+    # Check if AIDER_DOCKER environment variable is set and create client accordingly
+    if "OPENAI_API_BASE" in os.environ and "openrouter.ai" in os.environ["OPENAI_API_BASE"]:
+        client = openai.OpenAI(
+            api_key=os.environ["OPENAI_API_KEY"],
+            base_url=os.environ.get("OPENAI_API_BASE", "https://api.openai.com"),
+            default_headers={
+                "HTTP-Referer": "http://aider.chat",
+                "X-Title": "Aider",
+            }
+        )
+    else:
+        client = openai.OpenAI(
+            api_key=os.environ["OPENAI_API_KEY"],
+        )
+
+    main_model = models.Model.create(model_name, client)
     edit_format = edit_format or main_model.edit_format
 
     dump(main_model)
@@ -790,8 +805,6 @@ def run_test(
     show_fnames = ",".join(map(str, fnames))
     print("fnames:", show_fnames)
 
-    client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
-
     coder = Coder.create(
         main_model,
         edit_format,