Skip to content

Refactor response handling in websocket_wiki.py #459

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sng-asyncfunc merged 2 commits into from
Jan 23, 2026
Merged

Refactor response handling in websocket_wiki.py #459

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c1b46af
Refactor response handling in websocket_wiki.py
loon-hejw Jan 19, 2026
f7cd837
Merge pull request #1 from loon-hejw/loon-hejw-patch-1
loon-hejw Jan 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 22 additions & 4 deletions api/websocket_wiki.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment on lines +581 to +602
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Jan 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic for extracting text from the response chunk is quite complex and is now spread across many lines within the async for loop. This reduces the readability of the loop's main purpose.

To improve maintainability and separation of concerns, consider refactoring this extraction logic into a separate helper function, for example, _extract_text_from_chunk(chunk).

Within this new function, you could also:

  • Simplify the complex expression on line 583 into a more readable if/else block.
  • Re-evaluate the necessity of the __dict__ access on lines 595-598. Direct __dict__ access is an anti-pattern and might be redundant given the getattr call that precedes it. If it's a necessary workaround, a comment explaining why would be valuable.

Here's an example of how the refactoring could look:

from typing import Any, Optional

def _extract_text_from_chunk(chunk: Any) -> Optional[str]:
    """Extracts text content from a response chunk of unknown structure."""
    text = None
    if isinstance(chunk, dict):
        message = chunk.get("message")
        if isinstance(message, dict):
            text = message.get("content")
        else:
            text = 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)

    # This part is questionable and might be removable.
    if not text and hasattr(chunk, "__dict__"):
        message = chunk.__dict__.get("message")
        if isinstance(message, dict):
            text = message.get("content")
            
    return text

async def handle_websocket_chat(websocket: WebSocket):
    # ...
    async for chunk in response:
        text = _extract_text_from_chunk(chunk)
        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":
Expand Down