openassistant/app/core/tools/file_search_tool.py

from typing import Type, List

from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

from app.core.tools.base_tool import BaseTool
from app.models.run import Run
from app.services.file.file import FileService
from app.services.assistant.assistant import AssistantService

# import asyncio
import nest_asyncio

# 使得异步代码可以在已运行的事件循环中嵌套
nest_asyncio.apply()


# asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
# query: str = Field(..., description="query to look up in retrieval")
class FileSearchToolInput(BaseModel):
    # query: str = Field(
    #    ...,
    #    description="query to look up in retrieval",
    # )
    query: str = Field(
        ...,
        description="query to look up in retrieval",
    )


class FileSearchTool(BaseTool):
    name: str = "file_search"
    description: str = (
        "## 工具说明：这里是用户上传的文件汇总成知识库，可以根据用户输入的问题从文件知识库中检索相关的内容。## 注意：仅允许每次调用一次。"
    )

    args_schema: Type[BaseModel] = FileSearchToolInput

    def __init__(self) -> None:
        super().__init__()
        self.__filenames = []
        self.__keys = []
        self.__dirkeys = []
        self.loop = None
        self.index = 0

    def configure(self, session: Session, run: Run, **kwargs):
        # 获取当前事件循环
        # document_id = []
        file_key = []
        # filesinfo = []
        # 后语要从知识库里选择文件，所以在openassistant的数据库里可能不存在
        for key in run.file_ids:
            if len(key) == 36:
                self.__keys.append(key)  # 添加文件id 作为检索
            else:
                file_key.append(
                    key
                )  ## assiatant的id数据，在r2r里没办法检索需要提取filekey字段

        print(
            "document_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_iddocument_id"
        )
        # print(document_id)
        print(file_key)
        files = []
        # 这种情况是uuid.ex 这种格式的在最早的时候存在的，后续要去掉
        if len(file_key) > 0:
            ## 获取文件信息
            files = FileService.get_file_list_by_ids(session=session, file_ids=file_key)
            for file in files:
                self.__keys.append(file.key)
            print(files)
        # r2r接口不提供多条件，否则上面没必要存在
        """
        if len(document_id) > 0:
            filesinfo += FileService.list_in_files(ids=document_id, offset=0, limit=100)
            # asyncio.run(
            #    FileService.list_in_files(ids=document_id, offset=0, limit=100)
            # )
            for file in filesinfo:
                self.__filenames.append(file.get("title"))
                self.__keys.append(file.get("id"))
            print(filesinfo)
        """

        # files = FileService.list_in_files(ids=run.file_ids, offset=0, limit=100)

        # 读取assistant的数据，获取文件夹的id
        db_asst = AssistantService.get_assistant_sync(
            session=session, assistant_id=run.assistant_id
        )

        if db_asst.tool_resources and "file_search" in db_asst.tool_resources:
            ##{"file_search": {"vector_store_ids": [{"file_ids": []}]}}
            asst_folder_ids = (
                db_asst.tool_resources.get("file_search")
                .get("vector_stores")[0]
                .get("folder_ids")
            )
            print(asst_folder_ids)
            # folder_fileinfo = []
            if asst_folder_ids:
                self.__dirkeys = asst_folder_ids
            """
                for fid in asst_folder_ids:
                    folder_fileinfo += FileService.list_documents(
                        id=fid, offset=0, limit=100
                    )
                    # folder_fileinfo += asyncio.run(
                    #    FileService.list_documents(id=fid, offset=0, limit=100)
                    # )
                print(folder_fileinfo)
                for file in folder_fileinfo:
                    self.__filenames.append(file.get("title"))
                    self.__keys.append(file.get("id"))
            """
        # pre-cache data to prevent thread conflicts that may occur later on.
        print(
            "---------ssssssssssss-----------------sssssssssssss---------------ssssssssssssss-------------sssssssssssss-------------ss-------"
        )
        print(self.__dirkeys)
        """
        for file in files:
            self.__filenames.append(file.filename)
            self.__keys.append(file.key)
        """
        print(self.__keys)

    # indexes: List[int],
    def run(self, query: str) -> dict:
        print(
            "file_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keysfile_keys"
        )

        print(self.__keys)
        print(self.__dirkeys)
        files = []
        if self.index == 0:
            files = FileService.search_in_files(
                query=query, file_keys=self.__keys, folder_keys=self.__dirkeys
            )
            self.index = 1
        print(files)
        return files

    def instruction_supplement(self) -> str:
        """
        为 Retrieval 提供文件选择信息，用于 llm 调用抉择
        """
        if (self.__keys and len(self.__keys) > 0) or (
            self.__dirkeys and len(self.__dirkeys) > 0
        ):
            return ""  # "## 能力限制：每次请求将会限制为最多使用三种工具，并且每种工具每次只能使用一次。这样可以保证每次操作都更加简洁和有效。"  # "## Top Important: Please use the 'file_search' tool to search for relevant content or keywords. Summarize the content or keywords and provide them for 'file_search' tool usage. Please ensure that the 'file_search' tool is used only once per request."
        else:
            return ""

        # return '## important：You can use the "retrieval" tool to search for relevant information.\n If you are asking about the content of the files, please specify any keywords, topics, or context you are looking for to help retrieve the most relevant content.'
        """
            return (
                "## 工具使用规范"
                + "可调用工具："
                + "- file_search：根据关键词或短语检索指定文件中的内容，返回匹配的文本片段的内容"
                + "**调用规则**："
                + "1. 当问题涉及以下情况时必须调用本工具："
                + "   - 询问文件/文档中的具体内容"
                + "   - 需要查找数据、条款或技术细节"
                + '   - 用户明确要求"查文件"或"搜索资料"'
                + "2. 调用时需遵循："
                + "   ```json"
                + "   {"
                + '     "action": "file_search",'
                + '     "action_input": {'
                + '       "query": "精炼后的要搜索的关键词或短语，支持精确匹配"'
                + "     }"
                + "   }"
            )
            """