Eb cherry (#322)

* resolve conflicts

* remove unused comments

* resolve conflicts

* reformat

* resolve conflicts

* remove lines

* resolve conflicts

* Add ernieb speed

* resolve conflicts

* resolve conflicts

* Update

* restore handle response

* Update format

* Update http_client

* Update format

Author: w5688414

erniebot/src/erniebot/backends/bce.py

@@ -351,7 +351,6 @@ def handle_response(self, resp: EBResponse) -> EBResponse:
         if "error_code" in resp and "error_msg" in resp:
             ecode = resp["error_code"]
             emsg = resp["error_msg"]
-            print(ecode)
             if ecode in (4, 17):
                 raise errors.RequestLimitError(emsg, ecode=ecode)
             elif ecode in (13, 15, 18):

erniebot/src/erniebot/backends/custom.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from typing import Any, AsyncIterator, ClassVar, Dict, Iterator, Optional, Union
 
 import erniebot.errors as errors
+import erniebot.utils.logging as logging
 from erniebot.api_types import APIType
 from erniebot.response import EBResponse
 from erniebot.types import FilesType, HeadersType, ParamsType
@@ -31,6 +33,10 @@ class CustomBackend(EBBackend):
 
     def __init__(self, config_dict: Dict[str, Any]) -> None:
         super().__init__(config_dict=config_dict)
+        access_token = self._cfg.get("access_token", None)
+        if access_token is None:
+            access_token = os.environ.get("AISTUDIO_ACCESS_TOKEN", None)
+        self._access_token = access_token
 
     def request(
         self,
@@ -79,7 +85,8 @@ async def arequest(
             params=params,
             files=files,
         )
-
+        if self._access_token is not None:
+            headers = self._add_aistudio_fields_to_headers(headers)
         return await self._client.asend_request(
             method,
             url,
@@ -110,3 +117,12 @@ def handle_response(self, resp: EBResponse) -> EBResponse:
                 raise errors.APIError(emsg, ecode=ecode)
         else:
             return resp
+
+    def _add_aistudio_fields_to_headers(self, headers: HeadersType) -> HeadersType:
+        if "Authorization" in headers:
+            logging.warning(
+                "Key 'Authorization' already exists in `headers`: %r",
+                headers["Authorization"],
+            )
+        headers["Authorization"] = f"{self._access_token}"
+        return headers

erniebot/src/erniebot/resources/chat_completion.py

@@ -92,6 +92,15 @@ class ChatCompletion(EBResource, CreatableWithStreaming):
                 "ernie-3.5": {
                     "model_id": "completions",
                 },
+                "ernie-4.0": {
+                    "model_id": "completions_pro",
+                },
+                "ernie-longtext": {
+                    "model_id": "ernie_bot_8k",
+                },
+                "ernie-speed": {
+                    "model_id": "ernie_speed",
+                },
             },
         },
     }
@@ -512,8 +521,11 @@ def _set_val_if_key_exists(src: dict, dst: dict, key: str) -> None:
         if "extra_params" in kwargs:
             params.update(kwargs["extra_params"])
 
-        # headers
-        headers = kwargs.get("headers", None)
+        headers: HeadersType = {}
+        if self.api_type is APIType.AISTUDIO or self.api_type is APIType.CUSTOM:
+            headers["Content-Type"] = "application/json"
+        if "headers" in kwargs:
+            headers.update(kwargs["headers"])
 
         # request_timeout
         request_timeout = kwargs.get("request_timeout", None)