add new functions.

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.10.10
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+WORKDIR $HOME/app
+COPY --chown=user . $HOME/app
+COPY ./requirements.txt ~/app/requirements.txt
+RUN pip install -r requirements.txt
+COPY . .
+CMD ["chainlit", "run", "cyber-supervisor-openai.py", "--port", "7860"]

auto_draft.py

@@ -1,145 +0,0 @@
-# from utils.references import References
-# from utils.prompts import generate_paper_prompts, generate_keywords_prompts, generate_experiments_prompts
-# from utils.gpt_interaction import get_responses, extract_responses, extract_keywords, extract_json
-# from utils.tex_processing import replace_title
-# from utils.figures import generate_random_figures
-# import datetime
-# import shutil
-# import time
-# import logging
-# import os
-#
-# TOTAL_TOKENS = 0
-# TOTAL_PROMPTS_TOKENS = 0
-# TOTAL_COMPLETION_TOKENS = 0
-#
-# def make_archive(source, destination):
-#     base = os.path.basename(destination)
-#     name = base.split('.')[0]
-#     format = base.split('.')[1]
-#     archive_from = os.path.dirname(source)
-#     archive_to = os.path.basename(source.strip(os.sep))
-#     shutil.make_archive(name, format, archive_from, archive_to)
-#     shutil.move('%s.%s'%(name,format), destination)
-#     return destination
-#
-#
-# def log_usage(usage, generating_target, print_out=True):
-#     global TOTAL_TOKENS
-#     global TOTAL_PROMPTS_TOKENS
-#     global TOTAL_COMPLETION_TOKENS
-#
-#     prompts_tokens = usage['prompt_tokens']
-#     completion_tokens = usage['completion_tokens']
-#     total_tokens = usage['total_tokens']
-#
-#     TOTAL_TOKENS += total_tokens
-#     TOTAL_PROMPTS_TOKENS += prompts_tokens
-#     TOTAL_COMPLETION_TOKENS += completion_tokens
-#
-#     message = f"For generating {generating_target}, {total_tokens} tokens have been used ({prompts_tokens} for prompts; {completion_tokens} for completion). " \
-#               f"{TOTAL_TOKENS} tokens have been used in total."
-#     if print_out:
-#         print(message)
-#     logging.info(message)
-#
-# def pipeline(paper, section, save_to_path, model):
-#     """
-#     The main pipeline of generating a section.
-#         1. Generate prompts.
-#         2. Get responses from AI assistant.
-#         3. Extract the section text.
-#         4. Save the text to .tex file.
-#     :return usage
-#     """
-#     print(f"Generating {section}...")
-#     prompts = generate_paper_prompts(paper, section)
-#     gpt_response, usage = get_responses(prompts, model)
-#     output = extract_responses(gpt_response)
-#     paper["body"][section] = output
-#     tex_file = save_to_path + f"{section}.tex"
-#     if section == "abstract":
-#         with open(tex_file, "w") as f:
-#             f.write(r"\begin{abstract}")
-#         with open(tex_file, "a") as f:
-#             f.write(output)
-#         with open(tex_file, "a") as f:
-#             f.write(r"\end{abstract}")
-#     else:
-#         with open(tex_file, "w") as f:
-#             f.write(f"\section{{{section}}}\n")
-#         with open(tex_file, "a") as f:
-#             f.write(output)
-#     time.sleep(5)
-#     print(f"{section} has been generated. Saved to {tex_file}.")
-#     return usage
-#
-#
-#
-# def generate_draft(title, description="", template="ICLR2022", model="gpt-4"):
-#     """
-#     The main pipeline of generating a paper.
-#         1. Copy everything to the output folder.
-#         2. Create references.
-#         3. Generate each section using `pipeline`.
-#         4. Post-processing: check common errors, fill the title, ...
-#     """
-#     paper = {}
-#     paper_body = {}
-#
-#     # Create a copy in the outputs folder.
-#     # todo: use copy_templates function instead.
-#     now = datetime.datetime.now()
-#     target_name = now.strftime("outputs_%Y%m%d_%H%M%S")
-#     source_folder = f"latex_templates/{template}"
-#     destination_folder = f"outputs/{target_name}"
-#     shutil.copytree(source_folder, destination_folder)
-#
-#     bibtex_path = destination_folder + "/ref.bib"
-#     save_to_path = destination_folder +"/"
-#     replace_title(save_to_path, title)
-#     logging.basicConfig( level=logging.INFO, filename=save_to_path+"generation.log")
-#
-#     # Generate keywords and references
-#     print("Initialize the paper information ...")
-#     prompts = generate_keywords_prompts(title, description)
-#     gpt_response, usage = get_responses(prompts, model)
-#     keywords = extract_keywords(gpt_response)
-#     log_usage(usage, "keywords")
-#     ref = References(load_papers = "") #todo: allow users to upload bibfile.
-#     ref.collect_papers(keywords, method="arxiv") #todo: add more methods to find related papers
-#     all_paper_ids = ref.to_bibtex(bibtex_path) #todo: this will used to check if all citations are in this list
-#
-#     print(f"The paper information has been initialized. References are saved to {bibtex_path}.")
-#
-#     paper["title"] = title
-#     paper["description"] = description
-#     paper["references"] = ref.to_prompts() #todo: see if this prompts can be compressed.
-#     paper["body"] = paper_body
-#     paper["bibtex"] = bibtex_path
-#
-#     print("Generating figures ...")
-#     prompts = generate_experiments_prompts(paper)
-#     gpt_response, usage = get_responses(prompts, model)
-#     list_of_methods = list(extract_json(gpt_response))
-#     log_usage(usage, "figures")
-#     generate_random_figures(list_of_methods, save_to_path + "comparison.png")
-#
-#     for section in ["introduction", "related works", "backgrounds", "methodology", "experiments", "conclusion", "abstract"]:
-#         try:
-#             usage = pipeline(paper, section, save_to_path, model=model)
-#             log_usage(usage, section)
-#         except Exception as e:
-#             print(f"Failed to generate {section} due to the error: {e}")
-#     print(f"The paper {title} has been generated. Saved to {save_to_path}.")
-#     return make_archive(destination_folder, "output.zip")
-#
-# if __name__ == "__main__":
-#     # title = "Training Adversarial Generative Neural Network with Adaptive Dropout Rate"
-#     title = "Playing Atari Game with Deep Reinforcement Learning"
-#     description = ""
-#     template = "ICLR2022"
-#     model = "gpt-4"
-#     # model = "gpt-3.5-turbo"
-#
-#     generate_draft(title, description, template, model)

