Upload proxy_handler.py

proxy_handler.py

@@ -25,21 +25,17 @@ class ProxyHandler:
         )
 
     async def aclose(self):
-        """Closes the underlying httpx client."""
         if not self.client.is_closed:
             await self.client.aclose()
 
-    # ---------- text utilities ----------
+    # _balance_think_tag, _clean_thinking, _clean_answer, _serialize_msgs, _prep_upstream 保持不變
     def _balance_think_tag(self, txt: str) -> str:
         opens = len(re.findall(r"<think>", txt))
         closes = len(re.findall(r"</think>", txt))
-        if opens > closes:
-            txt += "</think>" * (opens - closes)
+        if opens > closes: txt += "</think>" * (opens - closes)
         elif closes > opens:
-            for _ in range(closes - opens):
-                txt = re.sub(r"</think>(?!.*</think>)", "", txt, 1)
+            for _ in range(closes - opens): txt = re.sub(r"</think>(?!.*</think>)", "", txt, 1)
         return txt
-
     def _clean_thinking(self, s: str) -> str:
         if not s: return s
         s = re.sub(r'<details[^>]*>.*?</details>', '', s, flags=re.DOTALL)
@@ -47,11 +43,9 @@ class ProxyHandler:
         s = re.sub(r'<[^>]+>', '', s)
         s = re.sub(r'^\s*>\s*', '', s, flags=re.MULTILINE)
         return s.strip()
-
     def _clean_answer(self, s: str) -> str:
         if not s: return s
         return re.sub(r"<details[^>]*>.*?</details>", "", s, flags=re.DOTALL)
-
     def _serialize_msgs(self, msgs) -> list:
         out = []
         for m in msgs:
@@ -60,31 +54,15 @@ class ProxyHandler:
             elif isinstance(m, dict): out.append(m)
             else: out.append({"role": getattr(m, "role", "user"), "content": getattr(m, "content", str(m))})
         return out
-
-    # ---------- upstream ----------
     async def _prep_upstream(self, req: ChatCompletionRequest) -> Tuple[Dict[str, Any], Dict[str, str], str]:
         ck = await cookie_manager.get_next_cookie()
         if not ck: raise HTTPException(503, "No available cookies")
-
         model = settings.UPSTREAM_MODEL if req.model == settings.MODEL_NAME else req.model
-        body = {
-            "stream": True, "model": model, "messages": self._serialize_msgs(req.messages),
-            "background_tasks": {"title_generation": True, "tags_generation": True}, "chat_id": str(uuid.uuid4()),
-            "features": {"image_generation": False, "code_interpreter": False, "web_search": False, "auto_web_search": False, "enable_thinking": True,},
-            "id": str(uuid.uuid4()), "mcp_servers": ["deep-web-search"],
-            "model_item": {"id": model, "name": "GLM-4.5", "owned_by": "openai"}, "params": {}, "tool_servers": [],
-            "variables": {"{{USER_NAME}}": "User", "{{USER_LOCATION}}": "Unknown", "{{CURRENT_DATETIME}}": time.strftime("%Y-%m-%d %H:%M:%S"),},
-        }
-        headers = {
-            "Content-Type": "application/json", "Authorization": f"Bearer {ck}",
-            "User-Agent": ("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36"),
-            "Accept": "application/json, text/event-stream", "Accept-Language": "zh-CN",
-            "sec-ch-ua": '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"',
-            "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', "x-fe-version": "prod-fe-1.0.53",
-            "Origin": "https://chat.z.ai", "Referer": "https://chat.z.ai/",
-        }
+        body = { "stream": True, "model": model, "messages": self._serialize_msgs(req.messages), "background_tasks": {"title_generation": True, "tags_generation": True}, "chat_id": str(uuid.uuid4()), "features": {"image_generation": False, "code_interpreter": False, "web_search": False, "auto_web_search": False, "enable_thinking": True,}, "id": str(uuid.uuid4()), "mcp_servers": ["deep-web-search"], "model_item": {"id": model, "name": "GLM-4.5", "owned_by": "openai"}, "params": {}, "tool_servers": [], "variables": {"{{USER_NAME}}": "User", "{{USER_LOCATION}}": "Unknown", "{{CURRENT_DATETIME}}": time.strftime("%Y-%m-%d %H:%M:%S"),},}
+        headers = { "Content-Type": "application/json", "Authorization": f"Bearer {ck}", "User-Agent": ("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36"), "Accept": "application/json, text/event-stream", "Accept-Language": "zh-CN", "sec-ch-ua": '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', "x-fe-version": "prod-fe-1.0.53", "Origin": "https://chat.z.ai", "Referer": "https://chat.z.ai/",}
         return body, headers, ck
 
