|
|
|
@@ -3,6 +3,7 @@ import logging |
|
|
|
import re |
|
|
|
import secrets |
|
|
|
import string |
|
|
|
import struct |
|
|
|
import subprocess |
|
|
|
import time |
|
|
|
import uuid |
|
|
|
@@ -14,6 +15,7 @@ from zoneinfo import available_timezones |
|
|
|
|
|
|
|
from flask import Response, stream_with_context |
|
|
|
from flask_restful import fields |
|
|
|
from pydantic import BaseModel |
|
|
|
|
|
|
|
from configs import dify_config |
|
|
|
from core.app.features.rate_limiting.rate_limit import RateLimitGenerator |
|
|
|
@@ -206,6 +208,60 @@ def compact_generate_response(response: Union[Mapping, Generator, RateLimitGener |
|
|
|
return Response(stream_with_context(generate()), status=200, mimetype="text/event-stream") |
|
|
|
|
|
|
|
|
|
|
|
def length_prefixed_response(magic_number: int, response: Union[Mapping, Generator, RateLimitGenerator]) -> Response: |
|
|
|
""" |
|
|
|
This function is used to return a response with a length prefix. |
|
|
|
Magic number is a one byte number that indicates the type of the response. |
|
|
|
|
|
|
|
For a compatibility with latest plugin daemon https://github.com/langgenius/dify-plugin-daemon/pull/341 |
|
|
|
Avoid using line-based response, it leads a memory issue. |
|
|
|
|
|
|
|
We uses following format: |
|
|
|
| Field | Size | Description | |
|
|
|
|---------------|----------|---------------------------------| |
|
|
|
| Magic Number | 1 byte | Magic number identifier | |
|
|
|
| Reserved | 1 byte | Reserved field | |
|
|
|
| Header Length | 2 bytes | Header length (usually 0xa) | |
|
|
|
| Data Length | 4 bytes | Length of the data | |
|
|
|
| Reserved | 6 bytes | Reserved fields | |
|
|
|
| Data | Variable | Actual data content | |
|
|
|
|
|
|
|
| Reserved Fields | Header | Data | |
|
|
|
|-----------------|----------|----------| |
|
|
|
| 4 bytes total | Variable | Variable | |
|
|
|
|
|
|
|
all data is in little endian |
|
|
|
""" |
|
|
|
|
|
|
|
def pack_response_with_length_prefix(response: bytes) -> bytes: |
|
|
|
header_length = 0xA |
|
|
|
data_length = len(response) |
|
|
|
# | Magic Number 1byte | Reserved 1byte | Header Length 2bytes | Data Length 4bytes | Reserved 6bytes | Data |
|
|
|
return struct.pack("<BBHI", magic_number, 0, header_length, data_length) + b"\x00" * 6 + response |
|
|
|
|
|
|
|
if isinstance(response, dict): |
|
|
|
return Response( |
|
|
|
response=pack_response_with_length_prefix(json.dumps(jsonable_encoder(response)).encode("utf-8")), |
|
|
|
status=200, |
|
|
|
mimetype="application/json", |
|
|
|
) |
|
|
|
elif isinstance(response, BaseModel): |
|
|
|
return Response( |
|
|
|
response=pack_response_with_length_prefix(response.model_dump_json().encode("utf-8")), |
|
|
|
status=200, |
|
|
|
mimetype="application/json", |
|
|
|
) |
|
|
|
|
|
|
|
def generate() -> Generator: |
|
|
|
for chunk in response: |
|
|
|
if isinstance(chunk, str): |
|
|
|
yield pack_response_with_length_prefix(chunk.encode("utf-8")) |
|
|
|
else: |
|
|
|
yield pack_response_with_length_prefix(chunk) |
|
|
|
|
|
|
|
return Response(stream_with_context(generate()), status=200, mimetype="text/event-stream") |
|
|
|
|
|
|
|
|
|
|
|
class TokenManager: |
|
|
|
@classmethod |
|
|
|
def generate_token( |