diff --git a/src/main/java/com/coco/boot/aspect/ResponseHeaderAspect.java b/src/main/java/com/coco/boot/aspect/ResponseHeaderAspect.java index 3d65edd..9bdd4d2 100644 --- a/src/main/java/com/coco/boot/aspect/ResponseHeaderAspect.java +++ b/src/main/java/com/coco/boot/aspect/ResponseHeaderAspect.java @@ -1,34 +1,34 @@ -package com.coco.boot.aspect; - -import jakarta.servlet.http.HttpServletResponse; -import org.aspectj.lang.ProceedingJoinPoint; -import org.aspectj.lang.annotation.Around; -import org.aspectj.lang.annotation.Aspect; -import org.aspectj.lang.annotation.Pointcut; -import org.springframework.stereotype.Component; -import org.springframework.web.context.request.RequestContextHolder; -import org.springframework.web.context.request.ServletRequestAttributes; - -@Aspect -@Component -public class ResponseHeaderAspect { - - @Pointcut("execution(public * *(..)) && @within(org.springframework.web.bind.annotation.RestController)") - public void controllerMethods() { - } - - @Around("controllerMethods()") - public Object apiResponseAdvice(ProceedingJoinPoint pjp) throws Throwable { - Object proceed = pjp.proceed(); - ServletRequestAttributes sra = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes(); - - if (sra != null) { - HttpServletResponse response = sra.getResponse(); - if (response != null) { - response.setHeader("Content-Type", "application/json; charset=utf-8"); - } - } - - return proceed; - } -} +//package com.coco.boot.aspect; +// +//import jakarta.servlet.http.HttpServletResponse; +//import org.aspectj.lang.ProceedingJoinPoint; +//import org.aspectj.lang.annotation.Around; +//import org.aspectj.lang.annotation.Aspect; +//import org.aspectj.lang.annotation.Pointcut; +//import org.springframework.stereotype.Component; +//import org.springframework.web.context.request.RequestContextHolder; +//import org.springframework.web.context.request.ServletRequestAttributes; +// +//@Aspect +//@Component +//public class ResponseHeaderAspect { +// +// @Pointcut("execution(public * *(..)) && @within(org.springframework.web.bind.annotation.RestController)") +// public void controllerMethods() { +// } +// +// @Around("controllerMethods()") +// public Object apiResponseAdvice(ProceedingJoinPoint pjp) throws Throwable { +// Object proceed = pjp.proceed(); +// ServletRequestAttributes sra = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes(); +// +// if (sra != null) { +// HttpServletResponse response = sra.getResponse(); +// if (response != null) { +// response.setHeader("Content-Type", "application/json; charset=utf-8"); +// } +// } +// +// return proceed; +// } +//} diff --git a/src/main/java/com/coco/boot/controller/CoCoPilotController.java b/src/main/java/com/coco/boot/controller/CoCoPilotController.java index 0dabc96..641338e 100644 --- a/src/main/java/com/coco/boot/controller/CoCoPilotController.java +++ b/src/main/java/com/coco/boot/controller/CoCoPilotController.java @@ -57,9 +57,9 @@ public ResponseEntity callback(@RequestParam String code, @RequestParam @RequestMapping(value = "/v1/**", method = {RequestMethod.GET, RequestMethod.POST}) - public ResponseEntity chat(@RequestBody Conversation requestBody, - @RequestHeader("Authorization") String auth, - HttpServletRequest request) { + public ResponseEntity chat(@RequestBody Conversation requestBody, + @RequestHeader("Authorization") String auth, + HttpServletRequest request) { try { System.out.println("request = " + request.getRequestURI()); return coCoPilotService.chat(requestBody, auth, request.getRequestURI()); diff --git a/src/main/java/com/coco/boot/service/CoCoPilotService.java b/src/main/java/com/coco/boot/service/CoCoPilotService.java index 2f9f8d2..717a4a1 100644 --- a/src/main/java/com/coco/boot/service/CoCoPilotService.java +++ b/src/main/java/com/coco/boot/service/CoCoPilotService.java @@ -33,7 +33,7 @@ public interface CoCoPilotService { /** * chat 接口 */ - ResponseEntity chat(@RequestBody Conversation requestBody, @RequestHeader("Authorization") String auth, String path); + ResponseEntity chat(@RequestBody Conversation requestBody, @RequestHeader("Authorization") String auth, String path); /** * 服务状态 diff --git a/src/main/java/com/coco/boot/service/impl/CoCoPilotServiceImpl.java b/src/main/java/com/coco/boot/service/impl/CoCoPilotServiceImpl.java index d662733..ff54a47 100644 --- a/src/main/java/com/coco/boot/service/impl/CoCoPilotServiceImpl.java +++ b/src/main/java/com/coco/boot/service/impl/CoCoPilotServiceImpl.java @@ -205,11 +205,15 @@ public ResponseEntity callback(String code, String state) { @Override - public ResponseEntity chat(Conversation requestBody, String auth, String path) { + public ResponseEntity chat(Conversation requestBody, String auth, String path) { JSONObject userInfo = ChatInterceptor.tl.get(); auth = auth.substring("Bearer ".length()); String userId = userInfo.getString("id"); int trustLevel = userInfo.getIntValue("trust_level"); + RSet ghuAliveKey = redissonClient.getSet(GHU_ALIVE_KEY, StringCodec.INSTANCE); + if (!ghuAliveKey.isExists()) { + return ResponseEntity.ok(NO_KEYS); + } // 根据用户信任级别限流 RRateLimiter rateLimiter = redissonClient.getRateLimiter(USER_RATE_LIMITER + userId); if (!rateLimiter.isExists()) { @@ -219,7 +223,7 @@ public ResponseEntity chat(Conversation requestBody, String auth, St String tokenKey = DigestUtils.md5DigestAsHex((auth + userId).getBytes()); if (rateLimiter.tryAcquire()) { // 调用 handleProxy 方法并获取响应 - ResponseEntity response = handleProxy(requestBody, path); + ResponseEntity response = getBaseProxyResponse(requestBody, path, ghuAliveKey); if (response.getStatusCode().is2xxSuccessful()) { //成功访问 long tokenSuccess = redissonClient.getAtomicLong(RC_TOKEN_SUCCESS_REQ + tokenKey).incrementAndGet(); @@ -230,9 +234,14 @@ public ResponseEntity chat(Conversation requestBody, String auth, St if (userSuccess > ((long) rcConfig.getUserMaxReq() * trustLevel)) { redissonClient.getBucket(RC_TEMPORARY_BAN + userId).set(true, Duration.ofHours(rcConfig.getUserMaxTime())); } - + } + MediaType contentType = response.getHeaders().getContentType(); + if (contentType.equals(MediaType.APPLICATION_JSON)) { + // 处理 application/json 类型数据 + return new ResponseEntity<>(JSON.parseObject(response.getBody().toString()),response.getStatusCode()); } return response; + } else { long l = redissonClient.getAtomicLong(RC_USER_TOKEN_LIMIT_NUM + tokenKey).incrementAndGet(); RAtomicLong userLimit = redissonClient.getAtomicLong(RC_USER_LIMIT_NUM + userId); @@ -270,22 +279,10 @@ public ServiceStatus getServiceStatus() { } - /** - * 核心代理 方法 - */ - private ResponseEntity handleProxy(Object requestBody, String path) { - // 实现 handleProxy 方法逻辑 - // 进来后再判断一次,拦截器判断通过,但万一其他线程正好用完了 - RSet ghuAliveKey = redissonClient.getSet(GHU_ALIVE_KEY, StringCodec.INSTANCE); - if (!ghuAliveKey.isExists()) { - return ResponseEntity.ok(NO_KEYS); - } - return getBaseProxyResponse(requestBody, path, ghuAliveKey); - } @NotNull - private ResponseEntity getBaseProxyResponse(Object requestBody, String path, RSet ghuAliveKey) { + private ResponseEntity getBaseProxyResponse(Object requestBody, String path, RSet ghuAliveKey) { int i = 0; while (i < 2) { String ghu = getGhu(ghuAliveKey); @@ -295,10 +292,11 @@ private ResponseEntity getBaseProxyResponse(Object requestBody, Stri log.info("{}可用令牌数量,当前选择{}", ghuAliveKey.size(), ghu); HttpHeaders headers = new HttpHeaders(); headers.setContentType(MediaType.APPLICATION_JSON); + headers.setAccept(Arrays.asList(MediaType.APPLICATION_JSON,MediaType.TEXT_EVENT_STREAM)); headers.set("Authorization", "Bearer " + ghu); StopWatch sw = new StopWatch(); sw.start("进入代理"); - ResponseEntity response = rest.postForEntity(coCoConfig.getBaseProxy() + path, new HttpEntity<>(requestBody, headers), JSONObject.class); + ResponseEntity response = rest.postForEntity(coCoConfig.getBaseProxy() + path, new HttpEntity<>(requestBody, headers), String.class); sw.stop(); log.info(sw.prettyPrint(TimeUnit.SECONDS)); if (response.getStatusCode().is2xxSuccessful()) {