From 2be6cfdd64002c59a9bed3a312adb765982a0a04 Mon Sep 17 00:00:00 2001
From: Liu Jun <liujun45@baidu.com>
Date: Wed, 10 Jan 2024 14:42:12 +0800
Subject: [PATCH] feat: optimize client traceback and input ak/sk (#182)

* optimize client traceback

* fix lint

* not showing message when accesskey is set
---
 docs/cli.md                       |  1 +
 src/qianfan/common/client/main.py | 56 +++++++++++++++++++++++++++++--
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/docs/cli.md b/docs/cli.md
index 8872f911..8ef62ccc 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -21,6 +21,7 @@ $ qianfan [OPTIONS] COMMAND [ARGS]...
 * `--secret-key TEXT`：百度智能云安全认证 Secret Key，获取方式参考 [文档](https://cloud.baidu.com/doc/Reference/s/9jwvz2egb)。
 * `--ak TEXT` [过时]：千帆平台应用的 API Key，仅能用于模型推理部分 API，获取方式参考 [文档](https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Slkkydake)。
 * `--sk TEXT` [过时]：千帆平台应用的 Secret Key，仅能用于模型推理部分 API，获取方式参考 [文档](https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Slkkydake)。
+* `--enable-traceback`：打印完整的错误堆栈信息，仅在发生异常时有效。
 * `--version -v`：打印版本信息。
 * `--install-completion`：为当前 shell 安装自动补全脚本。
 * `--show-completion`：展示自动补全脚本。
diff --git a/src/qianfan/common/client/main.py b/src/qianfan/common/client/main.py
index f28b0ce5..267f1ced 100644
--- a/src/qianfan/common/client/main.py
+++ b/src/qianfan/common/client/main.py
@@ -18,6 +18,7 @@
 from typing import Optional
 
 import typer
+from rich.prompt import Prompt
 
 import qianfan
 from qianfan.common.client.chat import chat_entry
@@ -25,6 +26,7 @@
 from qianfan.common.client.dataset import dataset_app
 from qianfan.common.client.embedding import embedding_entry
 from qianfan.common.client.txt2img import txt2img_entry
+from qianfan.common.client.utils import print_error_msg, print_info_msg
 
 app = typer.Typer(
     no_args_is_help=True,
@@ -37,6 +39,8 @@
 app.command(name="embedding", no_args_is_help=True)(embedding_entry)
 app.add_typer(dataset_app, name="dataset")
 
+_enable_traceback = False
+
 
 def version_callback(value: bool) -> None:
     """
@@ -65,7 +69,13 @@ def main() -> None:
     """
     Main function of qianfan client.
     """
-    app()
+    try:
+        app()
+    except Exception as e:
+        if _enable_traceback:
+            raise
+        else:
+            print_error_msg(str(e))
 
 
 @app.callback()
@@ -122,11 +132,53 @@ def entry(
         is_eager=True,
         help="Print version.",
     ),
+    enable_traceback: bool = typer.Option(
+        False, "--enable-traceback", help="Print traceback when exception is thrown."
+    ),
 ) -> None:
     """
     Qianfan CLI which provides access to various Qianfan services.
     """
-    pass
+    global _enable_traceback
+    _enable_traceback = enable_traceback
+
+    ak = qianfan.get_config().AK
+    sk = qianfan.get_config().SK
+    access_key = qianfan.get_config().ACCESS_KEY
+    secret_key = qianfan.get_config().SECRET_KEY
+
+    if ak is None or sk is None:
+        if access_key is None or secret_key is None:
+            print_info_msg(
+                'No enough credential found. Please provide your "access key" and'
+                ' "secret key".'
+            )
+            print_info_msg(
+                "You can find your key at"
+                " https://console.bce.baidu.com/iam/#/iam/accesslist"
+            )
+            print_info_msg(
+                "You can also set the credential using environment variable"
+                ' "QIANFAN_ACCESS_KEY" and "QIANFAN_SECRET_KEY".'
+            )
+            print()
+            if access_key is None:
+                while True:
+                    access_key = Prompt.ask("Please input your [b i]Access Key[/b i]")
+                    if len(access_key) != 0:
+                        qianfan.get_config().ACCESS_KEY = access_key
+                        break
+                    else:
+                        print_error_msg("Access key cannot be empty.")
+            if secret_key is None:
+                while True:
+                    secret_key = Prompt.ask("Please input your [b i]Secret Key[/b i]")
+                    if len(secret_key) != 0:
+                        qianfan.get_config().SECRET_KEY = secret_key
+                        break
+                    else:
+                        print_error_msg("Secret key cannot be empty.")
+            print()
 
 
 if __name__ == "__main__":