Upload proxy_handler.py

proxy_handler.py

@@ -83,13 +83,11 @@ class ProxyHandler:
             body, headers, ck = await self._prep_upstream(req)
             comp_id = f"chatcmpl-{uuid.uuid4().hex[:29]}"
             think_open = False
-            phase_cur = None
-            # Buffer for what has been sent to the client to calculate deltas.
             yielded_think_buffer = "" 
-            # Authoritative state of the current full thinking content.
-            current_thinking_content = ""
+            current_raw_thinking = ""
+            answer_started = False
 
-            async def yield_delta_content(content_type: str, text: str):
+            async def yield_delta(content_type: str, text: str):
                 nonlocal think_open, yielded_think_buffer
                 if content_type == "thinking" and settings.SHOW_THINK_TAGS:
                     if not think_open:
@@ -97,15 +95,10 @@ class ProxyHandler:
                         think_open = True
                     
                     cleaned_full_text = self._clean_thinking_content(text)
-                    
-                    if cleaned_full_text.startswith(yielded_think_buffer):
-                        delta_to_send = cleaned_full_text[len(yielded_think_buffer):]
-                    else:
-                        delta_to_send = cleaned_full_text
+                    delta_to_send = cleaned_full_text[len(yielded_think_buffer):] if cleaned_full_text.startswith(yielded_think_buffer) else cleaned_full_text
                     
                     if delta_to_send:
                         yield f"data: {json.dumps({'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': delta_to_send}, 'finish_reason': None}]})}\n\n"
-                    
                     yielded_think_buffer = cleaned_full_text
 
                 elif content_type == "answer":
@@ -138,25 +131,26 @@ class ProxyHandler:
                             dat = json.loads(payload_str).get("data", {})
                         except (json.JSONDecodeError, AttributeError): continue
                         
-                        new_phase = dat.get("phase")
-                        if new_phase: phase_cur = new_phase
-                        if not phase_cur: continue
-                        
-                        # Correctly handle delta_content (append) vs edit_content (replace).
-                        if phase_cur == "thinking":
-                            if "edit_content" in dat and dat["edit_content"] is not None:
-                                current_thinking_content = dat["edit_content"]
-                            elif "delta_content" in dat and dat["delta_content"] is not None:
-                                current_thinking_content += dat["delta_content"]
-                            
-                            async for item in yield_delta_content("thinking", current_thinking_content):
-                                yield item
-
-                        elif phase_cur == "answer":
+                        if answer_started:
                             content = dat.get("delta_content") or dat.get("edit_content")
-                            if content:
-                                async for item in yield_delta_content("answer", content):
-                                    yield item
+                            if content: async for item in yield_delta("answer", content): yield item
+                            continue
+
+                        content_chunk = dat.get("edit_content") or dat.get("delta_content") or ""
+                        if dat.get("edit_content") is not None:
+                            current_raw_thinking = content_chunk
+                        else:
+                            current_raw_thinking += content_chunk
+                        
+                        match = re.search(r'(.*</details>)(.*)', current_raw_thinking, flags=re.DOTALL)
+                        if match:
+                            thinking_part, answer_part = match.groups()
+                            async for item in yield_delta("thinking", thinking_part): yield item
+                            if answer_part:
+                                async for item in yield_delta("answer", answer_part): yield item
+                            answer_started = True
+                        else:
+                            async for item in yield_delta("thinking", current_raw_thinking): yield item
         except Exception:
             logger.exception("Stream error"); raise
 
@@ -164,42 +158,48 @@ class ProxyHandler:
         ck = None
         try:
             body, headers, ck = await self._prep_upstream(req)
-            # Use a single string to hold the latest full thinking content.
             last_thinking_content = ""
             raw_answer_parts = []
-            phase_cur = None
             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_failed(ck); error_detail = await resp.text()
                     raise HTTPException(resp.status_code, f"Upstream error: {error_detail}")
                 await cookie_manager.mark_cookie_success(ck)
                 
+                phase_cur = None
+                current_raw_thinking = ""
+                answer_started = False
+                
                 async for raw in resp.aiter_text():
                     for line in raw.strip().split('\n'):
                         line = line.strip()
                         if not line.startswith('data: '): continue
-                        
                         payload_str = line[6:]
                         if payload_str == '[DONE]': break
                         try:
                             dat = json.loads(payload_str).get("data", {})
                         except (json.JSONDecodeError, AttributeError): continue
                         
-                        new_phase = dat.get("phase")
-                        if new_phase: phase_cur = new_phase
-                        if not phase_cur: continue
-
-                        if phase_cur == "thinking":
-                            if "edit_content" in dat and dat["edit_content"] is not None:
-                                last_thinking_content = dat["edit_content"]
-                            elif "delta_content" in dat and dat["delta_content"] is not None:
-                                # In non-stream, we only need the final complete version.
-                                # 'edit_content' is usually that. We can build it as a fallback.
-                                last_thinking_content += dat["delta_content"]
-                        elif phase_cur == "answer":
+                        if answer_started:
                             content = dat.get("delta_content") or dat.get("edit_content")
-                            if content:
-                                raw_answer_parts.append(content)
+                            if content: raw_answer_parts.append(content)
+                            continue
+
+                        content_chunk = dat.get("edit_content") or dat.get("delta_content")
+                        if not content_chunk: continue
+
+                        if dat.get("edit_content") is not None:
+                             current_raw_thinking = content_chunk
+                        else:
+                             current_raw_thinking += content_chunk
+                        
+                        match = re.search(r'(.*</details>)(.*)', current_raw_thinking, flags=re.DOTALL)
+                        if match:
+                            last_thinking_content, answer_part = match.groups()
+                            if answer_part: raw_answer_parts.append(answer_part)
+                            answer_started = True
+                        else:
+                            last_thinking_content = current_raw_thinking
                     else: continue
                     break