+
     # ---------- stream ----------
     async def stream_proxy_response(self, req: ChatCompletionRequest) -> AsyncGenerator[str, None]:
         ck = None
@@ -97,73 +75,69 @@ class ProxyHandler:
             async with self.client.stream("POST", settings.UPSTREAM_URL, json=body, headers=headers) as resp:
                 if resp.status_code != 200:
                     await cookie_manager.mark_cookie_invalid(ck)
+                    # ... (error handling code remains the same)
                     err_body = await resp.aread()
                     err_msg = f"Error: {resp.status_code} - {err_body.decode(errors='ignore')}"
-                    err = {
-                        "id": comp_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": req.model,
-                        "choices": [{"index": 0, "delta": {"content": err_msg}, "finish_reason": "stop"}],
-                    }
+                    err = {"id": comp_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": req.model, "choices": [{"index": 0, "delta": {"content": err_msg}, "finish_reason": "stop"}],}
                     yield f"data: {json.dumps(err)}\n\n"; yield "data: [DONE]\n\n"; return
 
                 async for raw in resp.aiter_text():
-                    if not raw or raw.isspace(): continue
-                    for line in raw.split('\n'):
+                    for line in raw.strip().split('\n'):
                         line = line.strip()
                         if not line or not line.startswith('data: '): continue
 
-                        payload = line[6:]
-                        if payload == '[DONE]':
+                        payload_str = line[6:]
+                        if payload_str == '[DONE]':
                             if think_open:
-                                close_think_chunk = {'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': '</think>'}, 'finish_reason': None}]}
-                                yield f"data: {json.dumps(close_think_chunk)}\n\n"
-                            
-                            final_chunk = {'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]}
-                            yield f"data: {json.dumps(final_chunk)}\n\n"
+                                close_payload = {'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': '</think>'}, 'finish_reason': None}]}
+                                yield f"data: {json.dumps(close_payload)}\n\n"
+                            final_payload = {'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]}
+                            yield f"data: {json.dumps(final_payload)}\n\n"
                             yield "data: [DONE]\n\n"; return
 
                         try:
-                            parsed = json.loads(payload)
+                            parsed = json.loads(payload_str)
+                            dat = parsed.get("data", {})
                         except json.JSONDecodeError:
                             continue
-                        
-                        dat = parsed.get("data", {})
-                        delta, new_phase = dat.get("delta_content", ""), dat.get("phase")
 
+                        # --- FIX START: REVISED LOGIC ---
+                        delta = dat.get("delta_content", "")
+                        new_phase = dat.get("phase")
+
+                        # Step 1: Determine the effective phase for THIS chunk
+                        effective_phase = new_phase or phase_cur
+
+                        # Step 2: Handle phase transition side-effects (closing tags) BEFORE processing content
                         if new_phase and new_phase != phase_cur:
-                            if phase_cur == "thinking" and new_phase == "answer" and think_open and settings.SHOW_THINK_TAGS:
-                                # --- FIX START ---
-                                # 將包含反斜線的字典先創建好，再放入 f-string
-                                close_payload = {
-                                    'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()),
-                                    'model': req.model,
-                                    'choices': [{'index': 0, 'delta': {'content': '</think>\n'}, 'finish_reason': None}]
-                                }
-                                yield f"data: {json.dumps(close_payload)}\n\n"
-                                # --- FIX END ---
-                                think_open = False
+                            if phase_cur == "thinking" and new_phase == "answer":
+                                if think_open and settings.SHOW_THINK_TAGS:
+                                    close_payload = {'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': '</think>\n'}, 'finish_reason': None}]}
+                                    yield f"data: {json.dumps(close_payload)}\n\n"
+                                    think_open = False
+                            # Update phase_cur AFTER handling the transition
                             phase_cur = new_phase
                         
-                        if not delta: continue
-
+                        # Step 3: Process the delta content of THIS chunk based on its effective phase
                         text_to_yield = ""
-                        if phase_cur == "thinking":
+                        if effective_phase == "thinking":
                             if settings.SHOW_THINK_TAGS:
                                 if not think_open:
-                                    open_think_chunk = {'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': '<think>'}, 'finish_reason': None}]}
-                                    yield f"data: {json.dumps(open_think_chunk)}\n\n"
+                                    open_payload = {'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': '<think>'}, 'finish_reason': None}]}
+                                    yield f"data: {json.dumps(open_payload)}\n\n"
                                     think_open = True
                                 text_to_yield = self._clean_thinking(delta)
