Files
deer-flow/src/prose/graph/prose_continue_node.py

26 lines
784 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.prompts.template import get_prompt_template
from src.prose.graph.state import ProseState
logger = logging.getLogger(__name__)
def prose_continue_node(state: ProseState):
logger.info("Generating prose continue content...")
model = get_llm_by_type(AGENT_LLM_MAP["prose_writer"])
prose_content = model.invoke(
[
SystemMessage(content=get_prompt_template("prose/prose_continue")),
HumanMessage(content=state["content"]),
],
)
return {"output": prose_content.content}