Repository: xorbitsai/inference
Pull request: ENH: Implement RESTful API stream generate
Resolves #166.

Problem statement:
ENH: RESTful API stream generate
### Is your feature request related to a problem? Please describe
Provide stream generateion in RESTful API.

### Describe the solution you'd like
A clear and concise description of what you want to happen.

### Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

### Additional context
Add any other context or screenshots about the feature request here.

Hints from issue discussion:

Commit message: exception handling

Patch to review:
diff --git a/setup.cfg b/setup.cfg
index 8c1a49e68d..627bbb8d12 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -34,6 +34,7 @@ install_requires =
     pydantic
     fastapi
     uvicorn
+    sse_starlette
 
 [options.packages.find]
 exclude =
diff --git a/xinference/client.py b/xinference/client.py
index 2d99035cff..c3461f957c 100644
--- a/xinference/client.py
+++ b/xinference/client.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import json
 import uuid
 from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Union
 
@@ -80,6 +81,34 @@ def chat(
         return self._isolation.call(coro)
 
 
+def streaming_response_iterator(response_lines):
+    for line in response_lines:
+        line = line.strip()
+        if line:
+            if line == b"data: [DONE]":
+                yield {"data": "End of Response"}
+                break
+            try:
+                data = json.loads(line.decode("utf-8").replace("data: ", "", 1))
+                yield {"data": data}
+            except json.JSONDecodeError:
+                print("Error decoding JSON for line:", line)
+
+
+def chat_streaming_response_iterator(response_lines):
+    for line in response_lines:
+        line = line.strip()
+        if line:
+            if line == b"data: [DONE]":
+                yield {"data": "End of Response"}
+                break
+            try:
+                data = json.loads(line.decode("utf-8").replace("data: ", "", 1))
+                yield {"data": data}
+            except json.JSONDecodeError:
+                print("Error decoding JSON for line:", line)
+
+
 class RESTfulModelHandle:
     """
     A sync model interface (for RESTful client) which provides type hints that makes it much easier to use xinference
@@ -113,6 +142,10 @@ def generate(
             raise RuntimeError(
                 f"Failed to generate completion, detail: {response.json()['detail']}"
             )
+
+        if generate_config and generate_config.get("stream"):
+            return streaming_response_iterator(response.iter_lines())
+
         response_data = response.json()
         return response_data
 
@@ -157,6 +190,10 @@ def chat(
             raise RuntimeError(
                 f"Failed to generate chat completion, detail: {response.json()['detail']}"
             )
+
+        if generate_config and generate_config.get("stream"):
+            return chat_streaming_response_iterator(response.iter_lines())
+
         response_data = response.json()
         return response_data
 
@@ -192,6 +229,10 @@ def chat(
             raise RuntimeError(
                 f"Failed to generate chat completion, detail: {response.json()['detail']}"
             )
+
+        if generate_config and generate_config.get("stream"):
+            return chat_streaming_response_iterator(response.iter_lines())
+
         response_data = response.json()
         return response_data
 
@@ -270,7 +311,7 @@ def list_models(self) -> Dict[str, Dict[str, Any]]:
         response = requests.get(url)
         if response.status_code != 200:
             raise RuntimeError(
-                f"Failed to launch model, detail: {response.json()['detail']}"
+                f"Failed to list model, detail: {response.json()['detail']}"
             )
 
         response_data = response.json()
diff --git a/xinference/core/restful_api.py b/xinference/core/restful_api.py
index 7453857c24..21a44fbc7b 100644
--- a/xinference/core/restful_api.py
+++ b/xinference/core/restful_api.py
@@ -12,17 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import logging
 import socket
 import threading
+from functools import partial
 from typing import Any, Dict, List, Literal, Optional, Union
 
+import anyio
 import gradio as gr
 import xoscar as xo
+from anyio.streams.memory import MemoryObjectSendStream
 from fastapi import APIRouter, FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel, Field
+from sse_starlette.sse import EventSourceResponse
 from typing_extensions import NotRequired, TypedDict
 from uvicorn import Config, Server
 
@@ -294,6 +299,11 @@ async def list_models(self) -> Dict[str, Dict[str, Any]]:
     async def describe_model(self, model_uid: str):
         try:
             return await self._supervisor_ref.describe_model(model_uid)
+
+        except ValueError as ve:
+            logger.error(str(ve), exc_info=True)
+            raise HTTPException(status_code=400, detail=str(ve))
+
         except Exception as e:
             logger.error(e, exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
@@ -307,6 +317,12 @@ async def launch_model(self, request: Request) -> JSONResponse:
         quantization = payload.get("quantization")
         kwargs = payload.get("kwargs", {}) or {}
 
+        if model_uid is None or model_uid is None:
+            raise HTTPException(
+                status_code=400,
+                detail="Invalid input. Please specify the model UID and the model name",
+            )
+
         try:
             await self._supervisor_ref.launch_builtin_model(
                 model_uid=model_uid,
@@ -316,14 +332,29 @@ async def launch_model(self, request: Request) -> JSONResponse:
                 quantization=quantization,
                 **kwargs,
             )
+
+        except ValueError as ve:
+            logger.error(str(ve), exc_info=True)
+            raise HTTPException(status_code=400, detail=str(ve))
+
+        except RuntimeError as re:
+            logger.error(str(re), exc_info=True)
+            raise HTTPException(status_code=503, detail=str(re))
+
         except Exception as e:
-            logger.error(e, exc_info=True)
+            logger.error(str(e), exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
+
         return JSONResponse(content={"model_uid": model_uid})
 
     async def terminate_model(self, model_uid: str):
         try:
             await self._supervisor_ref.terminate_model(model_uid)
+
+        except ValueError as ve:
+            logger.error(str(ve), exc_info=True)
+            raise HTTPException(status_code=400, detail=str(ve))
+
         except Exception as e:
             logger.error(e, exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
@@ -344,17 +375,49 @@ async def create_completion(self, request: Request, body: CreateCompletionReques
         kwargs = body.dict(exclude=exclude)
 
         if body.logit_bias is not None:
-            raise NotImplementedError
+            raise HTTPException(status_code=501, detail="Not implemented")
+
         model_uid = body.model
 
         try:
             model = await self._supervisor_ref.get_model(model_uid)
+
+        except ValueError as ve:
+            logger.error(str(ve), exc_info=True)
+            raise HTTPException(status_code=400, detail=str(ve))
+
         except Exception as e:
             logger.error(e, exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
 
         if body.stream:
-            raise NotImplementedError
+            # create a pair of memory object streams
+            send_chan, recv_chan = anyio.create_memory_object_stream(10)
+
+            async def event_publisher(inner_send_chan: MemoryObjectSendStream):
+                async with inner_send_chan:
+                    try:
+                        iterator = await model.generate(body.prompt, kwargs)
+                        async for chunk in iterator:
+                            await inner_send_chan.send(dict(data=json.dumps(chunk)))
+                            if await request.is_disconnected():
+                                raise anyio.get_cancelled_exc_class()()
+                        await inner_send_chan.send(dict(data="[DONE]"))
+                    except anyio.get_cancelled_exc_class() as e:
+                        print("disconnected")
+                        with anyio.move_on_after(1, shield=True):
+                            print(
+                                f"Disconnected from client (via refresh/close) {request.client}"
+                            )
+                            await inner_send_chan.send(dict(closing=True))
+                            raise e
+                    except Exception as e:
+                        raise HTTPException(status_code=500, detail=str(e))
+
+            return EventSourceResponse(
+                recv_chan, data_sender_callable=partial(event_publisher, send_chan)
+            )
+
         else:
             try:
                 return await model.generate(body.prompt, kwargs)
@@ -363,7 +426,7 @@ async def create_completion(self, request: Request, body: CreateCompletionReques
                 raise HTTPException(status_code=500, detail=str(e))
 
     async def create_embedding(self, request: CreateEmbeddingRequest):
-        raise NotImplementedError
+        raise HTTPException(status_code=501, detail="Not implemented")
 
     async def create_chat_completion(
         self,
@@ -381,15 +444,19 @@ async def create_chat_completion(
         kwargs = body.dict(exclude=exclude)
 
         if body.logit_bias is not None:
-            raise NotImplementedError
+            raise HTTPException(status_code=501, detail="Not implemented")
+
+        if (
+            not body.messages
+            or body.messages[-1].get("role") != "user"
+            or not body.messages[-1].get("content")
+        ):
+            raise HTTPException(
+                status_code=400, detail="Invalid input. Please specify the prompt"
+            )
+
+        prompt = body.messages[-1]["content"]
 
-        user_messages = [
-            msg["content"] for msg in body.messages if msg["role"] == "user"
-        ]
-        if user_messages:
-            prompt = user_messages[-1]
-        else:
-            raise HTTPException(status_code=400, detail="No prompt given")
         system_prompt = next(
             (msg["content"] for msg in body.messages if msg["role"] == "system"), None
         )
@@ -397,14 +464,47 @@ async def create_chat_completion(
         chat_history = body.messages
 
         model_uid = body.model
+
         try:
             model = await self._supervisor_ref.get_model(model_uid)
+
+        except ValueError as ve:
+            logger.error(str(ve), exc_info=True)
+            raise HTTPException(status_code=400, detail=str(ve))
         except Exception as e:
             logger.error(e, exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
 
         if body.stream:
-            raise NotImplementedError
+            # create a pair of memory object streams
+            send_chan, recv_chan = anyio.create_memory_object_stream(10)
+
+            async def event_publisher(inner_send_chan: MemoryObjectSendStream):
+                async with inner_send_chan:
+                    try:
+                        iterator = await model.chat(
+                            prompt, system_prompt, chat_history, kwargs
+                        )
+                        async for chunk in iterator:
+                            await inner_send_chan.send(dict(data=json.dumps(chunk)))
+                            if await request.is_disconnected():
+                                raise anyio.get_cancelled_exc_class()()
+                        await inner_send_chan.send(dict(data="[DONE]"))
+                    except anyio.get_cancelled_exc_class() as e:
+                        print("disconnected")
+                        with anyio.move_on_after(1, shield=True):
+                            print(
+                                f"Disconnected from client (via refresh/close) {request.client}"
+                            )
+                            await inner_send_chan.send(dict(closing=True))
+                            raise e
+                    except Exception as e:
+                        raise HTTPException(status_code=500, detail=str(e))
+
+            return EventSourceResponse(
+                recv_chan, data_sender_callable=partial(event_publisher, send_chan)
+            )
+
         else:
             try:
                 return await model.chat(prompt, system_prompt, chat_history, kwargs)
diff --git a/xinference/core/service.py b/xinference/core/service.py
index c7d424a515..bf8452f592 100644
--- a/xinference/core/service.py
+++ b/xinference/core/service.py
@@ -111,7 +111,8 @@ async def launch_builtin_model(
             quantization,
         )
 
-        assert model_uid not in self._model_uid_to_worker
+        if model_uid in self._model_uid_to_worker:
+            raise ValueError(f"Model is already in the model list, uid: {model_uid}")
 
         worker_ref = await self._choose_worker()
         model_ref = yield worker_ref.launch_builtin_model(
@@ -146,7 +147,8 @@ async def _check_dead_nodes(self):
 
     @log
     async def terminate_model(self, model_uid: str):
-        assert model_uid in self._model_uid_to_worker
+        if model_uid not in self._model_uid_to_worker:
+            raise ValueError(f"Model not found in the model list, uid: {model_uid}")
 
         worker_ref = self._model_uid_to_worker[model_uid]
         await worker_ref.terminate_model(model_uid=model_uid)
@@ -154,11 +156,17 @@ async def terminate_model(self, model_uid: str):
 
     @log
     async def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
+        if model_uid not in self._model_uid_to_worker:
+            raise ValueError(f"Model not found in the model list, uid: {model_uid}")
+
         worker_ref = self._model_uid_to_worker[model_uid]
         return await worker_ref.get_model(model_uid=model_uid)
 
     @log
     async def describe_model(self, model_uid: str):
+        if model_uid not in self._model_uid_to_worker:
+            raise ValueError(f"Model not found in the model list, uid: {model_uid}")
+
         worker_ref = self._model_uid_to_worker[model_uid]
         return await worker_ref.describe_model(model_uid=model_uid)
 
@@ -288,9 +296,11 @@ async def launch_builtin_model(
 
     @log
     async def terminate_model(self, model_uid: str):
-        assert model_uid in self._model_uid_to_model
+        if model_uid not in self._model_uid_to_model:
+            raise ValueError(f"Model not found in the model list, uid: {model_uid}")
 
         model_ref = self._model_uid_to_model[model_uid]
+
         await xo.destroy_actor(model_ref)
         del self._model_uid_to_model[model_uid]
         del self._model_uid_to_model_spec[model_uid]
@@ -307,10 +317,16 @@ async def list_models(self) -> List[Tuple[str, ModelSpec]]:
 
     @log
     async def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
+        if model_uid not in self._model_uid_to_model:
+            raise ValueError(f"Model not found in the model list, uid: {model_uid}")
+
         return self._model_uid_to_model[model_uid]
 
     @log
     async def describe_model(self, model_uid: str) -> ModelSpec:
+        if model_uid not in self._model_uid_to_model:
+            raise ValueError(f"Model not found in the model list, uid: {model_uid}")
+
         return self._model_uid_to_model_spec[model_uid]
 
     async def report_status(self):
diff --git a/xinference/core/tests/test_restful_api.py b/xinference/core/tests/test_restful_api.py
index 13824b0847..5e656bafc1 100644
--- a/xinference/core/tests/test_restful_api.py
+++ b/xinference/core/tests/test_restful_api.py
@@ -35,6 +35,14 @@ async def test_restful_api(setup):
     model_uid_res = response_data["model_uid"]
     assert model_uid_res == "test"
 
+    payload = {"model_uid": "test", "model_name": "orca", "quantization": "q4_0"}
+    response = requests.post(url, json=payload)
+    assert response.status_code == 400
+
+    payload = {"model_name": "orca", "quantization": "q4_0"}
+    response = requests.post(url, json=payload)
+    assert response.status_code == 400
+
     # list
     response = requests.get(url)
     response_data = response.json()
@@ -43,9 +51,11 @@ async def test_restful_api(setup):
     # describe
     response = requests.get(f"{endpoint}/v1/models/test")
     response_data = response.json()
-    print(response_data)
     assert response_data["model_name"] == "orca"
 
+    response = requests.delete(f"{endpoint}/v1/models/bogus")
+    assert response.status_code == 400
+
     # generate
     url = f"{endpoint}/v1/completions"
     payload = {
@@ -56,6 +66,19 @@ async def test_restful_api(setup):
     completion = response.json()
     assert "text" in completion["choices"][0]
 
+    payload = {
+        "model": "bogus",
+        "prompt": "Once upon a time, there was a very old computer.",
+    }
+    response = requests.post(url, json=payload)
+    assert response.status_code == 400
+
+    payload = {
+        "prompt": "Once upon a time, there was a very old computer.",
+    }
+    response = requests.post(url, json=payload)
+    assert response.status_code == 422
+
     # chat
     url = f"{endpoint}/v1/chat/completions"
     payload = {
@@ -71,6 +94,40 @@ async def test_restful_api(setup):
     completion = response.json()
     assert "content" in completion["choices"][0]["message"]
 
+    payload = {
+        "messages": [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Hello!"},
+            {"role": "assistant", "content": "Hi what can I help you?"},
+            {"role": "user", "content": "What is the capital of France?"},
+        ],
+    }
+    response = requests.post(url, json=payload)
+    assert response.status_code == 422
+
+    payload = {
+        "model": "bogus",
+        "messages": [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Hello!"},
+            {"role": "assistant", "content": "Hi what can I help you?"},
+            {"role": "user", "content": "What is the capital of France?"},
+        ],
+    }
+    response = requests.post(url, json=payload)
+    assert response.status_code == 400
+
+    payload = {
+        "model": model_uid_res,
+        "messages": [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Hello!"},
+            {"role": "assistant", "content": "Hi what can I help you?"},
+        ],
+    }
+    response = requests.post(url, json=payload)
+    assert response.status_code == 400
+
     # delete
     url = f"{endpoint}/v1/models/test"
     response = requests.delete(url)
@@ -83,4 +140,4 @@ async def test_restful_api(setup):
     # delete again
     url = f"{endpoint}/v1/models/test"
     response = requests.delete(url)
-    assert response.status_code != 200
+    assert response.status_code == 400
diff --git a/xinference/tests/test_client.py b/xinference/tests/test_client.py
index 23d02c4ed6..9d489b0b09 100644
--- a/xinference/tests/test_client.py
+++ b/xinference/tests/test_client.py
@@ -50,17 +50,45 @@ async def test_RESTful_client(setup):
 
     model = client.get_model(model_uid=model_uid)
 
+    with pytest.raises(RuntimeError):
+        model = client.get_model(model_uid="test")
+
+    with pytest.raises(RuntimeError):
+        completion = model.generate({"max_tokens": 64})
+
     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}
+        "Once upon a time, there was a very old computer", {"max_tokens": 64}
     )
     assert "text" in completion["choices"][0]
 
-    completion = model.chat("write a poem.")
+    streaming_response = model.generate(
+        "Once upon a time, there was a very old computer",
+        {"max_tokens": 64, "stream": True},
+    )
+
+    for chunk in streaming_response:
+        assert (
+            chunk["data"] == "End of Response" or "text" in chunk["data"]["choices"][0]
+        )
+
+    with pytest.raises(RuntimeError):
+        completion = model.chat({"max_tokens": 64})
+
+    completion = model.chat("What is the capital of France?")
     assert "content" in completion["choices"][0]["message"]
 
+    streaming_response = model.chat(
+        prompt="What is the capital of France?", generate_config={"stream": True}
+    )
+
+    for chunk in streaming_response:
+        assert chunk["data"] == "End of Response" or (
+            "content" or "role" in chunk["data"]["choices"][0]["delta"]
+        )
+
     client.terminate_model(model_uid=model_uid)
     assert len(client.list_models()) == 0
 


Code review comments:
diff_hunk: @@ -80,6 +81,34 @@ def chat(
         return self._isolation.call(coro)
 
 
+def streaming_response_iterator(response_lines):
+    for line in response_lines:
+        line = line.strip()
+        if line:
+            if line == b"data: [DONE]":
path: xinference/client.py
text: @user1:
The ending message is not necessary.

diff_hunk: @@ -80,6 +81,34 @@ def chat(
         return self._isolation.call(coro)
 
 
+def streaming_response_iterator(response_lines):
+    for line in response_lines:
+        line = line.strip()
+        if line:
+            if line == b"data: [DONE]":
+                yield {"data": "End of Response"}
+                break
+            try:
+                data = json.loads(line.decode("utf-8").replace("data: ", "", 1))
+                yield {"data": data}
+            except json.JSONDecodeError:
+                print("Error decoding JSON for line:", line)
path: xinference/client.py
text: @user1:
The error should not be captured here.

diff_hunk: @@ -316,14 +332,29 @@ async def launch_model(self, request: Request) -> JSONResponse:
                 quantization=quantization,
                 **kwargs,
             )
+
+        except ValueError as ve:
+            logger.error(str(ve), exc_info=True)
+            raise HTTPException(status_code=400, detail=str(ve))
+
+        except RuntimeError as re:
+            logger.error(str(re), exc_info=True)
+            raise HTTPException(status_code=503, detail=str(re))
+
         except Exception as e:
-            logger.error(e, exc_info=True)
+            logger.error(str(e), exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
+
         return JSONResponse(content={"model_uid": model_uid})
 
     async def terminate_model(self, model_uid: str):
         try:
             await self._supervisor_ref.terminate_model(model_uid)
+
path: xinference/core/restful_api.py
text: @user1:
```suggestion
```

diff_hunk: @@ -344,17 +375,49 @@ async def create_completion(self, request: Request, body: CreateCompletionReques
         kwargs = body.dict(exclude=exclude)
 
         if body.logit_bias is not None:
-            raise NotImplementedError
+            raise HTTPException(status_code=501, detail="Not implemented")
+
         model_uid = body.model
 
         try:
             model = await self._supervisor_ref.get_model(model_uid)
+
+        except ValueError as ve:
+            logger.error(str(ve), exc_info=True)
+            raise HTTPException(status_code=400, detail=str(ve))
+
         except Exception as e:
             logger.error(e, exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
 
         if body.stream:
-            raise NotImplementedError
+            # create a pair of memory object streams
+            send_chan, recv_chan = anyio.create_memory_object_stream(10)
+
+            async def event_publisher(inner_send_chan: MemoryObjectSendStream):
+                async with inner_send_chan:
+                    try:
+                        iterator = await model.generate(body.prompt, kwargs)
+                        async for chunk in iterator:
+                            await inner_send_chan.send(dict(data=json.dumps(chunk)))
+                            if await request.is_disconnected():
+                                raise anyio.get_cancelled_exc_class()()
+                        await inner_send_chan.send(dict(data="[DONE]"))
path: xinference/core/restful_api.py
text: @user1:
The ending message is not compatible with the OpenAI API.

diff_hunk: @@ -344,17 +375,49 @@ async def create_completion(self, request: Request, body: CreateCompletionReques
         kwargs = body.dict(exclude=exclude)
 
         if body.logit_bias is not None:
-            raise NotImplementedError
+            raise HTTPException(status_code=501, detail="Not implemented")
+
         model_uid = body.model
 
         try:
             model = await self._supervisor_ref.get_model(model_uid)
+
path: xinference/core/restful_api.py
text: @user1:
```suggestion
```

diff_hunk: @@ -80,6 +81,34 @@ def chat(
         return self._isolation.call(coro)
 
 
+def streaming_response_iterator(response_lines):
path: xinference/client.py
text: @user1:
Type hint.