Update cache if cost not calculated (microsoft#903)

* update cost for cache * fix test
rlam3 · Dec 19, 2023 · 8eba4e8 · 8eba4e8
1 parent f89aa98
commit 8eba4e8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/autogen/oai/client.py b/autogen/oai/client.py
@@ -235,9 +235,15 @@ def yes_or_no_filter(context, response):
                     # Try to get the response from cache
                     key = get_key(params)
                     response = cache.get(key, None)
+
                     if response is not None:
+                        try:
+                            response.cost
+                        except AttributeError:
+                            # update atrribute if cost is not calculated
+                            response.cost = self.cost(response)
+                            cache.set(key, response)
                         self._update_usage_summary(response, use_cache=True)
-                    if response is not None:
                         # check the filter
                         pass_filter = filter_func is None or filter_func(context=context, response=response)
                         if pass_filter or i == last:

diff --git a/test/oai/test_client.py b/test/oai/test_client.py
@@ -80,7 +80,7 @@ def test_usage_summary():
     # check update
     client._update_usage_summary(response, use_cache=True)
     assert (
-        client.actual_usage_summary["total_cost"] == response.cost * 2
+        client.total_usage_summary["total_cost"] == response.cost * 2
     ), "total_cost should be equal to response.cost * 2"
 
     # check clear