fix: error code in connection_error, return token_usage in last frame…

… during using stream/astream
iflytek · May 13, 2024 · beaf8f6 · beaf8f6
1 parent 8aefe32
commit beaf8f6
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/sparkai/errors/__init__.py b/sparkai/errors/__init__.py
@@ -61,4 +61,5 @@ class SparkAIClientConfigurationError(SparkAIClientError):
 class SparkAIConnectionError(ConnectionError):
     def __init__(self, error_code, message):
         self.error_code = error_code
+        self.message = message
         super().__init__(message)
diff --git a/tests/examples/llm_test.py b/tests/examples/llm_test.py
@@ -419,15 +419,15 @@ def test_error():
     try:
         error_func()
     except ConnectionError as e:
-        print(e.error_code)
+        print(e.error_code, str(e))
 
 if __name__ == '__main__':
     # import asyncio
     # test_26b()
     # asyncio.run(test_26b())
     # test_once()
     #test_stream()
-    #test_error()
+    test_error()
     # test_function_call_once()
     # test_function_call_stream()
     # test_image()