chainlit.md

@@ -0,0 +1,15 @@
+# GPT赛博导师 (Cyber-Supervisor) 🚀🤖
+
+让ChatGPT为你的研究助力！提供研究课题，提供参考文献，帮助分析论文创新点. 
+
+## 使用OpenAI API来运行这个项目
+1. 在环境变量中添加`OPENAI_API_KEY`.
+2. 默认模型使用`gpt-3.5-turbo-16k`. 可以通过修改环境变量中的`DEFAULT_MODEL`来进行修改. 
+3. 在命令行中运行`chainlit run cyber-supervisor-openai.py`.
+
+## 基本原理
+目前提供了三个函数
+1. `find_research_directions`: 为你的研究课题寻找研究方向
+2. `find_references`: 为你的论文提供参考文献
+3. `judge_novelty`: 让赛博导师帮助分析你提出的想法的创新性 
+基于OpenAI API的Function Call功能, ChatGPT会自主选择调用哪一个工具.

cyber-supervisor-langchain.py

@@ -0,0 +1,48 @@
+import os
+
+from tools import FindResearchDirectionsTool, JudgeNoveltyTool, FindReferencesTool
+from langchain.chat_models import ChatOpenAI
+from langchain.agents import initialize_agent
+from langchain.agents import AgentType
+import openai
+from langchain.schema import SystemMessage
+from langchain.memory import ConversationBufferMemory
+
+openai.api_key = os.getenv("OPENAI_API_KEY")
+print(os.getenv("OPENAI_API_KEY"))
+default_model = os.getenv("DEFAULT_MODEL")
+if default_model is None:
+    default_model = "gpt-3.5-turbo-16k"
+
+import chainlit as cl
+
+agent_kwargs = {
+    "system_message": SystemMessage(content="You are a mighty cyber professor. "
+                                            "Your task is to assist your student to find an idea of research including:"
+                                            "1. Search related references."
+                                            "2. Propose potential research directions."
+                                            "3. Evaluate the novelty of any research direction."
+                                            "Follow the following instructions: "
+                                            "1. You always response in the same language as your student."
+                                            "2. Ask your student for further information if necessary to provide more assistance. ")
+}
+memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+
+@cl.langchain_factory(use_async=False)
+def main():
+    tools = [FindResearchDirectionsTool(), JudgeNoveltyTool(), FindReferencesTool()]
+    llm = ChatOpenAI(temperature=0.9, model=default_model, streaming=True)
+    open_ai_agent = initialize_agent(tools,
+                            llm,
+                            agent=AgentType.OPENAI_FUNCTIONS,
+                            verbose=True,
+                            agent_kwargs=agent_kwargs,
+                            memory=memory)
+    return open_ai_agent
+
+
+@cl.langchain_run
+async def run(agent, input_str):
+    res = await cl.make_async(agent)(input_str, callbacks=[cl.LangchainCallbackHandler()])
+    print(res)
+    await cl.Message(content=res["output"]).send()

