agent-writer/tools/agent/updateDB.py

from bson import ObjectId
from tools.database.mongo import mainDB
from langchain.tools import tool

@tool
def UpdateDiagnosisAndAssessmentTool(session_id: str, diagnosis_and_assessment: str):
    """
    更新诊断与资产评估报告
    Args:
        session_id: 会话id   
        diagnosis_and_assessment: 诊断与资产评估报告
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    return UpdateDiagnosisAndAssessment(session_id, diagnosis_and_assessment)

def UpdateDiagnosisAndAssessment(session_id: str, diagnosis_and_assessment: str):
    """
    更新诊断与资产评估报告
    Args:
        session_id: 会话id   
        diagnosis_and_assessment: 诊断与资产评估报告
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    script = mainDB.agent_writer_session.update_one({"_id": ObjectId(session_id)},{"$set":{"diagnosis_and_assessment":diagnosis_and_assessment}})
    return {
        "success": script.modified_count > 0,
    }

@tool
def UpdateAdaptationIdeasTool(session_id: str, adaptation_ideas: str):
    """
    更新改编思路
    Args:
        session_id: 会话id   
        adaptation_ideas: 改编思路
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    return UpdateAdaptationIdeas(session_id, adaptation_ideas)

def UpdateAdaptationIdeas(session_id: str, adaptation_ideas: str):
    """
    更新改编思路
    Args:
        session_id: 会话id   
        adaptation_ideas: 改编思路
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    script = mainDB.agent_writer_session.update_one({"_id": ObjectId(session_id)},{"$set":{"adaptation_ideas":adaptation_ideas}})
    return {
        "success": script.modified_count > 0,
    }

@tool
def UpdateScriptBibleTool(
    session_id: str, 
    core_outline:str|None = None,
    character_profile:str|None = None,
    core_event_timeline:str|None = None,
    character_list:str|None = None,
    ):
    """
    更新剧本圣经
    Args:
        session_id: 会话id   
        core_outline: 核心大纲
        character_profile: 核心人物小传
        core_event_timeline: 核心事件时间线
        character_list: 总角色表
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    return UpdateScriptBible(session_id, core_outline, character_profile, core_event_timeline, character_list)

def UpdateScriptBible(
    session_id: str, 
    core_outline:str|None = None,
    character_profile:str|None = None,
    core_event_timeline:str|None = None,
    character_list:str|None = None,
    ):
    """
    更新剧本圣经
    Args:
        session_id: 会话id   
        core_outline: 核心大纲
        character_profile: 核心人物小传
        core_event_timeline: 核心事件时间线
        character_list: 总角色表
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    if core_outline is None and character_profile is None and core_event_timeline is None and character_list is None:
        return {
            "success": False,
        }
    update_dict = {}
    if core_outline is not None:
        update_dict["script_bible.core_outline"] = core_outline
    if character_profile is not None:
        update_dict["script_bible.character_profile"] = character_profile
    if core_event_timeline is not None:
        update_dict["script_bible.core_event_timeline"] = core_event_timeline
    if character_list is not None:
        update_dict["script_bible.character_list"] = character_list
    script = mainDB.agent_writer_session.update_one({"_id": ObjectId(session_id)},{"$set":update_dict})
    return {
        "success": script.modified_count > 0,
    }

@tool
def SetTotalEpisodeNumTool(session_id: str, total_episode_num: int):
    """
    设置总集数
    Args:
        session_id: 会话id   
        total_episode_num: 总集数
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    return SetTotalEpisodeNum(session_id, total_episode_num)

def SetTotalEpisodeNum(session_id: str, total_episode_num: int):
    """
    设置总集数
    Args:
        session_id: 会话id   
        total_episode_num: 总集数
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    script = mainDB.agent_writer_session.update_one({"_id": ObjectId(session_id)},{"$set":{"total_episode_num":total_episode_num}})
    return {
        "success": script.modified_count > 0,
    }

@tool
def UpdateOneEpisodeTool(session_id: str, episode_num:int, content: str):
    """
    更新单集内容
    Args:
        session_id: 会话id   
        episode_num: 剧集编号
        content: 剧集内容
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    return UpdateOneEpisode(session_id, episode_num, content)

def UpdateOneEpisode(session_id: str, episode_num:int, content: str):
    """
    更新单集内容
    Args:
        session_id: 会话id   
        episode_num: 剧集编号
        content: 剧集内容
    Returns:
        Dict: 返回一个包含以下字段的字典：
            success (bool): 是否更新成功    
    """
    script = mainDB.agent_writer_episodes.update_one({"session_id": session_id, "episode_num": episode_num},{"$set":{"content":content}}, upsert=True)
    return {
        "success": script.modified_count > 0,
    }