From 9754499e814b88b70b309a46f7dfe4aa328bf8a3 Mon Sep 17 00:00:00 2001
From: Lu Yao Chen <luyao.chen@agilebits.com>
Date: Sun, 19 Mar 2023 17:17:18 -0700
Subject: [PATCH] simple send req non chatgpt test

---
 .gitignore             |  3 ++-
 models/openai_model.py | 16 +++++++++-------
 tests/__init__.py      |  0
 tests/test_requests.py | 21 +++++++++++++++++++++
 4 files changed, 32 insertions(+), 8 deletions(-)
 create mode 100644 tests/__init__.py
 create mode 100644 tests/test_requests.py

diff --git a/.gitignore b/.gitignore
index 5d320ce..6f7812f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,4 +11,5 @@ usage.txt
 /dalleimages
 /indexes
 /audiotemp
-/pickles
\ No newline at end of file
+/pickles
+.idea
\ No newline at end of file
diff --git a/models/openai_model.py b/models/openai_model.py
index f31a11b..b8cfced 100644
--- a/models/openai_model.py
+++ b/models/openai_model.py
@@ -633,10 +633,14 @@ class Model:
             else:
                 await self.usage_service.update_usage(tokens_used)
         except Exception as e:
-            raise ValueError(
-                "The API returned an invalid response: "
-                + str(response["error"]["message"])
-            ) from e
+            traceback.print_exc()
+            if 'error' in response:
+                raise ValueError(
+                    "The API returned an invalid response: "
+                    + str(response["error"]["message"])
+                ) from e
+            else:
+                raise ValueError("The API returned an invalid response") from e
 
     @backoff.on_exception(
         backoff.expo,
@@ -1010,8 +1014,6 @@ class Model:
         stop=None,
         custom_api_key=None,
         is_chatgpt_request=False,
-    ) -> (
-        Tuple[dict, bool]
     ):  # The response, and a boolean indicating whether or not the context limit was reached.
         # Validate that  all the parameters are in a good state before we send the request
 
@@ -1057,7 +1059,7 @@ class Model:
                     headers=headers,
                 ) as resp:
                     response = await resp.json()
-                    # print(f"Payload -> {payload}")
+                    print(f"Payload -> {payload}")
                     # Parse the total tokens used for this request and response pair from the response
                     await self.valid_text_request(
                         response, model=self.model if model is None else model
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_requests.py b/tests/test_requests.py
new file mode 100644
index 0000000..6226bba
--- /dev/null
+++ b/tests/test_requests.py
@@ -0,0 +1,21 @@
+from pathlib import Path
+
+import pytest
+from models.openai_model import Models, Model
+from transformers import GPT2TokenizerFast
+import asyncio
+
+from services.usage_service import UsageService
+import os
+
+# Non-ChatGPT
+@pytest.mark.asyncio
+async def test_send_req():
+
+    usage_service = UsageService(Path("../tests"))
+    model = Model(usage_service)
+    prompt = 'how many hours are in a day?'
+    tokens = len(GPT2TokenizerFast.from_pretrained("gpt2")(prompt)["input_ids"])
+    # tokens = 60
+    res = await model.send_request(prompt, tokens)
+    assert '24' in res['choices'][0]['text']