feat: Suport JSON mode response format

Author: hustshawn

src/api/routers/chat.py

@@ -1,11 +1,11 @@
-from typing import Annotated
+from typing import Annotated, Final
 
 from fastapi import APIRouter, Depends, Body
 from fastapi.responses import StreamingResponse
 
 from api.auth import api_key_auth
 from api.models.bedrock import BedrockModel
-from api.schema import ChatRequest, ChatResponse, ChatStreamResponse
+from api.schema import ChatRequest, ChatResponse, ChatStreamResponse, SystemMessage
 from api.setting import DEFAULT_MODEL
 import logging
 
@@ -15,6 +15,12 @@
     # responses={404: {"description": "Not found"}},
 )
 
+JSON_MODE_SYSTEM_PROMPT: Final = SystemMessage(
+    name="json_mode",
+    role="system",
+    content="You are a helpful assistant designed to output JSON without extra text",
+)
+
 
 @router.post(
     "/completions",
@@ -43,6 +49,16 @@ async def chat_completions(
     else:
         chat_request.model = deployment
 
+    # JSON mode handling
+    # Ref: https://platform.openai.com/docs/guides/structured-outputs/json-mode
+    response_format = chat_request.response_format
+    if (
+        chat_request.stream
+        or response_format
+        and response_format["type"] == "json_object"
+    ):
+        chat_request.messages.append(JSON_MODE_SYSTEM_PROMPT)
+
     logging.debug(f"chat_request: {chat_request}")
 
     # Exception will be raised if model not supported.

src/api/schema.py

@@ -98,6 +98,7 @@ class ChatRequest(BaseModel):
     n: int | None = 1  # Not used
     tools: list[Tool] | None = None
     tool_choice: str | object = "auto"
+    response_format: dict | None = None
 
 
 class Usage(BaseModel):

test/test.ipynb renamed to tests/test.ipynb

@@ -0,0 +1,98 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "isInteractiveWindowMessageCell": true
+   },
+   "source": [
+    "Connected to openai2 (Python 3.11.4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from openai import AzureOpenAI\n",
+    "\n",
+    "client = AzureOpenAI(\n",
+    "    azure_endpoint=os.getenv(\"AZURE_OPENAI_ENDPOINT\", \"http://localhost:8000\"),\n",
+    "    api_key=os.getenv(\"AZURE_OPENAI_API_KEY\", \"bedrock\"),\n",
+    "    azure_deployment=\"gpt-4o\",\n",
+    "    api_version=\"2024-06-01\",\n",
+    ")\n",
+    "\n",
+    "response = client.chat.completions.create(\n",
+    "    model=\"gpt-4o\",\n",
+    "    response_format={\"type\": \"json_object\"},\n",
+    "    messages=[\n",
+    "        # {\n",
+    "        #     \"role\": \"system\",\n",
+    "        #     \"content\": \"You are a helpful assistant designed to output JSON without extra text\",\n",
+    "        # },\n",
+    "        {\"role\": \"user\", \"content\": \"Who won the world series in 2020?\"},\n",
+    "    ],\n",
+    ")\n",
+    "\n",
+    "content = response.choices[0].message.content"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "<class 'str'>\n",
+      "{\n",
+      "  \"winner\": \"Los Angeles Dodgers\"\n",
+      "}\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(type(content))  # The type of content is a JSON String\n",
+    "print(content)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "<class 'dict'>\n"
+     ]
+    }
+   ],
+   "source": [
+    "import json\n",
+    "\n",
+    "json_data = json.loads(content)\n",
+    "print(type(json_data))  # Customer can later parse the string to a JSON format."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "openai2",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}