From 872f3f9f8555ca0811d464a957b95aecf25c55ba Mon Sep 17 00:00:00 2001
From: Jintao Huang <huangjintao.hjt@alibaba-inc.com>
Date: Tue, 29 Oct 2024 14:31:28 +0800
Subject: [PATCH] update client

---
 swift/llm/__init__.py                    | 4 ++--
 swift/llm/infer/__init__.py              | 6 ++++--
 swift/llm/infer/infer_engine/__init__.py | 4 +++-
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/swift/llm/__init__.py b/swift/llm/__init__.py
index ab56f7db2..acbcfb6aa 100644
--- a/swift/llm/__init__.py
+++ b/swift/llm/__init__.py
@@ -6,7 +6,7 @@
 if TYPE_CHECKING:
     # Recommend using `xxx_main`
     from .infer import (VllmEngine, InferRequest, RequestConfig, InferStats, LmdeployEngine, PtEngine, infer_main,
-                        deploy_main, PtLoRARequest)
+                        deploy_main, PtLoRARequest, InferClient)
     from .export import export_main, merge_lora
     from .eval import eval_main
     from .train import sft_main, pt_main, rlhf_main
@@ -30,7 +30,7 @@
         'rlhf': ['rlhf_main'],
         'infer': [
             'deploy_main', 'VllmEngine', 'InferRequest', 'RequestConfig', 'InferStats', 'LmdeployEngine', 'PtEngine',
-            'infer_main', 'PtLoRARequest'
+            'infer_main', 'PtLoRARequest', 'InferClient'
         ],
         'export': ['export_main', 'merge_lora'],
         'eval': ['eval_main'],
diff --git a/swift/llm/infer/__init__.py b/swift/llm/infer/__init__.py
index 0aa45cdd6..c3e3e27eb 100644
--- a/swift/llm/infer/__init__.py
+++ b/swift/llm/infer/__init__.py
@@ -7,14 +7,16 @@
     from .infer import infer_main
     from .deploy import deploy_main
     from .protocol import InferRequest, RequestConfig
-    from .infer_engine import (InferEngine, VllmEngine, LmdeployEngine, PtEngine, InferStats, PtLoRARequest)
+    from .infer_engine import (InferEngine, VllmEngine, LmdeployEngine, PtEngine, InferStats, PtLoRARequest,
+                               InferClient)
 else:
     _extra_objects = {k: v for k, v in globals().items() if not k.startswith('_')}
     _import_structure = {
         'deploy': ['deploy_main'],
         'infer': ['infer_main'],
         'protocol': ['InferRequest', 'RequestConfig'],
-        'infer_engine': ['InferEngine', 'VllmEngine', 'LmdeployEngine', 'PtEngine', 'InferStats', 'PtLoRARequest'],
+        'infer_engine': ['InferEngine', 'VllmEngine', 'LmdeployEngine', 'PtEngine', 'InferStats', 'PtLoRARequest',
+                         'InferClient'],
     }
 
     import sys
diff --git a/swift/llm/infer/infer_engine/__init__.py b/swift/llm/infer/infer_engine/__init__.py
index 916de19cd..958aa3374 100644
--- a/swift/llm/infer/infer_engine/__init__.py
+++ b/swift/llm/infer/infer_engine/__init__.py
@@ -7,6 +7,7 @@
     from .vllm_engine import VllmEngine
     from .lmdeploy_engine import LmdeployEngine
     from .pt_engine import PtEngine, PtLoRARequest
+    from .infer_client import InferClient
     from .infer_engine import InferEngine
     from .base import BaseInferEngine
     from .utils import InferStats
@@ -16,9 +17,10 @@
         'vllm_engine': ['VllmEngine'],
         'lmdeploy_engine': ['LmdeployEngine'],
         'pt_engine': ['PtEngine', 'PtLoRARequest'],
+        'infer_client': ['InferClient'],
         'infer_engine': ['InferEngine'],
         'base': ['BaseInferEngine'],
-        'utils': ['InferStats']
+        'utils': ['InferStats'],
     }
 
     import sys