refactor(api): replace provider if-else chains in simple_chat with a ChatStreamer#542
refactor(api): replace provider if-else chains in simple_chat with a ChatStreamer#542GdoongMathew wants to merge 7 commits into
Conversation
![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 chat completion streaming logic by extracting provider-specific implementations from api/simple_chat.py into a clean, registry-based ChatStreamer class hierarchy in api/chat.py. This significantly simplifies the main API endpoint and centralizes the streaming logic. The review feedback highlights a case-sensitivity bug in the token limit error detection helper, as well as multiple potential KeyError risks across the new streamer classes when directly accessing optional configuration keys like temperature, top_p, top_k, and num_ctx from model_config.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def _is_token_limit_error(exc: Exception) -> bool: | ||
| error_message = str(exc) | ||
| return any( | ||
| k in error_message for k in ("maximum context length", "token limit", "too many tokens") | ||
| ) |
There was a problem hiding this comment.
The case-sensitive check in _is_token_limit_error will fail to detect token limit errors if the exception message contains any capital letters (e.g., 'Token limit exceeded' or 'Token Limit Exceeded'). Convert the error message to lowercase before performing the substring checks to ensure the fallback mechanism is reliably triggered.
| def _is_token_limit_error(exc: Exception) -> bool: | |
| error_message = str(exc) | |
| return any( | |
| k in error_message for k in ("maximum context length", "token limit", "too many tokens") | |
| ) | |
| def _is_token_limit_error(exc: Exception) -> bool: | |
| error_message = str(exc).lower() | |
| return any( | |
| k in error_message for k in ("maximum context length", "token limit", "too many tokens") | |
| ) |
| self.model_kwargs = { | ||
| "model": model, | ||
| "stream": True, | ||
| "options": { | ||
| "temperature": model_config["temperature"], | ||
| "top_p": model_config["top_p"], | ||
| "num_ctx": model_config["num_ctx"] | ||
| } | ||
| } |
There was a problem hiding this comment.
Directly accessing keys like temperature, top_p, and num_ctx from model_config can raise a KeyError if they are not defined in the configuration file. It is safer to dynamically build the options dictionary only with the keys that are actually present in model_config.
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| "options": { | |
| "temperature": model_config["temperature"], | |
| "top_p": model_config["top_p"], | |
| "num_ctx": model_config["num_ctx"] | |
| } | |
| } | |
| options = {} | |
| for key in ("temperature", "top_p", "num_ctx"): | |
| if key in model_config: | |
| options[key] = model_config[key] | |
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| "options": options | |
| } |
| self.model_kwargs = { | ||
| "model": model, | ||
| "stream": True, | ||
| "temperature": model_config["temperature"] | ||
| } | ||
| if "top_k" in model_config: | ||
| self.model_kwargs["top_k"] = model_config["top_k"] |
There was a problem hiding this comment.
Directly accessing model_config["temperature"] can raise a KeyError if it is not defined in the configuration. Use a safe dictionary check or fallback to ensure robustness.
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| "temperature": model_config["temperature"] | |
| } | |
| if "top_k" in model_config: | |
| self.model_kwargs["top_k"] = model_config["top_k"] | |
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| } | |
| if "temperature" in model_config: | |
| self.model_kwargs["temperature"] = model_config["temperature"] | |
| if "top_k" in model_config: | |
| self.model_kwargs["top_k"] = model_config["top_k"] |
| self.model_kwargs = { | ||
| "model": model, | ||
| "stream": True, | ||
| "temperature": model_config["temperature"] | ||
| } | ||
| # Only add top_p if it exists in the model config | ||
| if "top_p" in model_config: | ||
| self.model_kwargs["top_p"] = model_config["top_p"] |
There was a problem hiding this comment.
Directly accessing model_config["temperature"] can raise a KeyError if it is not defined in the configuration. Use a safe dictionary check or fallback to ensure robustness.
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| "temperature": model_config["temperature"] | |
| } | |
| # Only add top_p if it exists in the model config | |
| if "top_p" in model_config: | |
| self.model_kwargs["top_p"] = model_config["top_p"] | |
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| } | |
| if "temperature" in model_config: | |
| self.model_kwargs["temperature"] = model_config["temperature"] | |
| if "top_p" in model_config: | |
| self.model_kwargs["top_p"] = model_config["top_p"] |
| self.model_kwargs = { | ||
| "model": model, | ||
| "stream": True, | ||
| "temperature": model_config["temperature"], | ||
| "top_p": model_config["top_p"] | ||
| } |
There was a problem hiding this comment.
Directly accessing model_config["temperature"] and model_config["top_p"] can raise a KeyError if they are not defined in the configuration. Use safe dictionary checks to ensure robustness.
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| "temperature": model_config["temperature"], | |
| "top_p": model_config["top_p"] | |
| } | |
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| } | |
| if "temperature" in model_config: | |
| self.model_kwargs["temperature"] = model_config["temperature"] | |
| if "top_p" in model_config: | |
| self.model_kwargs["top_p"] = model_config["top_p"] |
| self.model_kwargs = { | ||
| "model": model, | ||
| "temperature": model_config["temperature"], | ||
| "top_p": model_config["top_p"] | ||
| } |
There was a problem hiding this comment.
Directly accessing model_config["temperature"] and model_config["top_p"] can raise a KeyError if they are not defined in the configuration. Use safe dictionary checks to ensure robustness.
| self.model_kwargs = { | |
| "model": model, | |
| "temperature": model_config["temperature"], | |
| "top_p": model_config["top_p"] | |
| } | |
| self.model_kwargs = { | |
| "model": model, | |
| } | |
| if "temperature" in model_config: | |
| self.model_kwargs["temperature"] = model_config["temperature"] | |
| if "top_p" in model_config: | |
| self.model_kwargs["top_p"] = model_config["top_p"] |
| self.model_kwargs = { | ||
| "model": model, | ||
| "stream": True, | ||
| "temperature": model_config["temperature"], | ||
| "top_p": model_config["top_p"], | ||
| } |
There was a problem hiding this comment.
Directly accessing model_config["temperature"] and model_config["top_p"] can raise a KeyError if they are not defined in the configuration. Use safe dictionary checks to ensure robustness.
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| "temperature": model_config["temperature"], | |
| "top_p": model_config["top_p"], | |
| } | |
| self.model_kwargs = { | |
| "model": model, | |
| "stream": True, | |
| } | |
| if "temperature" in model_config: | |
| self.model_kwargs["temperature"] = model_config["temperature"] | |
| if "top_p" in model_config: | |
| self.model_kwargs["top_p"] = model_config["top_p"] |
| self.client = genai.GenerativeModel( | ||
| model_name=model, | ||
| generation_config=GenerationConfig( | ||
| temperature=model_config["temperature"], | ||
| top_p=model_config["top_p"], | ||
| top_k=model_config["top_k"], | ||
| ) | ||
| ) |
There was a problem hiding this comment.
Directly accessing model_config["temperature"], model_config["top_p"], and model_config["top_k"] can raise a KeyError if they are not defined in the configuration. Use safe dictionary checks to build the GenerationConfig to ensure robustness.
| self.client = genai.GenerativeModel( | |
| model_name=model, | |
| generation_config=GenerationConfig( | |
| temperature=model_config["temperature"], | |
| top_p=model_config["top_p"], | |
| top_k=model_config["top_k"], | |
| ) | |
| ) | |
| generation_kwargs = {} | |
| for key in ("temperature", "top_p", "top_k"): | |
| if key in model_config: | |
| generation_kwargs[key] = model_config[key] | |
| self.client = genai.GenerativeModel( | |
| model_name=model, | |
| generation_config=GenerationConfig(**generation_kwargs) if generation_kwargs else None | |
| ) |
1 participant
Summary
Extracts provider-specific streaming into a
ChatStreamerhierarchy.Changes
api/chat.py: abstractChatStreamer+ one subclass per provider (Ollama,OpenRouter, OpenAI, Azure, Bedrock, DashScope, Google). Subclasses self-register via
__init_subclass__;ChatStreamer.create(...)returns the right one. Adding a provider= adding a class. Each exposes
respond_stream(prompt) -> AsyncIterator[str].simple_chat.py: reduced toprompt()/simplified_prompt()builders plusstream_and_fallback(), which retries once with the context-free prompt on token-limiterrors. Streamers now propagate exceptions, so the fallback works for all providers.