39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
from bson import ObjectId
|
|
from tools.database.mongo import mainDB
|
|
from langchain.tools import tool
|
|
|
|
@tool
|
|
def QueryOriginalScript(session_id: str, only_exist: bool = False):
|
|
"""
|
|
查询原始剧本内容或是否存在
|
|
Args:
|
|
session_id: 会话id
|
|
only_exist: 是否只查询存在的剧本
|
|
|
|
Returns:
|
|
Dict: 返回一个包含以下字段的字典：
|
|
original_script (str): 原始剧本内容。仅当 only_exist 为 False 时返回该字段。
|
|
exist (bool): 原始剧本内容是否存在。
|
|
|
|
"""
|
|
# c = mainDB.agent_writer_session.count_documents({})
|
|
# print(f"查询到的原始剧本session_id: {session_id}, only_exist: {only_exist} count:{c}")
|
|
if only_exist:
|
|
script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id), "original_script": {"$exists": True, "$ne": ""}})
|
|
# print(f"exist: {script}")
|
|
return {
|
|
"original_script": "",
|
|
"exist": script is not None,
|
|
}
|
|
else:
|
|
script = mainDB.agent_writer_session.find_one({"_id": ObjectId(session_id)}, {"original_script": 1})
|
|
original_script = ""
|
|
if script:
|
|
original_script = script["original_script"] or ''
|
|
print(f"查询到的原始剧本字符长度: {len(original_script)}")
|
|
return {
|
|
"original_script": original_script,
|
|
"exist": original_script != '',
|
|
}
|
|
|