-                        elif phase_cur == "answer":
-                            text_to_yield = self._clean_answer(delta)
 
+                        elif effective_phase == "answer":
+                            text_to_yield = self._clean_answer(delta)
+                        
+                        # Step 4: Yield the processed content if it exists
                         if text_to_yield:
-                            out_chunk = {
-                                "id": comp_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": req.model,
-                                "choices": [{"index": 0, "delta": {"content": text_to_yield}, "finish_reason": None}],
-                            }
-                            yield f"data: {json.dumps(out_chunk)}\n\n"
-
+                            chunk_payload = {"id": comp_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": req.model, "choices": [{"index": 0, "delta": {"content": text_to_yield}, "finish_reason": None}],}
+                            yield f"data: {json.dumps(chunk_payload)}\n\n"
+                        # --- FIX END ---
         except httpx.RequestError as e:
+            # ... (error handling code remains the same)
             if ck: await cookie_manager.mark_cookie_invalid(ck)
             logger.error(f"Request error: {e}")
             err_msg = f"Connection error: {e}"
@@ -174,13 +148,13 @@ class ProxyHandler:
             err = {"id": f"chatcmpl-{uuid.uuid4().hex[:29]}", "choices": [{"delta": {"content": f"Internal error in stream"}, "finish_reason": "stop"}]}
             yield f"data: {json.dumps(err)}\n\n"; yield "data: [DONE]\n\n"
 
+
     # ---------- non-stream ----------
     async def non_stream_proxy_response(self, req: ChatCompletionRequest) -> ChatCompletionResponse:
         ck = None
         try:
             body, headers, ck = await self._prep_upstream(req)
             think_buf, answer_buf = [], []
-            
             phase_cur = None
 
             async with self.client.stream("POST", settings.UPSTREAM_URL, json=body, headers=headers) as resp:
@@ -190,30 +164,38 @@ class ProxyHandler:
                     raise HTTPException(resp.status_code, f"Upstream error: {error_detail}")
 
                 async for raw in resp.aiter_text():
-                    if not raw or raw.isspace(): continue
-                    for line in raw.split('\n'):
+                    for line in raw.strip().split('\n'):
                         line = line.strip()
                         if not line or not line.startswith('data: '): continue
-                        payload = line[6:]
-                        if payload == '[DONE]': break
                         
-                        try: parsed = json.loads(payload)
-                        except json.JSONDecodeError: continue
+                        payload_str = line[6:]
+                        if payload_str == '[DONE]': break
                         
-                        dat = parsed.get("data", {})
-                        delta, new_phase = dat.get("delta_content", ""), dat.get("phase")
+                        try:
+                            parsed = json.loads(payload_str)
+                            dat = parsed.get("data", {})
+                        except json.JSONDecodeError:
+                            continue
                         
+                        # --- FIX START: REVISED LOGIC ---
+                        delta = dat.get("delta_content", "")
+                        new_phase = dat.get("phase")
+
+                        # Step 1: Update the current phase if a new one is provided
                         if new_phase:
                             phase_cur = new_phase
                         
+                        # Step 2: Append the delta to the correct buffer based on the updated phase
                         if not delta: continue
                         
-                        if phase_cur == "thinking": 
+                        if phase_cur == "thinking":
                             think_buf.append(delta)
-                        elif phase_cur == "answer": 
+                        elif phase_cur == "answer":
                             answer_buf.append(delta)
-                else:
-                    pass
+                        # --- FIX END ---
+                    else: # continue for inner loop
+                        continue
+                    break # break for outer loop
             
             raw_answer = ''.join(answer_buf)
             ans_text = self._clean_answer(raw_answer)
@@ -230,6 +212,7 @@ class ProxyHandler:
                 choices=[{"index": 0, "message": {"role": "assistant", "content": final_content}, "finish_reason": "stop"}],
             )
         except httpx.RequestError as e:
+            # ... (error handling code remains the same)
             if ck: await cookie_manager.mark_cookie_invalid(ck)
             logger.error(f"Non-stream request error: {e}")
             raise HTTPException(502, f"Connection error: {e}")
@@ -237,6 +220,7 @@ class ProxyHandler:
             logger.exception(f"Non-stream unexpected error")
             raise HTTPException(500, "Internal server error")
 
+
     # ---------- FastAPI entry ----------
     async def handle_chat_completion(self, req: ChatCompletionRequest):
         stream = bool(req.stream) if req.stream is not None else settings.DEFAULT_STREAM