Spaces:
Sleeping
Sleeping
File size: 4,957 Bytes
d9cc70c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """OpenAI-compatible chat completion routes."""
from __future__ import annotations
import json
import logging
import time
from typing import Literal
from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field
from app.utils.config import settings
logger = logging.getLogger(__name__)
router = APIRouter(tags=["chat"])
class ChatMessage(BaseModel):
"""OpenAI-compatible chat message format."""
role: Literal["system", "user", "assistant"]
content: str
class ChatCompletionRequest(BaseModel):
"""Subset of OpenAI chat completion request fields."""
model: str = Field(default_factory=lambda: settings.model_name)
messages: list[ChatMessage]
stream: bool = False
temperature: float | None = None
top_p: float | None = None
max_tokens: int | None = Field(default=None, ge=1)
def _sse_event(payload: dict) -> str:
"""Format one SSE data event."""
return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
@router.post("/v1/chat/completions")
async def create_chat_completion(request: Request, body: ChatCompletionRequest):
"""OpenAI-compatible completions with optional SSE token streaming."""
prompt_service = request.app.state.prompt_service
model_service = request.app.state.model_service
injected_messages = prompt_service.inject_system_prompt(
[message.model_dump() for message in body.messages]
)
temperature = body.temperature if body.temperature is not None else settings.default_temperature
top_p = body.top_p if body.top_p is not None else settings.default_top_p
max_tokens = body.max_tokens if body.max_tokens is not None else settings.default_max_tokens
created = int(time.time())
if body.stream:
async def event_generator():
request_id = None
try:
# Initial chunk with assistant role to follow OpenAI streaming style.
bootstrap_chunk = {
"id": "chatcmpl-bootstrap",
"object": "chat.completion.chunk",
"created": created,
"model": body.model,
"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
}
yield _sse_event(bootstrap_chunk)
async for stream_request_id, delta in model_service.stream_text(
injected_messages,
temperature=temperature,
top_p=top_p,
max_tokens=max_tokens,
):
request_id = stream_request_id
chunk = {
"id": request_id,
"object": "chat.completion.chunk",
"created": created,
"model": body.model,
"choices": [{"index": 0, "delta": {"content": delta}, "finish_reason": None}],
}
yield _sse_event(chunk)
final_chunk = {
"id": request_id or "chatcmpl-final",
"object": "chat.completion.chunk",
"created": created,
"model": body.model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}
yield _sse_event(final_chunk)
yield "data: [DONE]\n\n"
except (RuntimeError, ValueError): # pragma: no cover - runtime guard
logger.exception("Failed to stream completion for request")
error_payload = {
"error": {
"message": "Failed to stream completion for request",
"type": "server_error",
}
}
yield _sse_event(error_payload)
yield "data: [DONE]\n\n"
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
request_id, text = await model_service.complete_text(
injected_messages,
temperature=temperature,
top_p=top_p,
max_tokens=max_tokens,
)
response_payload = {
"id": request_id,
"object": "chat.completion",
"created": created,
"model": body.model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": text},
"finish_reason": "stop",
}
],
}
return JSONResponse(response_payload)
|