From 20d7deaccacd34988a4c5fd9e52d962e6d86792e Mon Sep 17 00:00:00 2001
From: StanChan03 <StanleyChan529@gmail.com>
Date: Thu, 26 Dec 2024 19:51:34 -0800
Subject: [PATCH] fix test

---
 .github/tests/lm_tests.py | 55 +++++++++++++++++++++++++++------------
 1 file changed, 38 insertions(+), 17 deletions(-)

diff --git a/.github/tests/lm_tests.py b/.github/tests/lm_tests.py
index 91fc43c..ce15fdd 100644
--- a/.github/tests/lm_tests.py
+++ b/.github/tests/lm_tests.py
@@ -483,15 +483,25 @@ def test_operator_cache(setup_models, model):
     lm = LM(model="gpt-4o-mini", cache=cache)
     lotus.settings.configure(lm=lm, enable_message_cache=True, enable_operator_cache=True)
 
-    batch = [
-        [{"role": "user", "content": "Hello, world!"}],
-        [{"role": "user", "content": "What is the capital of France?"}],
-    ]
-    first_responses = lm(batch).outputs
+    data = {
+        "Course Name": [
+            "Dynamics and Control of Chemical Processes",
+            "Optimization Methods in Engineering",
+            "Chemical Kinetics and Catalysis",
+            "Transport Phenomena and Separations",
+        ]
+    }
+
+    df = pd.DataFrame(data)
+    user_instruction = "What is a similar course to {Course Name}. Be concise?"
+
+    first_response = df.sem_map(df, user_instruction)
     assert lm.stats.total_usage.operator_cache_hits == 0
-    second_responses = lm(batch).outputs
-    assert lm.stats.total_usage.operator_cache_hits == 2
-    assert first_responses == second_responses
+
+    second_response = df.sem_map(df, user_instruction)
+    assert lm.stats.total_usage.operator_cache_hits == 1
+
+    assert first_response == second_response
 
 
 @pytest.mark.parametrize("model", get_enabled("gpt-4o-mini"))
@@ -502,19 +512,30 @@ def test_disable_operator_cache(setup_models, model):
     lm = LM(model="gpt-4o-mini", cache=cache)
     lotus.settings.configure(lm=lm, enable_message_cache=True, enable_operator_cache=False)
 
-    batch = [
-        [{"role": "user", "content": "Hello, world!"}],
-        [{"role": "user", "content": "What is the capital of France?"}],
-    ]
-    lm(batch)
+    data = {
+        "Course Name": [
+            "Dynamics and Control of Chemical Processes",
+            "Optimization Methods in Engineering",
+            "Chemical Kinetics and Catalysis",
+            "Transport Phenomena and Separations",
+        ]
+    }
+
+    df = pd.DataFrame(data)
+    user_instruction = "What is a similar course to {Course Name}. Be concise?"
+
+    first_response = df.sem_map(df, user_instruction)
     assert lm.stats.total_usage.operator_cache_hits == 0
-    lm(batch)
+
+    second_response = df.sem_map(df, user_instruction)
     assert lm.stats.total_usage.operator_cache_hits == 0
 
+    assert first_response == second_response
+
     # Now enable operator cache. Note that the first batch is not cached.
     lotus.settings.configure(enable_operator_cache=True)
-    first_responses = lm(batch).outputs
+    first_responses = df.sem_map(df, user_instruction)
     assert lm.stats.total_usage.operator_cache_hits == 0
-    second_responses = lm(batch).outputs
-    assert lm.stats.total_usage.operator_cache_hits == 2
+    second_responses = df.sem_map(df, user_instruction)
+    assert lm.stats.total_usage.operator_cache_hits == 1
     assert first_responses == second_responses