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

feat: add debug logging support for streaming models

Browse Source
This commit is contained in:
Makonike
2025-07-09 15:29:03 +08:00
parent 2d8bee0d6d
commit 85e1e2d4ee
1 changed files with 23 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -15,6 +15,23 @@ import json
OPENAI_RETRIES = 5
class MockResponse:
"""Mock response object for streaming models to enable consistent logging."""
def __init__(self, resp, finish_reason):
self._data = {
"choices": [
{
"message": {"content": resp},
"finish_reason": finish_reason
}
]
}
def dict(self):
return self._data
class LiteLLMAIHandler(BaseAiHandler):
"""
This class handles interactions with the OpenAI API for chat completions.
@ -401,9 +418,13 @@ class LiteLLMAIHandler(BaseAiHandler):
get_logger().debug(f"\nAI response:\n{resp}")
# log the full response for debugging
if not (model in self.streaming_required_models):
if model in self.streaming_required_models:
# for streaming, we don't have the full response object, so we create a mock one
mock_response = MockResponse(resp, finish_reason)
response_log = self.prepare_logs(mock_response, system, user, resp, finish_reason)
else:
response_log = self.prepare_logs(response, system, user, resp, finish_reason)
get_logger().debug("Full_response", artifact=response_log)
get_logger().debug("Full_response", artifact=response_log)
# for CLI debugging
if get_settings().config.verbosity_level >= 2: