We’ll start with implementing the non-streaming bit. Let’s start with modeling our request:

from typing import List, Optional
from pydantic import BaseModel
class ChatMessage(BaseModel):
role: str
content: str
class ChatCompletionRequest(BaseModel):
model: str = "mock-gpt-model"
messages: List[ChatMessage]
max_tokens: Optional[int] = 512
temperature: Optional[float] = 0.1
stream: Optional[bool] = False

The PyDantic model represents the request from the client, aiming to duplicate the API reference. For the sake of brevity, this model doesn’t implement all the specs, but reasonably the bare bones needed for it to work. In the event you’re missing a parameter that could be a a part of the API specs (like top_p), you possibly can simply add it to the model.

The ChatCompletionRequest models the parameters OpenAI uses of their requests. The chat API specs require specifying a listing of ChatMessage (like a chat history, the client is frequently accountable for keeping it and feeding back in at every request). Each chat message has a role attribute (often system, assistant , or user ) and a content attribute containing the actual message text.

Next, we’ll write our FastAPI chat completions endpoint:

import time
from fastapi import FastAPI
app = FastAPI(title="OpenAI-compatible API")
@app.post("/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
if request.messages and request.messages[0].role == 'user':
resp_content = "As a mock AI Assitant, I can only echo your last message:" + request.messages[-1].content
else:
resp_content = "As a mock AI Assitant, I can only echo your last message, but there have been no messages!"
return {
"id": "1337",
"object": "chat.completion",
"created": time.time(),
"model": request.model,
"decisions": [{
"message": ChatMessage(role="assistant", content=resp_content)
}]
}

That easy.

Testing our implementation

Assuming each code blocks are in a file called foremost.py, we’ll install two Python libraries in the environment of selection (all the time best to create a brand new one): pip install fastapi openai and launch the server from a terminal:

uvicorn foremost:app

Using one other terminal (or by launching the server within the background), we are going to open a Python console and copy-paste the next code, taken straight from OpenAI’s Python Client Reference:

from openai import OpenAI
# init client and hook up with localhost server
client = OpenAI(
api_key="fake-api-key",
base_url="http://localhost:8000" # change the default port if needed
)
# call API
chat_completion = client.chat.completions.create(
messages=[
{
"role": "user",
"content": "Say this is a test",
}
],
model="gpt-1337-turbo-pro-max",
)
# print the highest "selection" 
print(chat_completion.decisions[0].message.content)

In the event you’ve done all the things accurately, the response from the server ought to be accurately printed. It’s also value inspecting the chat_completion object to see that each one relevant attributes are as sent from our server. You must see something like this:

As LLM generation tends to be slow (computationally expensive), it’s value streaming your generated content back to the client, in order that the user can see the response because it’s being generated, without having to attend for it to complete. In the event you recall, we gave ChatCompletionRequest a boolean stream property — this lets the client request that the information be streamed back to it, reasonably than sent without delay.

This makes things only a bit more complex. We are going to create a generator function to wrap our mock response (in a real-world scenario, we are going to need a generator that is connected to our LLM generation)

import asyncio
import json
async def _resp_async_generator(text_resp: str):
# let's pretend every word is a token and return it over time
tokens = text_resp.split(" ")
for i, token in enumerate(tokens):
chunk = {
"id": i,
"object": "chat.completion.chunk",
"created": time.time(),
"model": "blah",
"decisions": [{"delta": {"content": token + " "}}],
}
yield f"data: {json.dumps(chunk)}nn"
await asyncio.sleep(1)
yield "data: [DONE]nn"

And now, we’d modify our original endpoint to return a StreamingResponse when stream==True

import time
from starlette.responses import StreamingResponse
app = FastAPI(title="OpenAI-compatible API")
@app.post("/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
if request.messages:
resp_content = "As a mock AI Assitant, I can only echo your last message:" + request.messages[-1].content
else:
resp_content = "As a mock AI Assitant, I can only echo your last message, but there wasn't one!"
if request.stream:
return StreamingResponse(_resp_async_generator(resp_content), media_type="application/x-ndjson")
return {
"id": "1337",
"object": "chat.completion",
"created": time.time(),
"model": request.model,
"decisions": [{
"message": ChatMessage(role="assistant", content=resp_content)        }]
}

Testing the streaming implementation

After restarting the uvicorn server, we’ll open up a Python console and put on this code (again, taken from OpenAI’s library docs)

from openai import OpenAI
# init client and hook up with localhost server
client = OpenAI(
api_key="fake-api-key",
base_url="http://localhost:8000" # change the default port if needed
)
stream = client.chat.completions.create(
model="mock-gpt-model",
messages=[{"role": "user", "content": "Say this is a test"}],
stream=True,
)
for chunk in stream:
print(chunk.decisions[0].delta.content or "")

You must see each word within the server’s response being slowly printed, mimicking token generation. We are able to inspect the last chunk object to see something like this:

Putting all of it together

Finally, within the gist below, you possibly can see all the code for the server.