cyber-supervisor-openai.py

@@ -0,0 +1,116 @@
+import os
+import openai
+import ast
+from tools import functions, TOOLS
+
+MAX_ITER = 5
+
+openai.api_key = os.getenv("OPENAI_API_KEY")
+default_model = os.getenv("DEFAULT_MODEL")
+if default_model is None:
+    default_model = "gpt-3.5-turbo-16k"
+
+import chainlit as cl
+
+async def process_new_delta(new_delta, openai_message, content_ui_message, function_ui_message):
+    if "role" in new_delta:
+        openai_message["role"] = new_delta["role"]
+    if "content" in new_delta:
+        new_content = new_delta.get("content") or ""
+        openai_message["content"] += new_content
+        await content_ui_message.stream_token(new_content)
+    if "function_call" in new_delta:
+        if "name" in new_delta["function_call"]:
+            openai_message["function_call"] = {
+                "name": new_delta["function_call"]["name"]}
+            await content_ui_message.send()
+            function_ui_message = cl.Message(
+                author=new_delta["function_call"]["name"],
+                content="", indent=1, language="json")
+            await function_ui_message.stream_token(new_delta["function_call"]["name"])
+
+        if "arguments" in new_delta["function_call"]:
+            if "arguments" not in openai_message["function_call"]:
+                openai_message["function_call"]["arguments"] = ""
+            openai_message["function_call"]["arguments"] += new_delta["function_call"]["arguments"]
+            await function_ui_message.stream_token(new_delta["function_call"]["arguments"])
+    return openai_message, content_ui_message, function_ui_message
+
+
+system_message = "You are a mighty cyber professor. Follow the following instructions: " \
+                "1. You always response in the same language as your student." \
+                "2. Ask your student for further information if necessary to provide more assistance. " \
+                "3. If your student asks you to do something out of your responsibility, please say no. "
+
+@cl.on_chat_start
+def start_chat():
+    cl.user_session.set(
+        "message_history",
+        [{"role": "system", "content": system_message}],
+    )
+
+
+@cl.on_message
+async def run_conversation(user_message: str):
+    message_history = cl.user_session.get("message_history")
+    message_history.append({"role": "user", "content": user_message})
+
+    cur_iter = 0
+
+    while cur_iter < MAX_ITER:
+        # OpenAI call
+        openai_message = {"role": "", "content": ""}
+        function_ui_message = None
+        content_ui_message = cl.Message(content="")
+        async for stream_resp in await openai.ChatCompletion.acreate(
+            model=default_model,
+            messages=message_history,
+            stream=True,
+            function_call="auto",
+            functions=functions,
+            temperature=0.9
+        ):
+
+            new_delta = stream_resp.choices[0]["delta"]
+            openai_message, content_ui_message, function_ui_message = await process_new_delta(
+                new_delta, openai_message, content_ui_message, function_ui_message)
+
+        message_history.append(openai_message)
+        if function_ui_message is not None:
+            await function_ui_message.send()
+
+        if stream_resp.choices[0]["finish_reason"] == "stop":
+            break
+
+        elif stream_resp.choices[0]["finish_reason"] != "function_call":
+            raise ValueError(stream_resp.choices[0]["finish_reason"])
+
+        # if code arrives here, it means there is a function call
+        function_name = openai_message.get("function_call").get("name")
+        arguments = ast.literal_eval(
+            openai_message.get("function_call").get("arguments"))
+
+        if function_name == "find_research_directions":
+            function_response = TOOLS[function_name](
+                research_field=arguments.get("research_description"),
+            )
+        else:
+            function_response = TOOLS[function_name](
+                title=arguments.get("title"),
+                contributions=arguments.get("contributions"),
+            )
+        message_history.append(
+            {
+                "role": "function",
+                "name": function_name,
+                "content": f"{function_response}",
+            }
+        )
+
+        await cl.Message(
+            author=function_name,
+            content=str(function_response),
+            language='json',
+            indent=1,
+        ).send()
+        cur_iter += 1

initialization.py

@@ -1,7 +0,0 @@
-"""Generate necessary components of prompts. """
-from utils.prompts import SYSTEM
-
-def get_keywords(model, title):
-    pass
-
-

kdb_test.py

@@ -7,14 +7,27 @@ import os
 import json
 from models import EMBEDDINGS
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-REPO_ID = os.getenv("KDB_REPO")
+# todo: 功能还没做
 
-snapshot_download(REPO_ID, repo_type="dataset", local_dir="knowledge_databases/",
-                  local_dir_use_symlinks=False, token=HF_TOKEN)
+HF_TOKEN = None # os.getenv("HF_TOKEN")
+REPO_ID = None # os.getenv("KDB_REPO")
+if HF_TOKEN is not None and REPO_ID is not None:
+    snapshot_download(REPO_ID, repo_type="dataset", local_dir="knowledge_databases/",
+                      local_dir_use_symlinks=False, token=HF_TOKEN)
 ALL_KDB = ["(None)"] + list_folders("knowledge_databases")
 
+ANNOUNCEMENT = """
+# Evaluate the quality of retrieved date from the FAISS database
 
+Use this space test the performance of some pre-constructed vector databases hosted at `shaocongma/kdb`. To use this space for your own FAISS database, follow this instruction:
+1. Duplicate this space.
+2. Add the secret key `HF_TOKEN` with your own Huggingface User Access Token. 
+3. Create a Huggingface Dataset. Put your FAISS database to it.
+4. Add the secret key `REPO_ID` as your dataset's address. 
+"""
+AUTODRAFT = """
+AutoDraft is a GPT-based project to generate an academic paper using the title and contributions. When generating specific sections, AutoDraft will query some necessary backgrounds in related fields from the pre-constructed vector database.  
+"""
 
 def query_from_kdb(input, kdb, query_counts):
     if kdb == "(None)":
@@ -37,25 +50,36 @@ def query_from_kdb(input, kdb, query_counts):
         raise RuntimeError(f"Failed to query from FAISS.")
     return domain_knowledge, ""
 
-ANNOUNCEMENT = """"""
-
 with gr.Blocks() as demo:
-    gr.HTML(ANNOUNCEMENT)
     with gr.Row():
         with gr.Column():
-            kdb_dropdown = gr.Dropdown(choices=ALL_KDB, value="(None)")
-            user_input = gr.Textbox(label="Input")
-            button_retrieval = gr.Button("Query", variant="primary")
+            gr.Markdown(ANNOUNCEMENT)
+
+            kdb_dropdown = gr.Dropdown(choices=ALL_KDB, value="(None)", label="Knowledge Databases",
+                                       info="Pre-defined knowledge databases utilized to aid in the generation of academic writing content. "
+                                            "Hosted at `shaocongma/kdb`.")
+            with gr.Tab("User's Input"):
+                user_input = gr.Textbox(label="Input", info="Input anything you like to test what will be retrived from the vector database.")
+                with gr.Row():
+                    button_clear = gr.Button("Clear")
+                    button_retrieval = gr.Button("Retrieve", variant="primary")
+            with gr.Tab("AutoDraft"):
+                gr.Markdown(AUTODRAFT)
+                title_input = gr.Textbox(label="Title")
+                contribution_input = gr.Textbox(label="Contributions", lines=5)
+                with gr.Row():
+                    button_clear_2 = gr.Button("Clear")
+                    button_retrieval_2 = gr.Button("Retrieve", variant="primary")
 
             with gr.Accordion("Advanced Setting", open=False):
