Repository: xorbitsai/inference

Issue:
FEAT: implement chat and generate in RESTful client



Discussion:

Resolving patch:
diff --git a/xinference/client.py b/xinference/client.py
--- a/xinference/client.py
+++ b/xinference/client.py
@@ -36,6 +36,166 @@
     )
 
 
+class ModelHandle:
+    """
+    A sync model interface (for rpc client) which provides type hints that makes it much easier to use xinference
+    programmatically.
+    """
+
+    def __init__(self, model_ref: xo.ActorRefType["ModelActor"], isolation: Isolation):
+        self._model_ref = model_ref
+        self._isolation = isolation
+
+
+class LlamaCppModelHandle(ModelHandle):
+    def generate(
+        self, prompt: str, generate_config: Optional["LlamaCppGenerateConfig"] = None
+    ) -> Union["Completion", Iterator["CompletionChunk"]]:
+        coro = self._model_ref.generate(prompt, generate_config)
+        return self._isolation.call(coro)
+
+
+class LlamaCppChatModelHandle(LlamaCppModelHandle):
+    def chat(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        chat_history: Optional[List["ChatCompletionMessage"]] = None,
+        generate_config: Optional["LlamaCppGenerateConfig"] = None,
+    ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
+        coro = self._model_ref.chat(
+            prompt, system_prompt, chat_history, generate_config
+        )
+        return self._isolation.call(coro)
+
+
+class ChatglmCppChatModelHandle(ModelHandle):
+    def chat(
+        self,
+        prompt: str,
+        chat_history: Optional[List["ChatCompletionMessage"]] = None,
+        generate_config: Optional["ChatglmCppGenerateConfig"] = None,
+    ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
+        coro = self._model_ref.chat(prompt, chat_history, generate_config)
+        return self._isolation.call(coro)
+
+
+class RESTfulModelHandle:
+    """
+    A sync model interface (for RESTful client) which provides type hints that makes it much easier to use xinference
+    programmatically.
+    """
+
+    def __init__(self, model_uid: str, base_url: str):
+        self._model_uid = model_uid
+        self._base_url = base_url
+
+
+class RESTfulLlamaCppModelHandle(RESTfulModelHandle):
+    def generate(
+        self, prompt: str, generate_config: Optional["LlamaCppGenerateConfig"] = None
+    ) -> Union["Completion", Iterator["CompletionChunk"]]:
+        url = f"{self._base_url}/v1/completions"
+        if generate_config is None:
+            request_body = {"model": self._model_uid, "prompt": prompt}
+        else:
+            generate_config_dict = {}
+            for key, value in generate_config.items():
+                generate_config_dict[str(key)] = str(value)
+            request_body = {
+                "model": self._model_uid,
+                "prompt": prompt,
+                **generate_config_dict,
+            }
+
+        response = requests.post(url, json=request_body)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to generate completion, detail: {response.json()['detail']}"
+            )
+        response_data = response.json()
+        return response_data
+
+
+class RESTfulLlamaCppChatModelHandle(RESTfulLlamaCppModelHandle):
+    def chat(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        chat_history: Optional[List["ChatCompletionMessage"]] = None,
+        generate_config: Optional["LlamaCppGenerateConfig"] = None,
+    ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
+        url = f"{self._base_url}/v1/chat/completions"
+
+        if chat_history is None:
+            chat_history = []
+
+        if chat_history and chat_history[0]["role"] == "system":
+            if system_prompt is not None:
+                chat_history[0]["content"] = system_prompt
+
+        else:
+            if system_prompt is not None:
+                chat_history.insert(0, {"role": "system", "content": system_prompt})
+
+        chat_history.append({"role": "user", "content": prompt})
+
+        if generate_config is None:
+            request_body = {"model": self._model_uid, "messages": chat_history}
+        else:
+            generate_config_dict = {}
+            for key, value in generate_config.items():
+                generate_config_dict[str(key)] = str(value)
+            request_body = {
+                "model": self._model_uid,
+                "messages": chat_history,
+                **generate_config_dict,
+            }
+
+        response = requests.post(url, json=request_body)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to generate chat completion, detail: {response.json()['detail']}"
+            )
+        response_data = response.json()
+        return response_data
+
+
+class RESTfulChatglmCppChatModelHandle(RESTfulModelHandle):
+    def chat(
+        self,
+        prompt: str,
+        chat_history: Optional[List["ChatCompletionMessage"]] = None,
+        generate_config: Optional["ChatglmCppGenerateConfig"] = None,
+    ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
+        url = f"{self._base_url}/v1/chat/completions"
+
+        if chat_history is None:
+            chat_history = []
+
+        chat_history.append({"role": "user", "content": prompt})
+
+        if generate_config is None:
+            request_body = {"model": self._model_uid, "messages": chat_history}
+        else:
+            generate_config_dict = {}
+            for key, value in generate_config.items():
+                generate_config_dict[str(key)] = str(value)
+            request_body = {
+                "model": self._model_uid,
+                "messages": chat_history,
+                **generate_config_dict,
+            }
+
+        response = requests.post(url, json=request_body)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to generate chat completion, detail: {response.json()['detail']}"
+            )
+        response_data = response.json()
+        return response_data
+
+
 class Client:
     def __init__(self, endpoint: str):
         restful_client = RESTfulClient(endpoint)
@@ -95,50 +255,6 @@ def get_model(self, model_uid: str) -> "ModelHandle":
             return LlamaCppChatModelHandle(model_ref, self._isolation)
 
 
-class ModelHandle:
-    """
-    A sync model interface which provides type hints that makes it much easier to use xinference
-    programmatically.
-    """
-
-    def __init__(self, model_ref: xo.ActorRefType["ModelActor"], isolation: Isolation):
-        self._model_ref = model_ref
-        self._isolation = isolation
-
-
-class LlamaCppModelHandle(ModelHandle):
-    def generate(
-        self, prompt: str, generate_config: Optional["LlamaCppGenerateConfig"] = None
-    ) -> Union["Completion", Iterator["CompletionChunk"]]:
-        coro = self._model_ref.generate(prompt, generate_config)
-        return self._isolation.call(coro)
-
-
-class LlamaCppChatModelHandle(LlamaCppModelHandle):
-    def chat(
-        self,
-        prompt: str,
-        system_prompt: Optional[str] = None,
-        chat_history: Optional[List["ChatCompletionMessage"]] = None,
-        generate_config: Optional["LlamaCppGenerateConfig"] = None,
-    ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
-        coro = self._model_ref.chat(
-            prompt, system_prompt, chat_history, generate_config
-        )
-        return self._isolation.call(coro)
-
-
-class ChatglmCppChatModelHandle(ModelHandle):
-    def chat(
-        self,
-        prompt: str,
-        chat_history: Optional[List["ChatCompletionMessage"]] = None,
-        generate_config: Optional["ChatglmCppGenerateConfig"] = None,
-    ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
-        coro = self._model_ref.chat(prompt, chat_history, generate_config)
-        return self._isolation.call(coro)
-
-
 class RESTfulClient:
     def __init__(self, base_url):
         self.base_url = base_url
@@ -205,3 +321,22 @@ def _get_supervisor_internal_address(self):
             raise RuntimeError(f"Failed to get supervisor internal address")
         response_data = response.json()
         return response_data
+
+    def get_model(self, model_uid: str) -> RESTfulModelHandle:
+        url = f"{self.base_url}/v1/models/{model_uid}"
+        response = requests.get(url)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get the model description, detail: {response.json()['detail']}"
+            )
+        model_spec = response.json()
+
+        if (
+            model_spec["model_name"] == "chatglm"
+            or model_spec["model_name"] == "chatglm2"
+        ):
+            return RESTfulChatglmCppChatModelHandle(model_uid, self.base_url)
+        elif model_spec["model_name"] == "baichuan":
+            return RESTfulLlamaCppModelHandle(model_uid, self.base_url)
+        else:
+            return RESTfulLlamaCppChatModelHandle(model_uid, self.base_url)
diff --git a/xinference/core/restful_api.py b/xinference/core/restful_api.py
--- a/xinference/core/restful_api.py
+++ b/xinference/core/restful_api.py
@@ -32,7 +32,7 @@
 logger = logging.getLogger(__name__)
 
 max_tokens_field = Field(
-    default=16, ge=1, le=2048, description="The maximum number of tokens to generate."
+    default=128, ge=1, le=2048, description="The maximum number of tokens to generate."
 )
 
 temperature_field = Field(
@@ -236,6 +236,9 @@ def serve(self):
         )
         self._router = APIRouter()
         self._router.add_api_route("/v1/models", self.list_models, methods=["GET"])
+        self._router.add_api_route(
+            "/v1/models/{model_uid}", self.describe_model, methods=["GET"]
+        )
         self._router.add_api_route("/v1/models", self.launch_model, methods=["POST"])
         self._router.add_api_route(
             "/v1/models/{model_uid}", self.terminate_model, methods=["DELETE"]
@@ -288,6 +291,13 @@ async def list_models(self) -> Dict[str, Dict[str, Any]]:
             }
         return models_dict
 
+    async def describe_model(self, model_uid: str):
+        try:
+            return await self._supervisor_ref.describe_model(model_uid)
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
     async def launch_model(self, request: Request) -> JSONResponse:
         payload = await request.json()
         model_uid = payload.get("model_uid")


Test patch:
diff --git a/xinference/core/tests/test_restful_api.py b/xinference/core/tests/test_restful_api.py
--- a/xinference/core/tests/test_restful_api.py
+++ b/xinference/core/tests/test_restful_api.py
@@ -40,6 +40,12 @@ async def test_restful_api(setup):
     response_data = response.json()
     assert len(response_data) == 1
 
+    # describe
+    response = requests.get(f"{endpoint}/v1/models/test")
+    response_data = response.json()
+    print(response_data)
+    assert response_data["model_name"] == "orca"
+
     # generate
     url = f"{endpoint}/v1/completions"
     payload = {
diff --git a/xinference/tests/test_client.py b/xinference/tests/test_client.py
--- a/xinference/tests/test_client.py
+++ b/xinference/tests/test_client.py
@@ -14,7 +14,7 @@
 
 import pytest
 
-from ..client import Client
+from ..client import Client, RESTfulClient
 
 
 @pytest.mark.asyncio
@@ -35,3 +35,34 @@ async def test_sync_client(setup):
 
     client.terminate_model(model_uid=model_uid)
     assert len(client.list_models()) == 0
+
+
+@pytest.mark.asyncio
+async def test_RESTful_client(setup):
+    endpoint, _ = setup
+    client = RESTfulClient(endpoint)
+    assert len(client.list_models()) == 0
+
+    model_uid = client.launch_model(
+        model_name="orca", model_size_in_billions=3, quantization="q4_0"
+    )
+    assert len(client.list_models()) == 1
+
+    model = client.get_model(model_uid=model_uid)
+
+    completion = model.generate("Once upon a time, there was a very old computer")
+    assert "text" in completion["choices"][0]
+
+    completion = model.generate(
+        "Once upon a time, there was a very old computer", {"max_tokens": 256}
+    )
+    assert "text" in completion["choices"][0]
+
+    completion = model.chat("write a poem.")
+    assert "content" in completion["choices"][0]["message"]
+
+    client.terminate_model(model_uid=model_uid)
+    assert len(client.list_models()) == 0
+
+    with pytest.raises(RuntimeError):
+        client.terminate_model(model_uid=model_uid)