update (#3175)

Co-authored-by: HRUSHIKESH DOKALA <[email protected]>
microsoft · Jul 30, 2024 · 4c9341a · 4c9341a
1 parent 574c8c5
commit 4c9341a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/autogen/oai/client.py b/autogen/oai/client.py
@@ -795,7 +795,7 @@ def _cost_with_customized_price(
         n_output_tokens = response.usage.completion_tokens if response.usage is not None else 0  # type: ignore [union-attr]
         if n_output_tokens is None:
             n_output_tokens = 0
-        return n_input_tokens * price_1k[0] + n_output_tokens * price_1k[1]
+        return (n_input_tokens * price_1k[0] + n_output_tokens * price_1k[1]) / 1000
 
     @staticmethod
     def _update_dict_from_chunk(chunk: BaseModel, d: Dict[str, Any], field: str) -> int:

diff --git a/test/oai/test_client.py b/test/oai/test_client.py
@@ -140,10 +140,10 @@ def test_customized_cost():
         env_or_file=OAI_CONFIG_LIST, file_location=KEY_LOC, filter_dict={"tags": ["gpt-3.5-turbo-instruct"]}
     )
     for config in config_list:
-        config.update({"price": [1, 1]})
+        config.update({"price": [1000, 1000]})
     client = OpenAIWrapper(config_list=config_list, cache_seed=None)
     response = client.create(prompt="1+3=")
-    assert response.cost >= 4, "Due to customized pricing, cost should be greater than 4"
+    assert response.cost >= 4 and response.cost < 10, "Due to customized pricing, cost should be > 4 and < 10"
 
 
 @pytest.mark.skipif(skip, reason="openai>=1 not installed")