112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
import re
|
|
from loguru import logger
|
|
|
|
|
|
class Config:
|
|
"""Configuration constants for response cleaning"""
|
|
# Patterns for cleaning debug output
|
|
ANSI_PATTERN = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
|
|
TUI_BORDER = '┃'
|
|
SKIP_PATTERNS = ['alt+enter', 'Enter your prompt']
|
|
DEBUG_LOG_PATTERN = re.compile(r'^\d{4}/\d{2}/\d{2} \d{2}:\d{2}:\d{2} \w+ <.*?>.*$')
|
|
THINKING_SPINNER_PATTERN = re.compile(r'[⣽⢿⡿⣟⣯⣷⣾⣻] Thinking\.\.\.')
|
|
ASSISTANT_MARKER = "Assistant:"
|
|
PROMPT_INDICATOR = "Enter your prompt"
|
|
|
|
|
|
def clean_response(response: str, original_prompt: str) -> str:
|
|
"""Clean and format MCP response"""
|
|
if not response:
|
|
return ""
|
|
|
|
# Debug log the raw response
|
|
logger.debug(f"Raw response before cleaning: {response}")
|
|
|
|
# Remove ANSI escape sequences
|
|
response = Config.ANSI_PATTERN.sub('', response)
|
|
|
|
# Look for the Assistant: marker and extract content after it
|
|
if Config.ASSISTANT_MARKER in response:
|
|
parts = response.split(Config.ASSISTANT_MARKER)
|
|
if len(parts) > 1:
|
|
assistant_section = parts[-1]
|
|
|
|
# Find the end of the assistant response
|
|
if Config.PROMPT_INDICATOR in assistant_section:
|
|
assistant_response = assistant_section.split(Config.PROMPT_INDICATOR)[0]
|
|
else:
|
|
assistant_response = assistant_section
|
|
|
|
# Clean and extract the response
|
|
return clean_assistant_section(assistant_response)
|
|
|
|
# Fallback to cleaning the entire response if no Assistant: marker found
|
|
return clean_entire_response(response, original_prompt)
|
|
|
|
|
|
def clean_assistant_section(assistant_response: str) -> str:
|
|
"""Clean the assistant section"""
|
|
lines = assistant_response.split('\n')
|
|
cleaned_lines = []
|
|
|
|
for line in lines:
|
|
stripped = line.strip()
|
|
|
|
# Skip empty lines
|
|
if not stripped:
|
|
continue
|
|
|
|
# Skip debug log lines
|
|
if Config.DEBUG_LOG_PATTERN.match(line):
|
|
continue
|
|
|
|
# Skip thinking spinner lines
|
|
if Config.THINKING_SPINNER_PATTERN.search(line):
|
|
continue
|
|
|
|
# Handle TUI borders
|
|
if stripped.startswith(Config.TUI_BORDER):
|
|
content = stripped.strip(Config.TUI_BORDER).strip()
|
|
if content:
|
|
cleaned_lines.append(content)
|
|
else:
|
|
cleaned_lines.append(stripped)
|
|
|
|
return '\n'.join(cleaned_lines).strip()
|
|
|
|
|
|
def clean_entire_response(response: str, original_prompt: str) -> str:
|
|
"""Clean the entire response when no Assistant: marker is found"""
|
|
lines = response.split('\n')
|
|
cleaned_lines = []
|
|
|
|
for line in lines:
|
|
stripped = line.strip()
|
|
|
|
# Skip empty lines and original prompt
|
|
if not stripped or stripped == original_prompt:
|
|
continue
|
|
|
|
# Skip debug log lines
|
|
if Config.DEBUG_LOG_PATTERN.match(line):
|
|
continue
|
|
|
|
# Skip thinking spinner lines
|
|
if Config.THINKING_SPINNER_PATTERN.search(line):
|
|
continue
|
|
|
|
# Handle TUI decorations
|
|
if stripped.startswith(Config.TUI_BORDER):
|
|
content = stripped.strip(Config.TUI_BORDER).strip()
|
|
if content and content != original_prompt:
|
|
cleaned_lines.append(content)
|
|
continue
|
|
|
|
# Skip navigation hints
|
|
if any(pattern in line for pattern in Config.SKIP_PATTERNS):
|
|
continue
|
|
|
|
# Add non-empty, non-decoration lines
|
|
cleaned_lines.append(stripped)
|
|
|
|
return '\n'.join(cleaned_lines).strip() |