refactor: replace compact response generation with length-prefixed response for backwards invocation api (#20903)
This commit is contained in:
@@ -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(
|
||||
|
Reference in New Issue
Block a user