From c1b46af22bf2f25dda0b20588033572222fcc0f4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BD=95=E5=98=89=E4=BC=9F?=
 <54312752+loon-hejw@users.noreply.github.com>
Date: Mon, 19 Jan 2026 15:58:24 +0800
Subject: [PATCH] Refactor response handling in websocket_wiki.py

---
 api/websocket_wiki.py | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py
index 5bd0c9ff2..27bf0d71d 100644
--- a/api/websocket_wiki.py
+++ b/api/websocket_wiki.py
@@ -578,10 +578,28 @@ async def handle_websocket_chat(websocket: WebSocket):
                 response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM)
                 # Handle streaming response from Ollama
                 async for chunk in response:
-                    text = getattr(chunk, 'response', None) or getattr(chunk, 'text', None) or str(chunk)
-                    if text and not text.startswith('model=') and not text.startswith('created_at='):
-                        text = text.replace('<think>', '').replace('</think>', '')
-                        await websocket.send_text(text)
+                    text = None
+                    if isinstance(chunk, dict):
+                        text = chunk.get("message", {}).get("content") if isinstance(chunk.get("message"), dict) else chunk.get("message")
+                    else:
+                        message = getattr(chunk, "message", None)
+                        if message is not None:
+                            if isinstance(message, dict):
+                                text = message.get("content")
+                            else:
+                                text = getattr(message, "content", None)
+
+                    if not text:
+                        text = getattr(chunk, 'response', None) or getattr(chunk, 'text', None)
+
+                    if not text and hasattr(chunk, "__dict__"):
+                        message = chunk.__dict__.get("message")
+                        if isinstance(message, dict):
+                            text = message.get("content")
+
+                    if isinstance(text, str) and text and not text.startswith('model=') and not text.startswith('created_at='):
+                        clean_text = text.replace('<think>', '').replace('</think>', '')
+                        await websocket.send_text(clean_text)
                 # Explicitly close the WebSocket connection after the response is complete
                 await websocket.close()
             elif request.provider == "openrouter":