Refactor response handling in websocket_wiki.py#459
Conversation
Refactor response handling in websocket_wiki.py
Summary of ChangesHello @loon-hejw, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request refactors the response handling mechanism within the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the response handling for Ollama to address an XML parsing issue. The new logic is more defensive and handles various response structures. My review focuses on improving the maintainability of this new, complex logic by suggesting a refactoring into a helper function and simplifying some of the internal expressions for better readability.
| 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) |
There was a problem hiding this comment.
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/elseblock. - Re-evaluate the necessity of the
__dict__access on lines 595-598. Direct__dict__access is an anti-pattern and might be redundant given thegetattrcall 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)
# ...
- Add multi-language README support (10 languages) - Add Bedrock model support (PR AsyncFuncAI#420) - Add embedding diagnostics improvements (PR AsyncFuncAI#423) - Add websocket_wiki.py refactoring (PR AsyncFuncAI#459) - Add README screenshot emoji cleanup (PR AsyncFuncAI#449) - Update current documentation timestamp
2 participants
"No valid XML" when using ollama
#390