JSON formatted response using OpenAI API types for server completion requests

Author: vmpuri

api/api.py

@@ -95,8 +95,7 @@ class CompletionRequest:
     """
 
     model: str
-    prompt: str
-    messages: Optional[List[_AbstractMessage]]
+    messages: List[_AbstractMessage]
     frequency_penalty: float = 0.0
     temperature: float = 0.0
     stop: Optional[List[str]] = None
@@ -121,10 +120,10 @@ class CompletionChoice:
     See the "The chat completion object >>> choices" section of the OpenAI API docs for more details.
     """
 
-    finish_reason: str
     index: int
     message: AssistantMessage
-    logprobs: Optional[List[Any]]
+    finish_reason: str = None
+    logprobs: Optional[List[Any]] = None
 
 
 @dataclass
@@ -151,7 +150,7 @@ class CompletionResponse:
     created: int
     model: str
     system_fingerprint: str
-    usage: UsageStats
+    usage: Optional[UsageStats] = None
     object: str = "chat.completion"
     service_tier: Optional[str] = None
 
@@ -220,8 +219,11 @@ def __init__(self, *args, **kwargs):
             if self.draft_model is not None
             else self.model.config.max_seq_length
         )
+        self.system_fingerprint = (
+            self.builder_args.device + type(self.builder_args.precision).__name__
+        )
 
-    def completion(self, completion_request: CompletionRequest):
+    def chunked_completion(self, completion_request: CompletionRequest):
         """Handle a chat completion request and yield a chunked response.
 
         ** Warning ** : Not all arguments of the CompletionRequest are consumed as the server isn't completely implemented.
@@ -246,13 +248,16 @@ def completion(self, completion_request: CompletionRequest):
 
         # Initialize counters for chunk responses and encode the prompt.
         id = str(uuid.uuid4())
+
         idx = 0
         buffer = []
         encoded = self.encode_tokens(
-            completion_request.prompt, bos=True, device=self.builder_args.device
+            completion_request.messages[-1].get("content"),
+            bos=True,
+            device=self.builder_args.device,
         )
         generator_args = GeneratorArgs(
-            completion_request.prompt,
+            completion_request.messages[-1].get("content"),
             encoded_prompt=encoded,
             chat_mode=False,
         )
@@ -302,21 +307,45 @@ def callback(x, *, done_generating=False):
                 choices=[choice_chunk],
                 created=int(time.time()),
                 model=completion_request.model,
-                system_fingerprint=uuid.UUID(int=uuid.getnode()),
+                system_fingerprint=self.system_fingerprint,
             )
             yield chunk_response
             self.start_pos += y.size(0)
             idx += 1
 
         # Yield an ending chunk indicating the generation has completed.
-        end_chunk = CompletionChoiceChunk(ChunkDelta(None, None, None), idx, "eos")
+        end_chunk = CompletionChoiceChunk(
+            ChunkDelta(None, None, None), idx, finish_reason="stop"
+        )
 
         yield CompletionResponseChunk(
             id=str(id),
             choices=[end_chunk],
             created=int(time.time()),
             model=completion_request.model,
-            system_fingerprint=uuid.UUID(int=uuid.getnode()),
+            system_fingerprint=self.system_fingerprint,
+        )
+
+    def sync_completion(self, request: CompletionRequest):
+        """Handle a chat completion request and yield a single, non-chunked response"""
+        output = ""
+        for chunk in self.chunked_completion(request):
+            if not chunk.choices[0].finish_reason:
+                output += chunk.choices[0].delta.content
+
+        message = AssistantMessage(content=output)
+        return CompletionResponse(
+            id=str(uuid.uuid4()),
+            choices=[
+                CompletionChoice(
+                    finish_reason="stop",
+                    index=0,
+                    message=message,
+                )
+            ],
+            created=int(time.time()),
+            model=request.model,
+            system_fingerprint=self.system_fingerprint,
         )
 
     def _callback(self, x, *, buffer, done_generating):

server.py

@@ -4,18 +4,35 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-from api.api import AssistantMessage, CompletionRequest, OpenAiApiGenerator
+import json
+
+from dataclasses import asdict
+
+from api.api import AssistantMessage, CompletionRequest, OpenAiApiGenerator, UserMessage
 
 from build.builder import BuilderArgs, TokenizerArgs
-from flask import Flask, jsonify, request, Response
+from flask import Flask, request, Response
 from generate import GeneratorArgs
 
+
+"""
+Creates a flask app that can be used to serve the model as a chat API.
+"""
 app = Flask(__name__)
 # Messages and gen are kept global so they can be accessed by the flask app endpoints.
 messages: list = []
 gen: OpenAiApiGenerator = None
 
 
+def _del_none(d: dict):
+    """Recursively delete None values from a dictionary."""
+    if type(d) is dict:
+        return {k: _del_none(v) for k, v in d.items() if v}
+    elif type(d) is list:
+        return [_del_none(v) for v in d if v]
+    return d
+
+
 @app.route("/chat", methods=["POST"])
 def chat_endpoint():
     """
@@ -26,45 +43,44 @@ def chat_endpoint():
 
     See https://github.com/pytorch/torchchat/issues/973 and the OpenAiApiGenerator class for more details.
 
+    If stream is set to true, the response will be streamed back as a series of CompletionResponseChunk objects. Otherwise,
+    a single CompletionResponse object will be returned.
     """
-    data = request.get_json()
-
-    # Add user message to chat history
-    messages.append(data["messages"][-1])
-    prompt = messages[-1]["content"]
 
-    # Generate the assistant response
-    req = CompletionRequest(
-        model=gen.builder_args.checkpoint_path,
-        prompt=prompt,
-        temperature=0,
-        messages=[],
-    )
+    print(" === Completion Request ===")
 
-    response = ""
+    # Parse the request in to a CompletionRequest object
+    data = request.get_json()
+    req = CompletionRequest(**data)
 
-    def unwrap(completion_generator):
-        token_count = 0
-        for chunk_response in completion_generator:
-            content = chunk_response.choices[0].delta.content
-            if not gen.is_llama3_model or content not in set(
-                gen.tokenizer.special_tokens.keys()
-            ):
-                yield content if content is not None else ""
-            if content == gen.tokenizer.eos_id():
-                yield "."
-            token_count += 1
+    # Add the user message to our internal message history.
+    messages.append(UserMessage(**req.messages[-1]))
 
     if data.get("stream") == "true":
-        return Response(unwrap(gen.completion(req)), mimetype="text/event-stream")
+
+        def chunk_processor(chunked_completion_generator):
+            """Inline function for postprocessing CompletionResponseChunk objects.
+
+            Here, we just jsonify the chunk and yield it as a string.
+            """
+            messages.append(AssistantMessage(content=""))
+            for chunk in chunked_completion_generator:
+                nextok = chunk.choices[0].delta.content
+                nextok = nextok if nextok is not None else ""
+                messages[-1].content += nextok
+                print(nextok, end="")
+                yield json.dumps(_del_none(asdict(chunk)))
+
+        return Response(
+            chunk_processor(gen.chunked_completion(req)), mimetype="text/event-stream"
+        )
     else:
-        for content in unwrap(gen.completion(req)):
-            response += content
+        response = gen.sync_completion(req)
 
-    # Add assistant response to chat history
-    messages.append(AssistantMessage(content=response))
+        messages.append(response.choices[0].message)
+        print(messages[-1].content)
 
-    return jsonify({"response": response})
+        return json.dumps(_del_none(asdict(response)))
 
 
 def initialize_generator(args) -> OpenAiApiGenerator: