feat: implement enhance prompt (#294)

* feat: implement enhance prompt

* add unit test

* fix prompt

* fix: fix eslint and compiling issues

* feat: add border-beam animation

* fix: fix importing issues

---------

Co-authored-by: Henry Li <henry1943@163.com>

src/config/agents.py

@@ -16,4 +16,5 @@ AGENT_LLM_MAP: dict[str, LLMType] = {
     "podcast_script_writer": "basic",
     "ppt_composer": "basic",
     "prose_writer": "basic",
+    "prompt_enhancer": "basic",
 }

src/prompt_enhancer/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: MIT
+
+"""Prompt enhancer module for improving user prompts."""

src/prompt_enhancer/graph/builder.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: MIT
+
+from langgraph.graph import StateGraph
+
+from src.prompt_enhancer.graph.enhancer_node import prompt_enhancer_node
+from src.prompt_enhancer.graph.state import PromptEnhancerState
+
+
+def build_graph():
+    """Build and return the prompt enhancer workflow graph."""
+    # Build state graph
+    builder = StateGraph(PromptEnhancerState)
+
+    # Add the enhancer node
+    builder.add_node("enhancer", prompt_enhancer_node)
+
+    # Set entry point
+    builder.set_entry_point("enhancer")
+
+    # Set finish point
+    builder.set_finish_point("enhancer")
+
+    # Compile and return the graph
+    return builder.compile()

src/prompt_enhancer/graph/enhancer_node.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: MIT
+
+import logging
+
+from langchain.schema import HumanMessage, SystemMessage
+
+from src.config.agents import AGENT_LLM_MAP
+from src.llms.llm import get_llm_by_type
+from src.prompts.template import env, apply_prompt_template
+from src.prompt_enhancer.graph.state import PromptEnhancerState
+
+logger = logging.getLogger(__name__)
+
+
+def prompt_enhancer_node(state: PromptEnhancerState):
+    """Node that enhances user prompts using AI analysis."""
+    logger.info("Enhancing user prompt...")
+
+    model = get_llm_by_type(AGENT_LLM_MAP["prompt_enhancer"])
+
+    try:
+
+        # Create messages with context if provided
+        context_info = ""
+        if state.get("context"):
+            context_info = f"\n\nAdditional context: {state['context']}"
+
+        original_prompt_message = HumanMessage(
+            content=f"Please enhance this prompt:{context_info}\n\nOriginal prompt: {state['prompt']}"
+        )
+
+        messages = apply_prompt_template(
+            "prompt_enhancer/prompt_enhancer",
+            {
+                "messages": [original_prompt_message],
+                "report_style": state.get("report_style"),
+            },
+        )
+
+        # 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()
+
+        # 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}")
+        return {"output": enhanced_prompt}
+    except Exception as e:
+        logger.error(f"Error in prompt enhancement: {str(e)}")
+        return {"output": state["prompt"]}

src/prompt_enhancer/graph/state.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: MIT
+
+from typing import TypedDict, Optional
+from src.config.report_style import ReportStyle
+
+
+class PromptEnhancerState(TypedDict):
+    """State for the prompt enhancer workflow."""
+
+    prompt: str  # Original prompt to enhance
+    context: Optional[str]  # Additional context
+    report_style: Optional[ReportStyle]  # Report style preference
+    output: Optional[str]  # Enhanced prompt result

src/prompts/prompt_enhancer/prompt_enhancer.md

@@ -0,0 +1,104 @@
+---
+CURRENT_TIME: {{ CURRENT_TIME }}
+---
+
+You are an expert prompt engineer. Your task is to enhance user prompts to make them more effective, specific, and likely to produce high-quality results from AI systems.
+
+# Your Role
+- Analyze the original prompt for clarity, specificity, and completeness
+- Enhance the prompt by adding relevant details, context, and structure
+- Make the prompt more actionable and results-oriented
+- Preserve the user's original intent while improving effectiveness
+
+{% if report_style == "academic" %}
+# Enhancement Guidelines for Academic Style
+1. **Add methodological rigor**: Include research methodology, scope, and analytical framework
+2. **Specify academic structure**: Organize with clear thesis, literature review, analysis, and conclusions
+3. **Clarify scholarly expectations**: Specify citation requirements, evidence standards, and academic tone
+4. **Add theoretical context**: Include relevant theoretical frameworks and disciplinary perspectives
+5. **Ensure precision**: Use precise terminology and avoid ambiguous language
+6. **Include limitations**: Acknowledge scope limitations and potential biases
+{% elif report_style == "popular_science" %}
+# Enhancement Guidelines for Popular Science Style
+1. **Add accessibility**: Transform technical concepts into relatable analogies and examples
+2. **Improve narrative structure**: Organize as an engaging story with clear beginning, middle, and end
+3. **Clarify audience expectations**: Specify general audience level and engagement goals
+4. **Add human context**: Include real-world applications and human interest elements
+5. **Make it compelling**: Ensure the prompt guides toward fascinating and wonder-inspiring content
+6. **Include visual elements**: Suggest use of metaphors and descriptive language for complex concepts
+{% elif report_style == "news" %}
+# Enhancement Guidelines for News Style
+1. **Add journalistic rigor**: Include fact-checking requirements, source verification, and objectivity standards
+2. **Improve news structure**: Organize with inverted pyramid structure (most important information first)
+3. **Clarify reporting expectations**: Specify timeliness, accuracy, and balanced perspective requirements
+4. **Add contextual background**: Include relevant background information and broader implications
+5. **Make it newsworthy**: Ensure the prompt focuses on current relevance and public interest
+6. **Include attribution**: Specify source requirements and quote standards
+{% elif report_style == "social_media" %}
+# Enhancement Guidelines for Social Media Style
+1. **Add engagement focus**: Include attention-grabbing elements, hooks, and shareability factors
+2. **Improve platform structure**: Organize for specific platform requirements (character limits, hashtags, etc.)
+3. **Clarify audience expectations**: Specify target demographic and engagement goals
+4. **Add viral elements**: Include trending topics, relatable content, and interactive elements
+5. **Make it shareable**: Ensure the prompt guides toward content that encourages sharing and discussion
+6. **Include visual considerations**: Suggest emoji usage, formatting, and visual appeal elements
+{% else %}
+# General Enhancement Guidelines
+1. **Add specificity**: Include relevant details, scope, and constraints
+2. **Improve structure**: Organize the request logically with clear sections if needed
+3. **Clarify expectations**: Specify desired output format, length, or style
+4. **Add context**: Include background information that would help generate better results
+5. **Make it actionable**: Ensure the prompt guides toward concrete, useful outputs
+{% endif %}
+
+# Output Requirements
+- Output ONLY the enhanced prompt
+- Do NOT include any explanations, comments, or meta-text
+- Do NOT use phrases like "Enhanced Prompt:" or "Here's the enhanced version:"
+- The output should be ready to use directly as a prompt
+
+{% if report_style == "academic" %}
+# Academic Style Examples
+
+**Original**: "Write about AI"
+**Enhanced**: "Conduct a comprehensive academic analysis of artificial intelligence applications across three key sectors: healthcare, education, and business. Employ a systematic literature review methodology to examine peer-reviewed sources from the past five years. Structure your analysis with: (1) theoretical framework defining AI and its taxonomies, (2) sector-specific case studies with quantitative performance metrics, (3) critical evaluation of implementation challenges and ethical considerations, (4) comparative analysis across sectors, and (5) evidence-based recommendations for future research directions. Maintain academic rigor with proper citations, acknowledge methodological limitations, and present findings with appropriate hedging language. Target length: 3000-4000 words with APA formatting."
+
+**Original**: "Explain climate change"
+**Enhanced**: "Provide a rigorous academic examination of anthropogenic climate change, synthesizing current scientific consensus and recent research developments. Structure your analysis as follows: (1) theoretical foundations of greenhouse effect and radiative forcing mechanisms, (2) systematic review of empirical evidence from paleoclimatic, observational, and modeling studies, (3) critical analysis of attribution studies linking human activities to observed warming, (4) evaluation of climate sensitivity estimates and uncertainty ranges, (5) assessment of projected impacts under different emission scenarios, and (6) discussion of research gaps and methodological limitations. Include quantitative data, statistical significance levels, and confidence intervals where appropriate. Cite peer-reviewed sources extensively and maintain objective, third-person academic voice throughout."
+
+{% elif report_style == "popular_science" %}
+# Popular Science Style Examples
+
+**Original**: "Write about AI"
+**Enhanced**: "Tell the fascinating story of how artificial intelligence is quietly revolutionizing our daily lives in ways most people never realize. Take readers on an engaging journey through three surprising realms: the hospital where AI helps doctors spot diseases faster than ever before, the classroom where intelligent tutors adapt to each student's learning style, and the boardroom where algorithms are making million-dollar decisions. Use vivid analogies (like comparing neural networks to how our brains work) and real-world examples that readers can relate to. Include 'wow factor' moments that showcase AI's incredible capabilities, but also honest discussions about current limitations. Write with infectious enthusiasm while maintaining scientific accuracy, and conclude with exciting possibilities that await us in the near future. Aim for 1500-2000 words that feel like a captivating conversation with a brilliant friend."
+
+**Original**: "Explain climate change"
+**Enhanced**: "Craft a compelling narrative that transforms the complex science of climate change into an accessible and engaging story for curious readers. Begin with a relatable scenario (like why your hometown weather feels different than when you were a kid) and use this as a gateway to explore the fascinating science behind our changing planet. Employ vivid analogies - compare Earth's atmosphere to a blanket, greenhouse gases to invisible heat-trapping molecules, and climate feedback loops to a snowball rolling downhill. Include surprising facts and 'aha moments' that will make readers think differently about the world around them. Weave in human stories of scientists making discoveries, communities adapting to change, and innovative solutions being developed. Balance the serious implications with hope and actionable insights, concluding with empowering steps readers can take. Write with wonder and curiosity, making complex concepts feel approachable and personally relevant."
+
+{% elif report_style == "news" %}
+# News Style Examples
+
+**Original**: "Write about AI"
+**Enhanced**: "Report on the current state and immediate impact of artificial intelligence across three critical sectors: healthcare, education, and business. Lead with the most newsworthy developments and recent breakthroughs that are affecting people today. Structure using inverted pyramid format: start with key findings and immediate implications, then provide essential background context, followed by detailed analysis and expert perspectives. Include specific, verifiable data points, recent statistics, and quotes from credible sources including industry leaders, researchers, and affected stakeholders. Address both benefits and concerns with balanced reporting, fact-check all claims, and provide proper attribution for all information. Focus on timeliness and relevance to current events, highlighting what's happening now and what readers need to know. Maintain journalistic objectivity while making the significance clear to a general news audience. Target 800-1200 words following AP style guidelines."
+
+**Original**: "Explain climate change"
+**Enhanced**: "Provide comprehensive news coverage of climate change that explains the current scientific understanding and immediate implications for readers. Lead with the most recent and significant developments in climate science, policy, or impacts that are making headlines today. Structure the report with: breaking developments first, essential background for understanding the issue, current scientific consensus with specific data and timeframes, real-world impacts already being observed, policy responses and debates, and what experts say comes next. Include quotes from credible climate scientists, policy makers, and affected communities. Present information objectively while clearly communicating the scientific consensus, fact-check all claims, and provide proper source attribution. Address common misconceptions with factual corrections. Focus on what's happening now, why it matters to readers, and what they can expect in the near future. Follow journalistic standards for accuracy, balance, and timeliness."
+
+{% elif report_style == "social_media" %}
+# Social Media Style Examples
+
+**Original**: "Write about AI"
+**Enhanced**: "Create engaging social media content about AI that will stop the scroll and spark conversations! Start with an attention-grabbing hook like 'You won't believe what AI just did in hospitals this week 🤯' and structure as a compelling thread or post series. Include surprising facts, relatable examples (like AI helping doctors spot diseases or personalizing your Netflix recommendations), and interactive elements that encourage sharing and comments. Use strategic hashtags (#AI #Technology #Future), incorporate relevant emojis for visual appeal, and include questions that prompt audience engagement ('Have you noticed AI in your daily life? Drop examples below! 👇'). Make complex concepts digestible with bite-sized explanations, trending analogies, and shareable quotes. Include a clear call-to-action and optimize for the specific platform (Twitter threads, Instagram carousel, LinkedIn professional insights, or TikTok-style quick facts). Aim for high shareability with content that feels both informative and entertaining."
+
+**Original**: "Explain climate change"
+**Enhanced**: "Develop viral-worthy social media content that makes climate change accessible and shareable without being preachy. Open with a scroll-stopping hook like 'The weather app on your phone is telling a bigger story than you think 📱🌡️' and break down complex science into digestible, engaging chunks. Use relatable comparisons (Earth's fever, atmosphere as a blanket), trending formats (before/after visuals, myth-busting series, quick facts), and interactive elements (polls, questions, challenges). Include strategic hashtags (#ClimateChange #Science #Environment), eye-catching emojis, and shareable graphics or infographics. Address common questions and misconceptions with clear, factual responses. Create content that encourages positive action rather than climate anxiety, ending with empowering steps followers can take. Optimize for platform-specific features (Instagram Stories, TikTok trends, Twitter threads) and include calls-to-action that drive engagement and sharing."
+
+{% else %}
+# General Examples
+
+**Original**: "Write about AI"
+**Enhanced**: "Write a comprehensive 1000-word analysis of artificial intelligence's current applications in healthcare, education, and business. Include specific examples of AI tools being used in each sector, discuss both benefits and challenges, and provide insights into future trends. Structure the response with clear sections for each industry and conclude with key takeaways."
+
+**Original**: "Explain climate change"
+**Enhanced**: "Provide a detailed explanation of climate change suitable for a general audience. Cover the scientific mechanisms behind global warming, major causes including greenhouse gas emissions, observable effects we're seeing today, and projected future impacts. Include specific data and examples, and explain the difference between weather and climate. Organize the response with clear headings and conclude with actionable steps individuals can take."
+{% endif %}

src/server/app.py

@@ -20,11 +20,13 @@ from src.graph.builder import build_graph_with_memory
 from src.podcast.graph.builder import build_graph as build_podcast_graph
 from src.ppt.graph.builder import build_graph as build_ppt_graph
 from src.prose.graph.builder import build_graph as build_prose_graph
+from src.prompt_enhancer.graph.builder import build_graph as build_prompt_enhancer_graph
 from src.rag.builder import build_retriever
 from src.rag.retriever import Resource
 from src.server.chat_request import (
     ChatMessage,
     ChatRequest,
+    EnhancePromptRequest,
     GeneratePodcastRequest,
     GeneratePPTRequest,
     GenerateProseRequest,
@@ -300,6 +302,50 @@ async def generate_prose(request: GenerateProseRequest):
         raise HTTPException(status_code=500, detail=INTERNAL_SERVER_ERROR_DETAIL)
 
 
+@app.post("/api/prompt/enhance")
+async def enhance_prompt(request: EnhancePromptRequest):
+    try:
+        sanitized_prompt = request.prompt.replace("\r\n", "").replace("\n", "")
+        logger.info(f"Enhancing prompt: {sanitized_prompt}")
+
+        # Convert string report_style to ReportStyle enum
+        report_style = None
+        if request.report_style:
+            try:
+                # Handle both uppercase and lowercase input
+                style_mapping = {
+                    "ACADEMIC": ReportStyle.ACADEMIC,
+                    "POPULAR_SCIENCE": ReportStyle.POPULAR_SCIENCE,
+                    "NEWS": ReportStyle.NEWS,
+                    "SOCIAL_MEDIA": ReportStyle.SOCIAL_MEDIA,
+                    "academic": ReportStyle.ACADEMIC,
+                    "popular_science": ReportStyle.POPULAR_SCIENCE,
+                    "news": ReportStyle.NEWS,
+                    "social_media": ReportStyle.SOCIAL_MEDIA,
+                }
+                report_style = style_mapping.get(
+                    request.report_style, ReportStyle.ACADEMIC
+                )
+            except Exception:
+                # If invalid style, default to ACADEMIC
+                report_style = ReportStyle.ACADEMIC
+        else:
+            report_style = ReportStyle.ACADEMIC
+
+        workflow = build_prompt_enhancer_graph()
+        final_state = workflow.invoke(
+            {
+                "prompt": request.prompt,
+                "context": request.context,
+                "report_style": report_style,
+            }
+        )
+        return {"result": final_state["output"]}
+    except Exception as e:
+        logger.exception(f"Error occurred during prompt enhancement: {str(e)}")
+        raise HTTPException(status_code=500, detail=INTERNAL_SERVER_ERROR_DETAIL)
+
+
 @app.post("/api/mcp/server/metadata", response_model=MCPServerMetadataResponse)
 async def mcp_server_metadata(request: MCPServerMetadataRequest):
     """Get information about an MCP server."""

src/server/chat_request.py

@@ -94,3 +94,13 @@ class GenerateProseRequest(BaseModel):
     command: Optional[str] = Field(
         "", description="The user custom command of the prose writer"
     )
+
+
+class EnhancePromptRequest(BaseModel):
+    prompt: str = Field(..., description="The original prompt to enhance")
+    context: Optional[str] = Field(
+        "", description="Additional context about the intended use"
+    )
+    report_style: Optional[str] = Field(
+        "academic", description="The style of the report"
+    )