mirror of
https://github.com/jie65535/JChatGPT.git
synced 2026-09-15 02:56:10 +08:00
llm: harden streaming responses
This commit is contained in:
@@ -17,9 +17,12 @@ import kotlinx.coroutines.flow.flow
|
|||||||
import kotlinx.coroutines.isActive
|
import kotlinx.coroutines.isActive
|
||||||
import kotlinx.coroutines.sync.Semaphore
|
import kotlinx.coroutines.sync.Semaphore
|
||||||
import kotlinx.coroutines.sync.withPermit
|
import kotlinx.coroutines.sync.withPermit
|
||||||
import kotlinx.coroutines.withTimeout
|
import kotlinx.coroutines.withTimeoutOrNull
|
||||||
import kotlinx.serialization.json.*
|
import kotlinx.serialization.json.*
|
||||||
import okhttp3.Dispatcher as OkHttpDispatcher
|
import okhttp3.Dispatcher as OkHttpDispatcher
|
||||||
|
import okhttp3.Protocol as OkHttpProtocol
|
||||||
|
import java.io.ByteArrayOutputStream
|
||||||
|
import java.io.IOException
|
||||||
import kotlin.time.Duration
|
import kotlin.time.Duration
|
||||||
|
|
||||||
class ModelService(
|
class ModelService(
|
||||||
@@ -35,15 +38,16 @@ class ModelService(
|
|||||||
|
|
||||||
val httpClient: HttpClient by lazy {
|
val httpClient: HttpClient by lazy {
|
||||||
HttpClient(OkHttp) {
|
HttpClient(OkHttp) {
|
||||||
this@ModelService.maxConcurrentRequests?.let { concurrencyLimit ->
|
engine {
|
||||||
engine {
|
config {
|
||||||
config {
|
protocols(MODEL_HTTP_PROTOCOLS)
|
||||||
|
this@ModelService.maxConcurrentRequests?.let { concurrencyLimit ->
|
||||||
dispatcher(createRequestDispatcher(concurrencyLimit))
|
dispatcher(createRequestDispatcher(concurrencyLimit))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
install(HttpTimeout) {
|
install(HttpTimeout) {
|
||||||
// 流式响应的「首 token」与「token 间隔」超时统一由应用层 withTimeout 管控(见 chatCompletions)。
|
// 流式响应的「首 token」与「token 间隔」超时统一由应用层计时管控(见 chatCompletions)。
|
||||||
// 这里特意不设 requestTimeoutMillis:否则正常但耗时较长的流式输出会被 Ktor 在中途整体掐断。
|
// 这里特意不设 requestTimeoutMillis:否则正常但耗时较长的流式输出会被 Ktor 在中途整体掐断。
|
||||||
// socket 超时作为字节级兜底,连接超时只覆盖 TCP 握手。
|
// socket 超时作为字节级兜底,连接超时只覆盖 TCP 握手。
|
||||||
socketTimeoutMillis = timeout.inWholeMilliseconds
|
socketTimeoutMillis = timeout.inWholeMilliseconds
|
||||||
@@ -86,24 +90,19 @@ class ModelService(
|
|||||||
request: ChatCompletionRequest,
|
request: ChatCompletionRequest,
|
||||||
onCacheUsage: ((CacheUsage) -> Unit)? = null
|
onCacheUsage: ((CacheUsage) -> Unit)? = null
|
||||||
): Flow<ChatCompletionChunk> {
|
): Flow<ChatCompletionChunk> {
|
||||||
val requestJson = json.encodeToJsonElement(ChatCompletionRequest.serializer(), request)
|
val body = buildRequestBody(request, stream = true)
|
||||||
.jsonObject.toMutableMap()
|
|
||||||
requestJson["stream"] = JsonPrimitive(true)
|
|
||||||
extraBody?.forEach { (key, value) ->
|
|
||||||
requestJson[key] = value
|
|
||||||
}
|
|
||||||
val body = JsonObject(requestJson).toString()
|
|
||||||
|
|
||||||
val responseFlow: Flow<ChatCompletionChunk> = flow {
|
val responseFlow: Flow<ChatCompletionChunk> = flow {
|
||||||
// 关键:服务器繁忙时会拖住「响应头」,使 httpClient.post() 自身阻塞在等待响应的阶段,
|
// 关键:服务器繁忙时会拖住「响应头」,使 httpClient.post() 自身阻塞在等待响应的阶段,
|
||||||
// 因此必须把 post() 连同首个 data 块的读取一起包进 withTimeout。
|
// 因此必须把 post() 连同首个 data 块的读取一起纳入同一个应用层超时。
|
||||||
// 否则首 token 超时永远不会触发(post() 还没返回,根本进不到读取循环),
|
// 否则首 token 超时永远不会触发(post() 还没返回,根本进不到读取循环),
|
||||||
// 只能落到 Ktor 的兜底超时(很久)后再重试,表现为「等很久才报异常」。
|
// 只能落到 Ktor 的兜底超时(很久)后再重试,表现为「等很久才报异常」。
|
||||||
// channel 在 withTimeout 外层持有:哪怕首块读取在 withTimeout 内超时,
|
// channel 在超时块外层持有:哪怕首块读取超时,
|
||||||
// 只要 response.body() 已拿到通道,finally 也能释放它,避免慢速 API 重试时连接泄漏。
|
// 只要 response.body() 已拿到通道,finally 也能释放它,避免慢速 API 重试时连接泄漏。
|
||||||
var channel: ByteReadChannel? = null
|
var channel: ByteReadChannel? = null
|
||||||
|
var lineReader: LenientUtf8LineReader? = null
|
||||||
try {
|
try {
|
||||||
val firstDataLine = withTimeout(firstChunkTimeout) {
|
val firstDataLine = withModelResponseTimeout(firstChunkTimeout, "首个响应数据块") {
|
||||||
val response = httpClient.post("chat/completions") {
|
val response = httpClient.post("chat/completions") {
|
||||||
setBody(body)
|
setBody(body)
|
||||||
contentType(ContentType.Application.Json)
|
contentType(ContentType.Application.Json)
|
||||||
@@ -115,9 +114,11 @@ class ModelService(
|
|||||||
}
|
}
|
||||||
val ch: ByteReadChannel = response.body()
|
val ch: ByteReadChannel = response.body()
|
||||||
channel = ch
|
channel = ch
|
||||||
|
val reader = LenientUtf8LineReader(ch)
|
||||||
|
lineReader = reader
|
||||||
var found: String? = null
|
var found: String? = null
|
||||||
while (currentCoroutineContext().isActive && !ch.isClosedForRead) {
|
while (currentCoroutineContext().isActive) {
|
||||||
val line = ch.readUTF8Line() ?: continue
|
val line = reader.readLine() ?: break
|
||||||
if (line.startsWith("data: ")) {
|
if (line.startsWith("data: ")) {
|
||||||
found = line
|
found = line
|
||||||
break
|
break
|
||||||
@@ -129,19 +130,22 @@ class ModelService(
|
|||||||
|
|
||||||
if (firstDataLine != null && !firstDataLine.startsWith("data: [DONE]")) {
|
if (firstDataLine != null && !firstDataLine.startsWith("data: [DONE]")) {
|
||||||
val firstRaw = firstDataLine.removePrefix("data: ")
|
val firstRaw = firstDataLine.removePrefix("data: ")
|
||||||
emit(json.decodeFromString(firstRaw))
|
decodeStreamChunk(firstRaw)?.let { emit(it) }
|
||||||
onCacheUsage?.let { cb -> extractCacheUsage(firstRaw)?.let(cb) }
|
onCacheUsage?.let { cb -> extractCacheUsage(firstRaw)?.let(cb) }
|
||||||
|
|
||||||
val ch = channel!!
|
val ch = channel!!
|
||||||
while (currentCoroutineContext().isActive && !ch.isClosedForRead) {
|
val reader = lineReader!!
|
||||||
|
while (currentCoroutineContext().isActive) {
|
||||||
// 流式期间同样对每次读取设「token 间隔」超时,避免中途卡死后干等兜底超时,
|
// 流式期间同样对每次读取设「token 间隔」超时,避免中途卡死后干等兜底超时,
|
||||||
// 从而能快速失败并交给上层重试。正常流式 token 间隔远小于 firstChunkTimeout。
|
// 从而能快速失败并交给上层重试。正常流式 token 间隔远小于 firstChunkTimeout。
|
||||||
val line = withTimeout(firstChunkTimeout) { ch.readUTF8Line() } ?: continue
|
val line = withModelResponseTimeout(firstChunkTimeout, "流式响应数据块") {
|
||||||
|
reader.readLine()
|
||||||
|
} ?: break
|
||||||
when {
|
when {
|
||||||
line.startsWith("data: [DONE]") -> break
|
line.startsWith("data: [DONE]") -> break
|
||||||
line.startsWith("data: ") -> {
|
line.startsWith("data: ") -> {
|
||||||
val raw = line.removePrefix("data: ")
|
val raw = line.removePrefix("data: ")
|
||||||
emit(json.decodeFromString(raw))
|
decodeStreamChunk(raw)?.let { emit(it) }
|
||||||
onCacheUsage?.let { cb -> extractCacheUsage(raw)?.let(cb) }
|
onCacheUsage?.let { cb -> extractCacheUsage(raw)?.let(cb) }
|
||||||
}
|
}
|
||||||
else -> continue
|
else -> continue
|
||||||
@@ -154,9 +158,215 @@ class ModelService(
|
|||||||
}
|
}
|
||||||
return responseFlow.withConcurrencyLimit(requestSemaphore)
|
return responseFlow.withConcurrencyLimit(requestSemaphore)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
internal fun buildRequestBody(request: ChatCompletionRequest, stream: Boolean): String {
|
||||||
|
val requestJson = json.encodeToJsonElement(ChatCompletionRequest.serializer(), request)
|
||||||
|
.jsonObject.toMutableMap()
|
||||||
|
requestJson["stream"] = JsonPrimitive(stream)
|
||||||
|
extraBody?.forEach { (key, value) ->
|
||||||
|
requestJson[key] = value
|
||||||
|
}
|
||||||
|
return JsonObject(requestJson).toString()
|
||||||
|
}
|
||||||
|
|
||||||
|
internal fun decodeStreamChunk(rawJson: String): ChatCompletionChunk? {
|
||||||
|
val element = try {
|
||||||
|
json.parseToJsonElement(rawJson)
|
||||||
|
} catch (cause: Exception) {
|
||||||
|
throw ModelStreamProtocolException(
|
||||||
|
"模型 SSE data 不是有效 JSON:${rawJson.toDiagnosticSnippet()}",
|
||||||
|
cause,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
val payload = element as? JsonObject ?: throw ModelStreamProtocolException(
|
||||||
|
"模型 SSE data 必须是 JSON object:${rawJson.toDiagnosticSnippet()}"
|
||||||
|
)
|
||||||
|
if (payload.isEmpty()) return null
|
||||||
|
|
||||||
|
payload["error"]?.let { error ->
|
||||||
|
val errorObject = error as? JsonObject
|
||||||
|
val errorType = (errorObject?.get("type") as? JsonPrimitive)?.contentOrNull
|
||||||
|
val errorCode = (errorObject?.get("code") as? JsonPrimitive)?.contentOrNull
|
||||||
|
val errorMessage = (errorObject?.get("message") as? JsonPrimitive)?.contentOrNull
|
||||||
|
val diagnostic = "模型 SSE 返回错误事件:${error.toString().toDiagnosticSnippet()}"
|
||||||
|
if (isSafetyRejection(errorCode, errorMessage)) {
|
||||||
|
throw ModelSafetyRejectionException(errorType, errorCode, diagnostic)
|
||||||
|
}
|
||||||
|
if (errorType == "invalid_request_error" &&
|
||||||
|
isDeterministicRequestRejection(errorCode, errorMessage)
|
||||||
|
) {
|
||||||
|
throw ModelRequestRejectedException(errorType, errorCode, diagnostic)
|
||||||
|
}
|
||||||
|
throw ModelStreamProtocolException(diagnostic)
|
||||||
|
}
|
||||||
|
|
||||||
|
val requiredChunkFields = listOf("id", "created", "model", "choices")
|
||||||
|
val normalized = if (requiredChunkFields.all(payload::containsKey)) {
|
||||||
|
payload
|
||||||
|
} else if ("usage" in payload) {
|
||||||
|
JsonObject(
|
||||||
|
buildMap {
|
||||||
|
putAll(payload)
|
||||||
|
putIfAbsent("id", JsonPrimitive("usage-only"))
|
||||||
|
putIfAbsent("object", JsonPrimitive("chat.completion.chunk"))
|
||||||
|
putIfAbsent("created", JsonPrimitive(0))
|
||||||
|
putIfAbsent("model", JsonPrimitive("unknown"))
|
||||||
|
putIfAbsent("choices", JsonArray(emptyList()))
|
||||||
|
}
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
throw ModelStreamProtocolException(
|
||||||
|
"模型 SSE 事件缺少 chunk 字段,keys=${payload.keys.sorted()}:" +
|
||||||
|
rawJson.toDiagnosticSnippet()
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
return try {
|
||||||
|
json.decodeFromJsonElement(ChatCompletionChunk.serializer(), normalized)
|
||||||
|
} catch (cause: Exception) {
|
||||||
|
throw ModelStreamProtocolException(
|
||||||
|
"模型 SSE chunk 结构无效,keys=${payload.keys.sorted()}:" +
|
||||||
|
rawJson.toDiagnosticSnippet(),
|
||||||
|
cause,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun String.toDiagnosticSnippet(maxChars: Int = 500): String =
|
||||||
|
replace(Regex("[\\r\\n]+"), " ").let { normalized ->
|
||||||
|
if (normalized.length <= maxChars) normalized else normalized.take(maxChars) + "...[截断]"
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun isSafetyRejection(code: String?, message: String?): Boolean {
|
||||||
|
val normalizedCode = code.orEmpty().lowercase()
|
||||||
|
val normalizedMessage = message.orEmpty().lowercase()
|
||||||
|
val explicitPolicyCode = normalizedCode == "cyber_policy" ||
|
||||||
|
normalizedCode == "safety_policy" ||
|
||||||
|
normalizedCode == "moderation_blocked" ||
|
||||||
|
normalizedCode == "content_filter" ||
|
||||||
|
normalizedCode.contains("content_policy")
|
||||||
|
val explicitContentSafetyMessage =
|
||||||
|
(normalizedMessage.contains("this content") &&
|
||||||
|
normalizedMessage.contains("safety reason")) ||
|
||||||
|
(normalizedMessage.contains("flagged for possible") &&
|
||||||
|
normalizedMessage.contains("risk")) ||
|
||||||
|
normalizedMessage.contains("biological risk") ||
|
||||||
|
normalizedMessage.contains("cybersecurity risk")
|
||||||
|
return explicitPolicyCode || explicitContentSafetyMessage
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun isDeterministicRequestRejection(code: String?, message: String?): Boolean {
|
||||||
|
val normalizedCode = code.orEmpty().lowercase()
|
||||||
|
val normalizedMessage = message.orEmpty().lowercase()
|
||||||
|
return normalizedCode.contains("invalid_prompt") ||
|
||||||
|
normalizedCode.contains("context_length") ||
|
||||||
|
normalizedCode.contains("invalid_parameter") ||
|
||||||
|
normalizedCode.contains("unsupported_parameter") ||
|
||||||
|
normalizedMessage.contains("invalid prompt") ||
|
||||||
|
normalizedMessage.contains("maximum context length") ||
|
||||||
|
normalizedMessage.contains("unsupported parameter") ||
|
||||||
|
normalizedMessage.contains("missing required parameter")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
internal class LenientUtf8LineReader(
|
||||||
|
private val channel: ByteReadChannel,
|
||||||
|
) {
|
||||||
|
private val readBuffer = ByteArray(DEFAULT_BUFFER_SIZE)
|
||||||
|
private val lineBuffer = ByteArrayOutputStream()
|
||||||
|
private var readOffset = 0
|
||||||
|
private var readLimit = 0
|
||||||
|
private var skipLeadingLineFeed = false
|
||||||
|
|
||||||
|
suspend fun readLine(): String? {
|
||||||
|
while (true) {
|
||||||
|
if (readOffset >= readLimit) {
|
||||||
|
val read = channel.readAvailable(readBuffer)
|
||||||
|
if (read < 0) {
|
||||||
|
channel.closedCause?.let { throw it }
|
||||||
|
return if (lineBuffer.size() == 0) null else decodeBufferedLine()
|
||||||
|
}
|
||||||
|
if (read == 0) continue
|
||||||
|
readOffset = 0
|
||||||
|
readLimit = read
|
||||||
|
}
|
||||||
|
|
||||||
|
if (skipLeadingLineFeed) {
|
||||||
|
skipLeadingLineFeed = false
|
||||||
|
if (readBuffer[readOffset] == LINE_FEED) {
|
||||||
|
readOffset++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
val segmentStart = readOffset
|
||||||
|
while (readOffset < readLimit) {
|
||||||
|
val current = readBuffer[readOffset]
|
||||||
|
if (current == CARRIAGE_RETURN || current == LINE_FEED) break
|
||||||
|
readOffset++
|
||||||
|
}
|
||||||
|
if (readOffset > segmentStart) {
|
||||||
|
lineBuffer.write(readBuffer, segmentStart, readOffset - segmentStart)
|
||||||
|
}
|
||||||
|
if (readOffset >= readLimit) continue
|
||||||
|
|
||||||
|
val delimiter = readBuffer[readOffset++]
|
||||||
|
if (delimiter == CARRIAGE_RETURN) skipLeadingLineFeed = true
|
||||||
|
return decodeBufferedLine()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun decodeBufferedLine(): String {
|
||||||
|
val bytes = lineBuffer.toByteArray()
|
||||||
|
lineBuffer.reset()
|
||||||
|
// String(byte[], UTF_8) uses replacement semantics for malformed input instead of aborting the SSE stream.
|
||||||
|
return String(bytes, Charsets.UTF_8)
|
||||||
|
}
|
||||||
|
|
||||||
|
private companion object {
|
||||||
|
const val DEFAULT_BUFFER_SIZE = 8 * 1024
|
||||||
|
const val CARRIAGE_RETURN: Byte = '\r'.code.toByte()
|
||||||
|
const val LINE_FEED: Byte = '\n'.code.toByte()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
internal class ModelResponseTimeoutException(
|
||||||
|
stage: String,
|
||||||
|
timeout: Duration,
|
||||||
|
) : IOException("等待模型${stage}超过 ${timeout.inWholeMilliseconds} ms")
|
||||||
|
|
||||||
|
internal class ModelStreamProtocolException(
|
||||||
|
message: String,
|
||||||
|
cause: Throwable? = null,
|
||||||
|
) : IOException(message, cause)
|
||||||
|
|
||||||
|
internal open class ModelRequestRejectedException(
|
||||||
|
val errorType: String?,
|
||||||
|
val errorCode: String?,
|
||||||
|
message: String,
|
||||||
|
) : IOException(message)
|
||||||
|
|
||||||
|
internal class ModelSafetyRejectionException(
|
||||||
|
errorType: String?,
|
||||||
|
errorCode: String?,
|
||||||
|
message: String,
|
||||||
|
) : ModelRequestRejectedException(errorType, errorCode, message)
|
||||||
|
|
||||||
|
internal suspend fun <T> withModelResponseTimeout(
|
||||||
|
timeout: Duration,
|
||||||
|
stage: String,
|
||||||
|
block: suspend () -> T,
|
||||||
|
): T {
|
||||||
|
val completed = withTimeoutOrNull(timeout) {
|
||||||
|
CompletedModelResponse(block())
|
||||||
|
} ?: throw ModelResponseTimeoutException(stage, timeout)
|
||||||
|
return completed.value
|
||||||
|
}
|
||||||
|
|
||||||
|
private data class CompletedModelResponse<T>(val value: T)
|
||||||
|
|
||||||
internal const val MAX_MODEL_CONCURRENT_REQUESTS = 512
|
internal const val MAX_MODEL_CONCURRENT_REQUESTS = 512
|
||||||
|
internal val MODEL_HTTP_PROTOCOLS = listOf(OkHttpProtocol.HTTP_1_1)
|
||||||
|
|
||||||
internal fun normalizeMaxConcurrentRequests(value: Int): Int =
|
internal fun normalizeMaxConcurrentRequests(value: Int): Int =
|
||||||
value.coerceIn(1, MAX_MODEL_CONCURRENT_REQUESTS)
|
value.coerceIn(1, MAX_MODEL_CONCURRENT_REQUESTS)
|
||||||
|
|||||||
@@ -56,10 +56,7 @@ class ProfileModelClient(
|
|||||||
profile: UserProfileSnapshot,
|
profile: UserProfileSnapshot,
|
||||||
batch: ProfileHistoryBatch,
|
batch: ProfileHistoryBatch,
|
||||||
): ProfileModelResult {
|
): ProfileModelResult {
|
||||||
val content = StringBuilder()
|
val completion = complete(
|
||||||
var lastUsage: Usage? = null
|
|
||||||
var cacheUsage: ModelService.CacheUsage? = null
|
|
||||||
endpoint.service.chatCompletions(
|
|
||||||
ChatCompletionRequest(
|
ChatCompletionRequest(
|
||||||
model = ModelId(endpoint.model),
|
model = ModelId(endpoint.model),
|
||||||
temperature = endpoint.temperature,
|
temperature = endpoint.temperature,
|
||||||
@@ -70,21 +67,13 @@ class ProfileModelClient(
|
|||||||
ChatMessage.User(ProfilePromptStore.buildUserPrompt(profile, batch)),
|
ChatMessage.User(ProfilePromptStore.buildUserPrompt(profile, batch)),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
) { cacheUsage = it }.collect { chunk ->
|
)
|
||||||
chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
|
val raw = completion.content.replace(THINK_REGEX, "").trim()
|
||||||
chunk.usage?.let { lastUsage = it }
|
|
||||||
}
|
|
||||||
|
|
||||||
val raw = content.toString().replace(THINK_REGEX, "").trim()
|
|
||||||
val response = parseResponse(raw)
|
val response = parseResponse(raw)
|
||||||
return ProfileModelResult(
|
return ProfileModelResult(
|
||||||
response = response,
|
response = response,
|
||||||
rawResponse = raw,
|
rawResponse = raw,
|
||||||
usage = ProfileTokenUsage(
|
usage = completion.usage,
|
||||||
promptTokens = lastUsage?.promptTokens ?: 0,
|
|
||||||
completionTokens = lastUsage?.completionTokens ?: 0,
|
|
||||||
cachedTokens = cacheUsage?.hitTokens ?: 0,
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -93,10 +82,7 @@ class ProfileModelClient(
|
|||||||
batch: ConversationProfileBatch,
|
batch: ConversationProfileBatch,
|
||||||
eligibleUserIds: Set<Long>,
|
eligibleUserIds: Set<Long>,
|
||||||
): ConversationProfileModelResult {
|
): ConversationProfileModelResult {
|
||||||
val content = StringBuilder()
|
val completion = complete(
|
||||||
var lastUsage: Usage? = null
|
|
||||||
var cacheUsage: ModelService.CacheUsage? = null
|
|
||||||
endpoint.service.chatCompletions(
|
|
||||||
ChatCompletionRequest(
|
ChatCompletionRequest(
|
||||||
model = ModelId(endpoint.model),
|
model = ModelId(endpoint.model),
|
||||||
temperature = endpoint.temperature,
|
temperature = endpoint.temperature,
|
||||||
@@ -109,20 +95,12 @@ class ProfileModelClient(
|
|||||||
),
|
),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
) { cacheUsage = it }.collect { chunk ->
|
)
|
||||||
chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
|
val raw = completion.content.replace(THINK_REGEX, "").trim()
|
||||||
chunk.usage?.let { lastUsage = it }
|
|
||||||
}
|
|
||||||
|
|
||||||
val raw = content.toString().replace(THINK_REGEX, "").trim()
|
|
||||||
return ConversationProfileModelResult(
|
return ConversationProfileModelResult(
|
||||||
response = parseObject(raw),
|
response = parseObject(raw),
|
||||||
rawResponse = raw,
|
rawResponse = raw,
|
||||||
usage = ProfileTokenUsage(
|
usage = completion.usage,
|
||||||
promptTokens = lastUsage?.promptTokens ?: 0,
|
|
||||||
completionTokens = lastUsage?.completionTokens ?: 0,
|
|
||||||
cachedTokens = cacheUsage?.hitTokens ?: 0,
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -130,10 +108,7 @@ class ProfileModelClient(
|
|||||||
profile: UserProfileSnapshot,
|
profile: UserProfileSnapshot,
|
||||||
supportStats: Map<String, ProfileItemSupportStats>,
|
supportStats: Map<String, ProfileItemSupportStats>,
|
||||||
): ProfileCompactionModelResult {
|
): ProfileCompactionModelResult {
|
||||||
val content = StringBuilder()
|
val completion = complete(
|
||||||
var lastUsage: Usage? = null
|
|
||||||
var cacheUsage: ModelService.CacheUsage? = null
|
|
||||||
endpoint.service.chatCompletions(
|
|
||||||
ChatCompletionRequest(
|
ChatCompletionRequest(
|
||||||
model = ModelId(endpoint.model),
|
model = ModelId(endpoint.model),
|
||||||
temperature = endpoint.temperature,
|
temperature = endpoint.temperature,
|
||||||
@@ -144,23 +119,36 @@ class ProfileModelClient(
|
|||||||
ChatMessage.User(ProfilePromptStore.buildCompactionUserPrompt(profile, supportStats)),
|
ChatMessage.User(ProfilePromptStore.buildCompactionUserPrompt(profile, supportStats)),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
) { cacheUsage = it }.collect { chunk ->
|
)
|
||||||
chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
|
val raw = completion.content.replace(THINK_REGEX, "").trim()
|
||||||
chunk.usage?.let { lastUsage = it }
|
|
||||||
}
|
|
||||||
|
|
||||||
val raw = content.toString().replace(THINK_REGEX, "").trim()
|
|
||||||
return ProfileCompactionModelResult(
|
return ProfileCompactionModelResult(
|
||||||
response = json.decodeFromString(extractObject(raw)),
|
response = json.decodeFromString(extractObject(raw)),
|
||||||
rawResponse = raw,
|
rawResponse = raw,
|
||||||
usage = ProfileTokenUsage(
|
usage = completion.usage,
|
||||||
promptTokens = lastUsage?.promptTokens ?: 0,
|
|
||||||
completionTokens = lastUsage?.completionTokens ?: 0,
|
|
||||||
cachedTokens = cacheUsage?.hitTokens ?: 0,
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private suspend fun complete(request: ChatCompletionRequest): CompletedProfileResponse {
|
||||||
|
val content = StringBuilder()
|
||||||
|
var lastUsage: Usage? = null
|
||||||
|
var cacheUsage: ModelService.CacheUsage? = null
|
||||||
|
endpoint.service.chatCompletions(request) { cacheUsage = it }.collect { chunk ->
|
||||||
|
chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
|
||||||
|
chunk.usage?.let { lastUsage = it }
|
||||||
|
}
|
||||||
|
require(content.isNotBlank()) { "模型流式响应没有文本内容" }
|
||||||
|
return CompletedProfileResponse(
|
||||||
|
content = content.toString(),
|
||||||
|
usage = lastUsage.toProfileUsage(cacheUsage),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
private fun Usage?.toProfileUsage(cacheUsage: ModelService.CacheUsage?) = ProfileTokenUsage(
|
||||||
|
promptTokens = this?.promptTokens ?: 0,
|
||||||
|
completionTokens = this?.completionTokens ?: 0,
|
||||||
|
cachedTokens = cacheUsage?.hitTokens ?: 0,
|
||||||
|
)
|
||||||
|
|
||||||
private fun parseResponse(raw: String): ProfileModelResponse {
|
private fun parseResponse(raw: String): ProfileModelResponse {
|
||||||
return parseObject(raw)
|
return parseObject(raw)
|
||||||
}
|
}
|
||||||
@@ -186,4 +174,9 @@ class ProfileModelClient(
|
|||||||
companion object {
|
companion object {
|
||||||
private val THINK_REGEX = Regex("<think>[\\s\\S]*?</think>")
|
private val THINK_REGEX = Regex("<think>[\\s\\S]*?</think>")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private data class CompletedProfileResponse(
|
||||||
|
val content: String,
|
||||||
|
val usage: ProfileTokenUsage,
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,18 +1,33 @@
|
|||||||
package top.jie65535.mirai.llm
|
package top.jie65535.mirai.llm
|
||||||
|
|
||||||
|
import com.aallam.openai.api.chat.ChatCompletionRequest
|
||||||
|
import com.aallam.openai.api.chat.ChatMessage
|
||||||
|
import com.aallam.openai.api.model.ModelId
|
||||||
|
import com.sun.net.httpserver.HttpServer
|
||||||
|
import io.ktor.utils.io.ByteReadChannel
|
||||||
import kotlinx.coroutines.CompletableDeferred
|
import kotlinx.coroutines.CompletableDeferred
|
||||||
import kotlinx.coroutines.async
|
import kotlinx.coroutines.async
|
||||||
import kotlinx.coroutines.awaitAll
|
import kotlinx.coroutines.awaitAll
|
||||||
import kotlinx.coroutines.coroutineScope
|
import kotlinx.coroutines.coroutineScope
|
||||||
|
import kotlinx.coroutines.delay
|
||||||
import kotlinx.coroutines.flow.collect
|
import kotlinx.coroutines.flow.collect
|
||||||
import kotlinx.coroutines.flow.flow
|
import kotlinx.coroutines.flow.flow
|
||||||
import kotlinx.coroutines.runBlocking
|
import kotlinx.coroutines.runBlocking
|
||||||
import kotlinx.coroutines.sync.Semaphore
|
import kotlinx.coroutines.sync.Semaphore
|
||||||
import kotlinx.coroutines.withTimeout
|
import kotlinx.coroutines.withTimeout
|
||||||
|
import kotlinx.serialization.json.Json
|
||||||
|
import kotlinx.serialization.json.jsonObject
|
||||||
|
import kotlinx.serialization.json.jsonPrimitive
|
||||||
|
import okhttp3.Protocol
|
||||||
|
import java.net.InetSocketAddress
|
||||||
import java.util.concurrent.atomic.AtomicInteger
|
import java.util.concurrent.atomic.AtomicInteger
|
||||||
import kotlin.test.Test
|
import kotlin.test.Test
|
||||||
import kotlin.test.assertEquals
|
import kotlin.test.assertEquals
|
||||||
|
import kotlin.test.assertFailsWith
|
||||||
|
import kotlin.test.assertNotNull
|
||||||
import kotlin.test.assertNull
|
import kotlin.test.assertNull
|
||||||
|
import kotlin.test.assertTrue
|
||||||
|
import kotlin.time.Duration.Companion.milliseconds
|
||||||
import kotlin.time.Duration.Companion.seconds
|
import kotlin.time.Duration.Companion.seconds
|
||||||
|
|
||||||
class ModelServiceTest {
|
class ModelServiceTest {
|
||||||
@@ -54,12 +69,169 @@ class ModelServiceTest {
|
|||||||
assertEquals(128, dispatcher.maxRequestsPerHost)
|
assertEquals(128, dispatcher.maxRequestsPerHost)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun forcesHttp11ForAllModelRequests() {
|
||||||
|
assertEquals(listOf(Protocol.HTTP_1_1), MODEL_HTTP_PROTOCOLS)
|
||||||
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
fun clampsConcurrencyToSafetyRange() {
|
fun clampsConcurrencyToSafetyRange() {
|
||||||
assertEquals(1, normalizeMaxConcurrentRequests(0))
|
assertEquals(1, normalizeMaxConcurrentRequests(0))
|
||||||
assertEquals(MAX_MODEL_CONCURRENT_REQUESTS, normalizeMaxConcurrentRequests(Int.MAX_VALUE))
|
assertEquals(MAX_MODEL_CONCURRENT_REQUESTS, normalizeMaxConcurrentRequests(Int.MAX_VALUE))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun buildsExplicitStreamingRequestBody() {
|
||||||
|
val body = service.buildRequestBody(
|
||||||
|
request = ChatCompletionRequest(
|
||||||
|
model = ModelId("test-model"),
|
||||||
|
messages = listOf(ChatMessage.User("hello")),
|
||||||
|
),
|
||||||
|
stream = true,
|
||||||
|
)
|
||||||
|
|
||||||
|
assertEquals(
|
||||||
|
true,
|
||||||
|
Json.parseToJsonElement(body).jsonObject.getValue("stream").jsonPrimitive.content.toBoolean(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun acceptsStandaloneUsageAndEmptyHeartbeatEvents() {
|
||||||
|
assertNull(service.decodeStreamChunk("{}"))
|
||||||
|
|
||||||
|
val chunk = assertNotNull(
|
||||||
|
service.decodeStreamChunk(
|
||||||
|
"""{"usage":{"prompt_tokens":10,"completion_tokens":4,"total_tokens":14}}"""
|
||||||
|
)
|
||||||
|
)
|
||||||
|
assertTrue(chunk.choices.isEmpty())
|
||||||
|
assertEquals(10, chunk.usage?.promptTokens)
|
||||||
|
assertEquals(4, chunk.usage?.completionTokens)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun reportsProviderErrorEventsInsteadOfMissingChunkFields() {
|
||||||
|
val failure = assertFailsWith<ModelStreamProtocolException> {
|
||||||
|
service.decodeStreamChunk(
|
||||||
|
"""{"error":{"message":"upstream unavailable","type":"gateway_error"}}"""
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
assertTrue(failure.message.orEmpty().contains("upstream unavailable"))
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun classifiesPolicyRejectionsAsNonRetryableSafetyFailures() {
|
||||||
|
val failure = assertFailsWith<ModelSafetyRejectionException> {
|
||||||
|
service.decodeStreamChunk(
|
||||||
|
"""{"error":{"code":"cyber_policy","message":"This content was flagged for possible cybersecurity risk.","type":"invalid_request_error"}}"""
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
assertEquals("cyber_policy", failure.errorCode)
|
||||||
|
assertEquals("invalid_request_error", failure.errorType)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun classifiesSafetyMessagesEvenWhenGatewayLabelsThemAsUpstreamErrors() {
|
||||||
|
val failure = assertFailsWith<ModelSafetyRejectionException> {
|
||||||
|
service.decodeStreamChunk(
|
||||||
|
"""{"error":{"message":"This content was flagged for possible biological risk.","type":"upstream_error"}}"""
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
assertEquals("upstream_error", failure.errorType)
|
||||||
|
assertNull(failure.errorCode)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun classifiesOtherInvalidRequestsAsNonRetryableRejections() {
|
||||||
|
val failure = assertFailsWith<ModelRequestRejectedException> {
|
||||||
|
service.decodeStreamChunk(
|
||||||
|
"""{"error":{"message":"unsupported parameter","type":"invalid_request_error"}}"""
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
assertTrue(failure !is ModelSafetyRejectionException)
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun leavesAmbiguousAccountAccessErrorsRetryableForUpstreamPoolSwitching() {
|
||||||
|
val failure = assertFailsWith<ModelStreamProtocolException> {
|
||||||
|
service.decodeStreamChunk(
|
||||||
|
"""{"error":{"message":"Account access is temporarily limited","type":"invalid_request_error"}}"""
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
assertTrue(failure.message.orEmpty().contains("temporarily limited"))
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun readsMalformedUtf8SseLinesWithReplacement() = runBlocking {
|
||||||
|
val bytes = "data: first".toByteArray() +
|
||||||
|
byteArrayOf(0xC3.toByte()) +
|
||||||
|
"\r\n\r\ndata: 第二行\n".toByteArray()
|
||||||
|
val reader = LenientUtf8LineReader(ByteReadChannel(bytes))
|
||||||
|
|
||||||
|
assertEquals("data: first\uFFFD", reader.readLine())
|
||||||
|
assertEquals("", reader.readLine())
|
||||||
|
assertEquals("data: 第二行", reader.readLine())
|
||||||
|
assertNull(reader.readLine())
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun drainsPrefetchedLinesAfterUnderlyingChannelCloses() = runBlocking {
|
||||||
|
val channel = ByteReadChannel("data: first\ndata: second\n".toByteArray())
|
||||||
|
val reader = LenientUtf8LineReader(channel)
|
||||||
|
|
||||||
|
assertEquals("data: first", reader.readLine())
|
||||||
|
assertTrue(channel.isClosedForRead)
|
||||||
|
assertEquals("data: second", reader.readLine())
|
||||||
|
assertNull(reader.readLine())
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun drainsCompleteSseResponseDeliveredBeforeFirstRead() = runBlocking {
|
||||||
|
val server = HttpServer.create(InetSocketAddress("127.0.0.1", 0), 0)
|
||||||
|
val responseBody = listOf(
|
||||||
|
"""data: {"id":"chunk-1","object":"chat.completion.chunk","created":1,"model":"test-model","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}""",
|
||||||
|
"""data: {"id":"chunk-2","object":"chat.completion.chunk","created":1,"model":"test-model","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}""",
|
||||||
|
"""data: {"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}""",
|
||||||
|
"data: [DONE]",
|
||||||
|
"",
|
||||||
|
).joinToString("\n\n").toByteArray()
|
||||||
|
server.createContext("/v1/chat/completions") { exchange ->
|
||||||
|
exchange.responseHeaders.add("Content-Type", "text/event-stream")
|
||||||
|
exchange.sendResponseHeaders(200, responseBody.size.toLong())
|
||||||
|
exchange.responseBody.use { body -> body.write(responseBody) }
|
||||||
|
}
|
||||||
|
server.start()
|
||||||
|
|
||||||
|
val localService = ModelService(
|
||||||
|
baseUrl = "http://127.0.0.1:${server.address.port}/v1/",
|
||||||
|
token = "test",
|
||||||
|
timeout = 1.seconds,
|
||||||
|
firstChunkTimeout = 1.seconds,
|
||||||
|
)
|
||||||
|
try {
|
||||||
|
val content = StringBuilder()
|
||||||
|
localService.chatCompletions(
|
||||||
|
ChatCompletionRequest(
|
||||||
|
model = ModelId("test-model"),
|
||||||
|
messages = listOf(ChatMessage.User("hello")),
|
||||||
|
)
|
||||||
|
).collect { chunk ->
|
||||||
|
chunk.choices.firstOrNull()?.delta?.content?.let(content::append)
|
||||||
|
}
|
||||||
|
|
||||||
|
assertEquals("hello world", content.toString())
|
||||||
|
} finally {
|
||||||
|
localService.httpClient.close()
|
||||||
|
server.stop(0)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
fun queuesFlowsBeforeStartingWork() = runBlocking {
|
fun queuesFlowsBeforeStartingWork() = runBlocking {
|
||||||
val active = AtomicInteger()
|
val active = AtomicInteger()
|
||||||
@@ -93,4 +265,28 @@ class ModelServiceTest {
|
|||||||
assertEquals(2, maximumActive.get())
|
assertEquals(2, maximumActive.get())
|
||||||
assertEquals(6, started.get())
|
assertEquals(6, started.get())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun convertsOwnedResponseTimeoutToRetryableException() {
|
||||||
|
assertFailsWith<ModelResponseTimeoutException> {
|
||||||
|
runBlocking {
|
||||||
|
withModelResponseTimeout(10.milliseconds, "测试响应") {
|
||||||
|
delay(1.seconds)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
fun preservesOuterCoroutineTimeoutCancellation() {
|
||||||
|
assertFailsWith<kotlinx.coroutines.TimeoutCancellationException> {
|
||||||
|
runBlocking {
|
||||||
|
withTimeout(10.milliseconds) {
|
||||||
|
withModelResponseTimeout(1.seconds, "测试响应") {
|
||||||
|
delay(1.seconds)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user