Django 流式调用

2024-08-20

import os
import json
from collections.abc import Generator

from openai import OpenAI
from django.conf import settings
from django.http import HttpRequest
from django.http import StreamingHttpResponse
from django.urls import path
from django.core.wsgi import get_wsgi_application
from django.views.decorators.csrf import csrf_exempt
from django.views.decorators.http import require_http_methods

# Django settings
settings.configure(
    DEBUG=True,
    SECRET_KEY="your_secret_key_here",
    ROOT_URLCONF=__name__,
    MIDDLEWARE=[
        "django.middleware.common.CommonMiddleware",
        "django.middleware.csrf.CsrfViewMiddleware",
    ],
    INSTALLED_APPS=[
        "django.contrib.auth",
        "django.contrib.contenttypes",
    ],
)

# OpenAI client setup
openai_client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL"),
)


def generate_stream(text: str) -> Generator[str, None, None]:
    stream = openai_client.chat.completions.create(
        model="yi-medium",
        messages=[{"role": "user", "content": text}],
        stream=True,
    )
    for chunk in stream:
        current_response = chunk.choices[0].delta.content
        if current_response:
            yield current_response


@csrf_exempt
@require_http_methods(["POST"])
def chat(request: HttpRequest) -> StreamingHttpResponse:
    data = json.loads(request.body)
    text = data.get("text")
    if not text:
        return StreamingHttpResponse("Error: No text provided", status=400)

    response = StreamingHttpResponse(generate_stream(text))
    response["Cache-Control"] = "no-cache"
    response["X-Accel-Buffering"] = "no"
    response["Content-Type"] = "text/event-stream; charset=utf-8"

    # 添加 CORS 头
    response["Access-Control-Allow-Origin"] = "*"
    response["Access-Control-Allow-Methods"] = "POST, OPTIONS"
    response["Access-Control-Allow-Headers"] = "Content-Type"
    return response


# URL configuration
urlpatterns = [
    path("chat", chat, name="chat"),
]

# WSGI application
application = get_wsgi_application()

if __name__ == "__main__":
    from django.core.management import execute_from_command_line

    execute_from_command_line(["manage.py", "runserver", "0.0.0.0:8000"])