Cómo crear una API compatible con OpenAI | de Saar Berkovich | marzo de 2024

Comenzaremos implementando el bit de no transmisión. Comencemos modelando nuestra solicitud:

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

El PyDantic El modelo representa la solicitud del cliente, con el objetivo de replicar la referencia API. En aras de la brevedad, este modelo no implementa todas las especificaciones, sino más bien lo básico necesario para que funcione. Si le falta un parámetro que forma parte del Especificaciones API (como top_p), simplemente puedes agregarlo al modelo.

El ChatCompletionRequest modela los parámetros que OpenAI utiliza en sus solicitudes. Las especificaciones de la API de chat requieren especificar una lista de ChatMessage (Al igual que un historial de chat, el cliente suele ser responsable de conservarlo y retroalimentarlo en cada solicitud). Cada mensaje de chat tiene un role atributo (generalmente system, assistant o user ) y un content atributo que contiene el texto del mensaje real.

A continuación, escribiremos nuestro punto final de finalización de chat FastAPI:

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 were no messages!"
return {
"id": "1337",
"object": "chat.completion",
"created": time.time(),
"model": request.model,
"choices": [{
"message": ChatMessage(role="assistant", content=resp_content)
}]
}

Así de sencillo.

Probando nuestra implementación

Suponiendo que ambos bloques de código están en un archivo llamado main.pyinstalaremos dos bibliotecas de Python en nuestro entorno de elección (siempre es mejor crear una nueva): pip install fastapi openai e iniciar el servidor desde una terminal:

uvicorn main:app

Usando otra terminal (o iniciando el servidor en segundo plano), abriremos una consola Python y copiaremos y pegaremos el siguiente código, tomado directamente de Referencia del cliente Python de OpenAI:

from openai import OpenAI
# init client and connect to 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 top "choice" 
print(chat_completion.choices[0].message.content)

Si ha hecho todo correctamente, la respuesta del servidor debería imprimirse correctamente. También vale la pena inspeccionar el chat_completion objeto para ver que todos los atributos relevantes sean tal como se envían desde nuestro servidor. Debería ver algo como esto:

Como la generación de LLM tiende a ser lenta (computacionalmente costosa), vale la pena transmitir el contenido generado al cliente, para que el usuario pueda ver la respuesta a medida que se genera, sin tener que esperar a que termine. Si recuerdas, le dimos ChatCompletionRequest un booleano stream propiedad: esto permite al cliente solicitar que los datos se le transmitan de nuevo, en lugar de enviarlos de inmediato.

Esto hace que las cosas sean un poco más complejas. Crearemos un función generadora para concluir nuestra respuesta simulada (en un escenario del mundo real, necesitaremos un generador que esté conectado a nuestra generación LLM)

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",
"choices": [{"delta": {"content": token + " "}}],
}
yield f"data: {json.dumps(chunk)}\n\n"
await asyncio.sleep(1)
yield "data: [DONE]\n\n"

Y ahora, modificaríamos nuestro punto final original para devolver un StreamingResponse cuando 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,
"choices": [{
"message": ChatMessage(role="assistant", content=resp_content)        }]
}

Probando la implementación de streaming

Después de reiniciar el servidor uvicorn, abriremos una consola Python e ingresaremos este código (nuevamente, tomado de los documentos de la biblioteca de OpenAI)

from openai import OpenAI
# init client and connect to 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.choices[0].delta.content or "")

Debería ver que cada palabra de la respuesta del servidor se imprime lentamente, imitando la generación de tokens. Podemos inspeccionar el último chunk objeto de ver algo como esto:

Poniendolo todo junto

Finalmente, en lo esencial a continuación, puede ver el código completo del servidor.