-                query_counts_slider = gr.Slider(minimum=1, maximum=20, value=10, step=1,
-                                                       interactive=True, label="QUERY_COUNTS",
-                                                       info="从知识库内检索多少条内容")
+                query_counts_slider = gr.Slider(minimum=1, maximum=50, value=10, step=1,
+                                                interactive=True, label="QUERY_COUNTS",
+                                                info="How many contents will be retrieved from the vector database.")
 
         retrieval_output = gr.JSON(label="Output")
 
-
     button_retrieval.click(fn=query_from_kdb, inputs=[user_input, kdb_dropdown, query_counts_slider], outputs=[retrieval_output, user_input])
+
 demo.queue(concurrency_count=1, max_size=5, api_open=False)
 demo.launch(show_error=True)
 

tools.py

@@ -0,0 +1,206 @@
+import os
+import openai
+from utils.references import References
+from utils.gpt_interaction import GPTModel
+from utils.prompts import SYSTEM
+from langchain.tools import BaseTool
+from pydantic import BaseModel, Field
+from typing import Optional, Type
+
+MAX_TOKENS = 2048
+
+openai.api_key = os.getenv("OPENAI_API_KEY")
+default_model = os.getenv("DEFAULT_MODEL")
+if default_model is None:
+    default_model = "gpt-3.5-turbo-16k"
+llm = GPTModel(model=default_model, delay=1)
+
+paper_system_prompt =  '''You are an assistant designed to propose choices of research direction.
+The user will input questions or some keywords of a fields. You need to generate some paper titles and main contributions. Ensure follow the following instructions:
+Instruction:
+- Your response should follow the JSON format.
+- Your response should have the following structure: 
+{
+    "your suggested paper title": 
+        {
+            "summary": "an overview introducing what this paper will include",
+            "contributions": {
+                "contribution1": {"statement": "briefly describe this contribution", "reason": "reason why this contribution can make this paper outstanding"}, 
+                "contribution2": {"statement": "briefly describe this contribution", "reason": "reason why this contribution can make this paper outstanding"}, 
+                ...
+            }
+        }
+    "your suggested paper title": 
+        {
+            "summary": "an overview introducing what this paper will include",
+            "contributions": {
+                "contribution1": {"statement": "briefly describe this contribution", "reason": "reason why this contribution can make this paper outstanding"}, 
+                "contribution2": {"statement": "briefly describe this contribution", "reason": "reason why this contribution can make this paper outstanding"}, 
+                ...
+            }
+        }
+    ...
+}
+- Please list three to five suggested title and at least three contributions for each paper.
+'''
+
+
+contribution_system_prompt = '''You are an assistant designed to criticize the contributions of a paper. You will be provided Paper's Title, References and Contributions. Ensure follow the following instructions:
+Instruction:
+- Your response should follow the JSON format.
+- Your response should have the following structure: 
+{
+    "title": "the title provided by the user",
+    "comment": "your thoughts on if this title clearly reflects the key ideas of this paper and explain why"
+    "contributions": {
+        "contribution1": {"statement": "briefly describe what the contribution is", 
+                          "reason": "reason why the user claims it is a contribution", 
+                          "judge": "your thought about if this is a novel contribution and explain why", 
+                          "suggestion": "your suggestion on how to modify the research direction to enhance the novelty "},
+        "contribution2": {"statement": "briefly describe what the contribution is", 
+                          "reason": "reason why the user claims it is a contribution", 
+                          "judge": "your thought about if this is a novel contribution and explain why", 
+                          "suggestion": "your suggestion on how to modify the research direction to enhance the novelty "},
+        ...
+    }
+} 
+- You need to carefully check if the claimed contribution has been made in the provided references, which makes the contribution not novel.  
+- You also need to propose your concerns on if any of contributions could be incremental or just a mild modification on an existing work.
+'''
+
+
+def find_research_directions(research_field):
+    output, _ = llm(systems=paper_system_prompt, prompts=research_field, return_json=False)
+    return output
+
+def find_references(title, contributions):
+    max_tokens = MAX_TOKENS
+    ref = References(title=title, description=f"{contributions}")
+    keywords, _ = llm(systems=SYSTEM["keywords"], prompts=title, return_json=True)
+    keywords = {keyword: 10 for keyword in keywords}
+    ref.collect_papers(keywords)
+    ref_prompt = ref.to_prompts(max_tokens=max_tokens)
+    return ref_prompt
+
+
+def judge_novelty(title, contributions):
+    max_tokens = MAX_TOKENS
+    ref = References(title=title, description=f"{contributions}")
+    keywords, _ = llm(systems=SYSTEM["keywords"], prompts=title, return_json=True)
+    keywords = {keyword: 10 for keyword in keywords}
+    ref.collect_papers(keywords)
+    ref_prompt = ref.to_prompts(max_tokens=max_tokens)
+    prompt = f"Title: {title}\n References: {ref_prompt}\n Contributions: {contributions}"
+    output, _ = llm(systems=contribution_system_prompt, prompts=prompt, return_json=False)
+    return output
+
+
+functions = [
+    {
+        "name": "find_research_directions",
+        "description": "when your student has already shown interests in a specific topic and provided a rough description of potential contributions, help your student to dive this direction deeper",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "research_description": {
+                    "type": "string",
+                    "description": "a paragraph with details in English describing "
+                                   "(1) what is the main problem you are trying to solve "
+                                   "(2) what is the main novelty of this idea (3) how to complete this research."
+                }
+            },
+            "required": ["research_description"],
+        },
+    },
+    {
+        "name": "find_references",
+        "description": "find references for given details of a paper",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "title": {
+                    "type": "string",
+                    "description": "the title (in English) of the academic paper your student will write.",
+                },
+                "contributions": {"type": "string",
+                         "description": "a general description on the contributions of this paper in English."
+                                                "If there are multiple contributions, index them with numbers."},
+            },
+            "required": ["title", "contributions"],
+        },
+    },
+    {
+        "name": "judge_novelty",
+        "description": "evaluate the novelty of a paper given its title and main contributions",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "title": {
+                    "type": "string",
+                    "description": "the title (in English) of the academic paper your student will write.",
+                },
+                "contributions": {"type": "string",
+                         "description": "a general description on the contributions of this paper in English."
+                                                "If there are multiple contributions, index them with numbers."},
+            },
+            "required": ["title", "contributions"],
+        },
+    }
+]
+
+TOOLS = {"find_research_directions": find_research_directions, "find_references": find_references, "judge_novelty": judge_novelty}
+
+class FindResearchDirectionsCheckInput(BaseModel):
+    research_description: str = Field(..., description="a paragraph with details in English describing (1) what is the main problem you are trying to solve "
+                                                       "(2) what is the main novelty of this idea (3) how to complete this research.")
+
+class TitleDescriptionCheckInput(BaseModel):
+    title: str = Field(..., description="the title of the academic paper your student will write in English.")
+    contributions: str = Field(..., description="a general description on the contributions of this paper in English."
+                                                "If there are multiple contributions, index them with numbers.")
+
+
+class FindResearchDirectionsTool(BaseTool):
+    name = "find_research_directions"
+    description = """Useful when your student has already shown interests in a specific topic and provided a rough description of
+    potential contributions and you need to help your student to dive this direction deeper for your student. 
+
+                  """
+    def _run(self, research_description: str):
+        response = find_research_directions(research_description)
+        return response
+
+    def _arun(self, research_field: str):
+        raise NotImplementedError("This tool does not support async")
+
+    args_schema: Optional[Type[BaseModel]] = FindResearchDirectionsCheckInput
+
+
+class JudgeNoveltyTool(BaseTool):
+    name = "judge_novelty"
+    description = """Useful when you need to evaluate the novelty of your student's idea.
+
+                  """
+    def _run(self, title: str, contributions: str):
+        response = judge_novelty(title, contributions)
+        return response
+
+    def _arun(self, title: str, contributions: str):
+        raise NotImplementedError("This tool does not support async")
+
+    args_schema: Optional[Type[BaseModel]] = TitleDescriptionCheckInput
+
+class FindReferencesTool(BaseTool):
+    name = "find_references"
+    description = """Useful when you need to find references for a paper.
+
+                  """
+    def _run(self, title: str, contributions: str):
+        response = find_references(title, contributions)
+        return response
+
+    def _arun(self, title: str, contributions: str):
+        raise NotImplementedError("This tool does not support async")
+
+    args_schema: Optional[Type[BaseModel]] = TitleDescriptionCheckInput
+