agent-writer/tools/agent/queryDB.py

from typing import Annotated
from bson import ObjectId
from langchain_core.messages import ToolMessage
from langchain_core.tools import InjectedToolCallId
from langgraph.types import Command
from tools.database.mongo import mainDB
from langchain.tools import tool
import json

@tool(return_direct=True)
def QueryOriginalScript(session_id: str):
    """
    查询原始剧本内容是否存在
    Args:
        session_id: 会话id
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 原始剧本内容是否存在。
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "original_script": {"$exists": True, "$ne": ""}},{"_id":1})
    is_original_script = script is not None
    # print(f"tool_call_id {tool_call_id}")
    Command(update={
        "is_original_script": is_original_script
    })
    return {
        "exist": is_original_script,
    }
    tool_message_content = json.dumps({"exist": is_original_script})
    # return Command(update={
    #     "is_original_script": is_original_script,
    #     "messages": [
    #         ToolMessage(
    #             tool_call_id=tool_call_id,
    #             content=tool_message_content
    #         )
    #     ]
    # })

def QueryOriginalScriptContent(session_id: str):
    """
    查询原始剧本内容
    Args:
        session_id: 会话id
    Returns:
        Dict: 返回一个包含以下字段的字典：
            content (str): 原始剧本内容。
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "original_script": {"$exists": True, "$ne": ""}},{"original_script":1})    
    return {
        "content": script["original_script"] if script else "",
    }


@tool(return_direct=True)
def QueryDiagnosisAndAssessment(session_id: str):
    """
    查询诊断与资产评估报告是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 诊断与资产评估报告是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "diagnosis_and_assessment": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryDiagnosisAndAssessmentContent(session_id: str):
    """
    查询诊断与资产评估报告内容
    Args:
        session_id: 会话id
    Returns:
        Dict: 返回一个包含以下字段的字典：
            content (str): 诊断与资产评估报告内容。
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "diagnosis_and_assessment": {"$exists": True, "$ne": ""}},{"diagnosis_and_assessment":1})    
    return {
        "content": script["diagnosis_and_assessment"] if script else "",
    }

@tool(return_direct=True)
def QueryAdaptationIdeas(session_id: str):
    """
    查询改编思路是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 改编思路是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "adaptation_ideas": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryAdaptationIdeasContent(session_id: str):
    """
    查询改编思路内容
    Args:
        session_id: 会话id
    Returns:
        Dict: 返回一个包含以下字段的字典：
            content (str): 改编思路内容。
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "adaptation_ideas": {"$exists": True, "$ne": ""}},{"adaptation_ideas":1})    
    return {
        "content": script["adaptation_ideas"] if script else "",
    }

@tool(return_direct=True)
def QueryScriptBible(session_id: str):
    """
    查询剧本圣经是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible": {"$exists": True}},{"_id":1})
    return {
        "exist": script is not None,
    }

def QueryScriptBibleContent(session_id: str):
    """
    查询剧本圣经内容
    Args:
        session_id: 会话id
    Returns:
        Dict: 返回一个包含以下字段的字典：
            content (str): 剧本圣经内容。
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible": {"$exists": True}},{"script_bible":1})    
    return {
        "content": script["script_bible"] if script else {},
    }

@tool(return_direct=True)
def QueryCoreOutline(session_id: str):
    """
    查询剧本圣经中的核心大纲是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经中的核心大纲是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible.core_outline": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

@tool(return_direct=True)
def QueryCharacterProfile(session_id: str):
    """
    查询剧本圣经中的核心人物小传是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经中的核心人物小传是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible.character_profile": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

@tool(return_direct=True)
def QueryCoreEventTimeline(session_id: str):
    """
    查询剧本圣经中的重大事件时间线是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经中的重大事件时间线是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible.core_event_timeline": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

@tool(return_direct=True)
def QueryCharacterList(session_id: str):
    """
    查询剧本圣经中的总人物表是否存在
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            exist (bool): 剧本圣经中的总人物表是否存在。    
    """
    script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "script_bible.character_list": {"$exists": True, "$ne": ""}},{"_id":1})
    return {
        "exist": script is not None,
    }

@tool(return_direct=True)
def QueryEpisodeCount(session_id: str):
    """
    查询剧集创作情况
    Args:
        session_id: 会话id        
    Returns:
        Dict: 返回一个包含以下字段的字典：
            completed (int): 已完成的集数
            total (int): 总集数
    """
    total = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id)},{"total_episode_num":1})
    if total is None:
        return {
            "completed": 0,
            "total": 0,
        }
    count = mainDB.agent_writer_episodes.count_documents({"session_id": session_id, "content": {"$exists": True, "$ne": ""}})
    return {
        "completed": count,
        "total": int(total["total_episode_num"]) or 0,
    }

# def QuerySingleEpisodeContent(session_id: str):
#     """
#     查询单集完整内容
#     Args:
#         session_id: 会话id        
#     Returns:
#         Dict: 返回一个包含以下字段的字典：
#             exist (bool): 剧本圣经中的总人物表是否存在。    
#     """
#     pass