diff --git a/pr_agent/algo/__init__.py b/pr_agent/algo/__init__.py index ed9edddc..0de83548 100644 --- a/pr_agent/algo/__init__.py +++ b/pr_agent/algo/__init__.py @@ -45,6 +45,7 @@ MAX_TOKENS = { 'command-nightly': 4096, 'deepseek/deepseek-chat': 128000, # 128K, but may be limited by config.max_model_tokens 'deepseek/deepseek-reasoner': 64000, # 64K, but may be limited by config.max_model_tokens + 'openai/qwq-plus': 131072, # 131K context length, but may be limited by config.max_model_tokens 'replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1': 4096, 'meta-llama/Llama-2-7b-chat-hf': 4096, 'vertex_ai/codechat-bison': 6144, @@ -193,3 +194,8 @@ CLAUDE_EXTENDED_THINKING_MODELS = [ "anthropic/claude-3-7-sonnet-20250219", "claude-3-7-sonnet-20250219" ] + +# Models that require streaming mode +STREAMING_REQUIRED_MODELS = [ + "openai/qwq-plus" +] diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py index ec96d952..59a00045 100644 --- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py +++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py @@ -5,7 +5,7 @@ import requests from litellm import acompletion from tenacity import retry, retry_if_exception_type, retry_if_not_exception_type, stop_after_attempt -from pr_agent.algo import CLAUDE_EXTENDED_THINKING_MODELS, NO_SUPPORT_TEMPERATURE_MODELS, SUPPORT_REASONING_EFFORT_MODELS, USER_MESSAGE_ONLY_MODELS +from pr_agent.algo import CLAUDE_EXTENDED_THINKING_MODELS, NO_SUPPORT_TEMPERATURE_MODELS, SUPPORT_REASONING_EFFORT_MODELS, USER_MESSAGE_ONLY_MODELS, STREAMING_REQUIRED_MODELS from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.utils import ReasoningEffort, get_version from pr_agent.config_loader import get_settings @@ -143,6 +143,9 @@ class LiteLLMAIHandler(BaseAiHandler): # Models that support extended thinking self.claude_extended_thinking_models = CLAUDE_EXTENDED_THINKING_MODELS + # Models that require streaming + self.streaming_required_models = STREAMING_REQUIRED_MODELS + def _get_azure_ad_token(self): """ Generates an access token using Azure AD credentials from settings. @@ -370,7 +373,21 @@ class LiteLLMAIHandler(BaseAiHandler): get_logger().info(f"\nSystem prompt:\n{system}") get_logger().info(f"\nUser prompt:\n{user}") - response = await acompletion(**kwargs) + # Check if model requires streaming + if model in self.streaming_required_models: + kwargs["stream"] = True + get_logger().info(f"Using streaming mode for model {model}") + response = await acompletion(**kwargs) + # Handle streaming response + resp, finish_reason = await self._handle_streaming_response(response) + else: + response = await acompletion(**kwargs) + # Handle non-streaming response + if response is None or len(response["choices"]) == 0: + raise openai.APIError + resp = response["choices"][0]['message']['content'] + finish_reason = response["choices"][0]["finish_reason"] + except openai.RateLimitError as e: get_logger().error(f"Rate limit error during LLM inference: {e}") raise @@ -380,19 +397,43 @@ class LiteLLMAIHandler(BaseAiHandler): except Exception as e: get_logger().warning(f"Unknown error during LLM inference: {e}") raise openai.APIError from e - if response is None or len(response["choices"]) == 0: - raise openai.APIError - else: - resp = response["choices"][0]['message']['content'] - finish_reason = response["choices"][0]["finish_reason"] - get_logger().debug(f"\nAI response:\n{resp}") - # log the full response for debugging + get_logger().debug(f"\nAI response:\n{resp}") + + # log the full response for debugging + if not (model in self.streaming_required_models): response_log = self.prepare_logs(response, system, user, resp, finish_reason) get_logger().debug("Full_response", artifact=response_log) - # for CLI debugging - if get_settings().config.verbosity_level >= 2: - get_logger().info(f"\nAI response:\n{resp}") + # for CLI debugging + if get_settings().config.verbosity_level >= 2: + get_logger().info(f"\nAI response:\n{resp}") return resp, finish_reason + + async def _handle_streaming_response(self, response): + """ + Handle streaming response from acompletion and collect the full response. + + Args: + response: The streaming response object from acompletion + + Returns: + tuple: (full_response_content, finish_reason) + """ + full_response = "" + finish_reason = None + + try: + async for chunk in response: + if chunk.choices and len(chunk.choices) > 0: + delta = chunk.choices[0].delta + if hasattr(delta, 'content') and delta.content: + full_response += delta.content + if chunk.choices[0].finish_reason: + finish_reason = chunk.choices[0].finish_reason + except Exception as e: + get_logger().error(f"Error handling streaming response: {e}") + raise + + return full_response, finish_reason