Merge pull request #41 from RexWzh/rex

fix tests

.github/workflows/test.yml

@@ -20,8 +20,6 @@ jobs:
             os: ubuntu-latest
           - python-version: 3.9
             os: ubuntu-latest
-          - python-version: '3.10'
-            os: ubuntu-latest
 
     steps:
     - uses: actions/checkout@v2

openai_api_call/__init__.py

@@ -2,7 +2,7 @@
 
 __author__ = """Rex Wang"""
 __email__ = '[email protected]'
-__version__ = '1.0.0'
+__version__ = '1.0.1'
 
 import os, requests
 from .chattool import Chat, Resp

openai_api_call/asynctool.py

@@ -110,7 +110,7 @@ async def chat_complete(ind, locker, chatlog, chkpoint, **options):
                                  , chatlog=chatlog
                                  , chkpoint=chkpoint
                                  , **options)))
-        responses = await tqdm.gather(*tasks)
+        responses = await tqdm.gather(tasks)
         return responses
 
 def async_chat_completion( chatlogs:List[List[Dict]]

openai_api_call/response.py

@@ -2,6 +2,22 @@
 
 from typing import Dict
 
+# model cost($ per 1K tokens)\
+## https://openai.com/pricing
+## model | input | output
+model_cost_perktoken ={
+    "gpt-3.5-turbo": (0.0015, 0.002),
+    "gpt-3.5-turbo-0613" : (0.0015, 0.002),
+    "gpt-3.5-turbo-0301" : (0.0015, 0.002),
+    "gpt-3.5-turbo-16k-0613" : (0.003, 0.004),
+    "gpt-3.5-turbo-16k" : (0.003, 0.004),
+    "gpt-4": (0.03, 0.06),
+    "gpt-4-0613": (0.03, 0.06),
+    "gpt-4-0301": (0.03, 0.06),
+    "gpt-4-32k-0613": (0.06, 0.12),
+    "gpt-4-32k": (0.06, 0.12),
+}
+
 class Resp():
 
     def __init__(self, response:Dict) -> None:
@@ -11,6 +27,10 @@ def is_valid(self):
         """Check if the response is an error"""
         return 'error' not in self.response
 
+    def cost(self):
+        """Calculate the cost of the response"""
+        return response_cost(self.model, self.prompt_tokens, self.completion_tokens)
+
     def __repr__(self) -> str:
         return f"`Resp`: {self.content}"
 
@@ -91,4 +111,19 @@ def error_code(self):
     @property
     def finish_reason(self):
         """Finish reason"""
-        return self.response['choices'][0]['finish_reason']
+        return self.response['choices'][0]['finish_reason']
+
+def response_cost(model:str, prompt_tokens:int, completion_tokens:int):
+    """Calculate the cost of the response
+
+    Args:
+        model (str): model name
+        prompt_tokens (int): number of tokens in the prompt
+        completion_tokens (int): number of tokens of the response
+
+    Returns:
+        float: cost of the response
+    """
+    assert model in model_cost_perktoken, f"Model {model} is not known!"
+    input_price, output_price = model_cost_perktoken[model]
+    return (input_price * prompt_tokens + output_price * completion_tokens) / 1000

setup.py

@@ -7,7 +7,7 @@
 with open('README.md') as readme_file:
     readme = readme_file.read()
 
-VERSION = '1.0.0'
+VERSION = '1.0.1'
 
 requirements = ['Click>=7.0', 'requests>=2.20', 'tqdm>=4.60', 'docstring_parser>=0.10', 'aiohttp>=3.8']
 test_requirements = ['pytest>=3', 'unittest']

tests/test_async.py

@@ -26,4 +26,4 @@ def data2chat(data):
         return chat
     t = time.time()
     process_chats(chatlogs, data2chat, chkpoint, clearfile=True)
-    print(f"Time elapsed: {time.time() - t:.2f}s")
+    print(f"Time elapsed: {time.time() - t:.2f}s")

tests/test_openai_api_call.py

@@ -194,6 +194,7 @@ def test_usage():
     assert resp.total_tokens == 18
     assert resp.prompt_tokens == 8
     assert resp.completion_tokens == 10
+    print(resp.cost())
 
 def test_content():
     resp = Resp(response=response)