xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-15 18:20:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(ai_handler): remove model parameter from _get_completion and handle it within the method

Browse Source
This commit is contained in:
mrT23
2025-07-13 21:29:53 +03:00
parent 0e9cf274ef
commit 8e0c5c8784
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
pr_agent/algo/ai_handlers/litellm_ai_handler.py
View File

@ -359,7 +359,7 @@ class LiteLLMAIHandler(BaseAiHandler):
get_logger().info(f"\nUser prompt:\n{user}") get_logger().info(f"\nUser prompt:\n{user}")
# Get completion with automatic streaming detection # Get completion with automatic streaming detection
resp, finish_reason, response_obj = await self._get_completion(model, **kwargs) resp, finish_reason, response_obj = await self._get_completion(**kwargs)
except openai.RateLimitError as e: except openai.RateLimitError as e:
get_logger().error(f"Rate limit error during LLM inference: {e}") get_logger().error(f"Rate limit error during LLM inference: {e}")
@ -383,10 +383,11 @@ class LiteLLMAIHandler(BaseAiHandler):
return resp, finish_reason return resp, finish_reason
async def _get_completion(self, model, **kwargs): async def _get_completion(self, **kwargs):
""" """
Wrapper that automatically handles streaming for required models. Wrapper that automatically handles streaming for required models.
""" """
model = kwargs["model"]
if model in self.streaming_required_models: if model in self.streaming_required_models:
kwargs["stream"] = True kwargs["stream"] = True
get_logger().info(f"Using streaming mode for model {model}") get_logger().info(f"Using streaming mode for model {model}")

1
pr_agent/algo/ai_handlers/litellm_helpers.py
View File

@ -1,7 +1,6 @@
import json import json
import openai import openai
from azure.identity import ClientSecretCredential
from pr_agent.config_loader import get_settings from pr_agent.config_loader import get_settings
from pr_agent.log import get_logger from pr_agent.log import get_logger