refactor: refine the prompt enhancer pipeline (#426)

This commit is contained in:
DanielWalnut
2025-07-15 19:21:59 +08:00
committed by GitHub
parent 448001f532
commit b155e1eca6
2 changed files with 76 additions and 29 deletions

View File

@@ -2,6 +2,7 @@
# SPDX-License-Identifier: MIT
import logging
import re
from langchain.schema import HumanMessage
@@ -41,23 +42,38 @@ def prompt_enhancer_node(state: PromptEnhancerState):
# Get the response from the model
response = model.invoke(messages)
# Clean up the response - remove any extra formatting or comments
enhanced_prompt = response.content.strip()
# Extract content from response
response_content = response.content.strip()
logger.debug(f"Response content: {response_content}")
# Remove common prefixes that might be added by the model
prefixes_to_remove = [
"Enhanced Prompt:",
"Enhanced prompt:",
"Here's the enhanced prompt:",
"Here is the enhanced prompt:",
"**Enhanced Prompt**:",
"**Enhanced prompt**:",
]
# Try to extract content from XML tags first
xml_match = re.search(
r"<enhanced_prompt>(.*?)</enhanced_prompt>", response_content, re.DOTALL
)
for prefix in prefixes_to_remove:
if enhanced_prompt.startswith(prefix):
enhanced_prompt = enhanced_prompt[len(prefix) :].strip()
break
if xml_match:
# Extract content from XML tags and clean it up
enhanced_prompt = xml_match.group(1).strip()
logger.debug("Successfully extracted enhanced prompt from XML tags")
else:
# Fallback to original logic if no XML tags found
enhanced_prompt = response_content
logger.warning("No XML tags found in response, using fallback parsing")
# Remove common prefixes that might be added by the model
prefixes_to_remove = [
"Enhanced Prompt:",
"Enhanced prompt:",
"Here's the enhanced prompt:",
"Here is the enhanced prompt:",
"**Enhanced Prompt**:",
"**Enhanced prompt**:",
]
for prefix in prefixes_to_remove:
if enhanced_prompt.startswith(prefix):
enhanced_prompt = enhanced_prompt[len(prefix) :].strip()
break
logger.info("Prompt enhancement completed successfully")
logger.debug(f"Enhanced prompt: {enhanced_prompt}")