130 lines
5.1 KiB
Kotlin
130 lines
5.1 KiB
Kotlin
package com.rainy.token.data.proxy
|
||
|
||
import com.rainy.token.data.repository.CredentialRepository
|
||
import com.rainy.token.domain.model.Credential
|
||
import com.rainy.token.domain.service.ServiceType
|
||
import javax.inject.Inject
|
||
import javax.inject.Singleton
|
||
import kotlinx.coroutines.Dispatchers
|
||
import kotlinx.coroutines.withContext
|
||
import okhttp3.MediaType.Companion.toMediaType
|
||
import okhttp3.OkHttpClient
|
||
import okhttp3.Request
|
||
import okhttp3.RequestBody.Companion.toRequestBody
|
||
|
||
/**
|
||
* Sub2API Chat 透传代理。
|
||
*
|
||
* Sub2API 实例本身暴露 OpenAI 兼容接口，因此这里不做协议转换：
|
||
* 取出用户保存的 Sub2ApiCredential（优先 sk- API Key，其次面板 authToken），
|
||
* 把客户端请求原样转发到 {base}/v1/chat/completions，
|
||
* 并把上游响应（含 SSE 流式内容）原样返回给本地客户端。
|
||
*
|
||
* M1a 阶段：先做整包透传（流式也先缓冲），后续由 StreamNormalizer 升级为逐块转发。
|
||
*/
|
||
@Singleton
|
||
class Sub2ApiChatProxy @Inject constructor(
|
||
private val okHttpClient: OkHttpClient,
|
||
private val credentialRepository: CredentialRepository,
|
||
) {
|
||
|
||
/** 转发 POST /v1/chat/completions。 */
|
||
suspend fun forwardChat(requestBody: String, accountId: String? = null): ProxyUpstreamResponse? {
|
||
val base = resolveBase(accountId) ?: return null
|
||
return forward(base = base, path = "/v1/chat/completions", requestBody = requestBody, accountId = accountId)
|
||
}
|
||
|
||
/** 转发流式 POST /v1/chat/completions（上游 SSE 原样转发）。 */
|
||
suspend fun openStreamingChat(
|
||
requestBody: String,
|
||
accountId: String? = null,
|
||
): ProxyUpstreamStream? = withContext(Dispatchers.IO) {
|
||
val credential = credentialRepository.get(ServiceType.SUB2API, accountId)
|
||
?: return@withContext null
|
||
if (credential !is Credential.Sub2ApiCredential) return@withContext null
|
||
val auth = resolveAuth(credential) ?: return@withContext null
|
||
val base = normalizeBase(credential.baseUrl) ?: return@withContext null
|
||
|
||
val builder = Request.Builder()
|
||
.url(base + "/v1/chat/completions")
|
||
.addHeader("Authorization", auth)
|
||
.addHeader("Content-Type", "application/json")
|
||
.post(requestBody.toRequestBody("application/json".toMediaType()))
|
||
|
||
val response = okHttpClient.newCall(builder.build()).execute()
|
||
val input = response.body?.byteStream() ?: run {
|
||
response.close()
|
||
return@withContext null
|
||
}
|
||
ProxyUpstreamStream(
|
||
status = response.code,
|
||
contentType = response.header("Content-Type") ?: "text/event-stream",
|
||
input = input,
|
||
close = { response.close() },
|
||
)
|
||
}
|
||
|
||
/** 转发 GET /v1/models。 */
|
||
suspend fun forwardModels(accountId: String? = null): ProxyUpstreamResponse? {
|
||
val base = resolveBase(accountId) ?: return null
|
||
return forward(base = base, path = "/v1/models", requestBody = null, accountId = accountId)
|
||
}
|
||
|
||
private suspend fun resolveBase(accountId: String?): String? {
|
||
val credential = credentialRepository.get(ServiceType.SUB2API, accountId) ?: return null
|
||
if (credential !is Credential.Sub2ApiCredential) return null
|
||
if (resolveAuth(credential) == null) return null
|
||
return normalizeBase(credential.baseUrl)
|
||
}
|
||
|
||
private suspend fun forward(
|
||
base: String,
|
||
path: String,
|
||
requestBody: String?,
|
||
accountId: String?,
|
||
): ProxyUpstreamResponse? = withContext(Dispatchers.IO) {
|
||
val credential = credentialRepository.get(ServiceType.SUB2API, accountId)
|
||
?: return@withContext null
|
||
if (credential !is Credential.Sub2ApiCredential) return@withContext null
|
||
val auth = resolveAuth(credential) ?: return@withContext null
|
||
|
||
val builder = Request.Builder()
|
||
.url(base + path)
|
||
.addHeader("Authorization", auth)
|
||
val body = requestBody?.takeIf { it.isNotBlank() }
|
||
if (body != null) {
|
||
builder
|
||
.addHeader("Content-Type", "application/json")
|
||
.post(body.toRequestBody("application/json".toMediaType()))
|
||
}
|
||
|
||
val response = okHttpClient.newCall(builder.build()).execute()
|
||
val bytes = response.body?.bytes() ?: ByteArray(0)
|
||
val contentType = response.header("Content-Type") ?: "application/json"
|
||
val status = response.code
|
||
response.close()
|
||
ProxyUpstreamResponse(status, contentType, bytes)
|
||
}
|
||
|
||
private fun resolveAuth(credential: Credential.Sub2ApiCredential): String? {
|
||
credential.apiKey?.trim()?.takeIf { it.isNotBlank() }?.let { return "Bearer $it" }
|
||
credential.authToken?.trim()?.takeIf { it.isNotBlank() }?.let { return "Bearer $it" }
|
||
return null
|
||
}
|
||
|
||
private fun normalizeBase(raw: String): String? {
|
||
var s = raw.trim()
|
||
while (s.endsWith("/")) s = s.dropLast(1)
|
||
return s.takeIf { it.isNotBlank() }
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 上游 HTTP 响应（透传用）。
|
||
*/
|
||
data class ProxyUpstreamResponse(
|
||
val status: Int,
|
||
val contentType: String,
|
||
val body: ByteArray,
|
||
)
|