mirror of
https://gitee.com/wanwujie/deer-flow
synced 2026-04-09 00:34:46 +08:00
32 lines
994 B
Python
32 lines
994 B
Python
# 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.prose.graph.state import ProseState
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
prompt = """
|
|
You are an AI writing assistant that improves existing text.
|
|
- Limit your response to no more than 200 characters, but make sure to construct complete sentences.
|
|
- Use Markdown formatting when appropriate.
|
|
"""
|
|
|
|
|
|
def prose_improve_node(state: ProseState):
|
|
logger.info("Generating prose improve content...")
|
|
model = get_llm_by_type(AGENT_LLM_MAP["prose_writer"])
|
|
prose_content = model.invoke(
|
|
[
|
|
SystemMessage(content=prompt),
|
|
HumanMessage(content=f"The existing text is: {state['content']}"),
|
|
],
|
|
)
|
|
logger.info(f"prose_content: {prose_content}")
|
|
return {"output": prose_content.content}
|