added notebook for testing lit search

mdagent/tools/base_tools/util_tools/search_tools.py

@@ -5,6 +5,7 @@
 import paperqa
 import paperscraper
 from langchain.base_language import BaseLanguageModel
+from langchain.tools import BaseTool
 from pypdf.errors import PdfReadError
 
 
@@ -20,22 +21,22 @@ def paper_search(llm, query):
         input_variables=["question"],
         template="""
         I would like to find scholarly papers to answer
-        this question: {question}.
+        this question: {question}. Your response must be at
+        most 10 words long.
         'A search query that would bring up papers that can answer
         this question would be: '""",
     )
 
     query_chain = langchain.chains.llm.LLMChain(llm=llm, prompt=prompt)
     if not os.path.isdir("./query"):  # todo: move to ckpt
         os.mkdir("query/")
-
     search = query_chain.run(query)
     print("\nSearch:", search)
     papers = paper_scraper(search, pdir=f"query/{re.sub(' ', '', search)}")
     return papers
 
 
-def scholar2result_llm(llm, query):
+def scholar2result_llm(llm, query, k=5, max_sources=2):
     """Useful to answer questions that require
     technical knowledge. Ask a specific question."""
     papers = paper_search(llm, query)
@@ -50,18 +51,20 @@ def scholar2result_llm(llm, query):
             not_loaded += 1
 
     print(f"\nFound {len(papers.items())} papers but couldn't load {not_loaded}")
-    return docs.query(query).formatted_answer
+    answer = docs.query(query, k=k, max_sources=max_sources).formatted_answer
+    return answer
 
 
-class Scholar2ResultLLM:
+class Scholar2ResultLLM(BaseTool):
     name = "LiteratureSearch"
     description = (
-        "Useful to answer questions that require technical ",
-        "knowledge. Ask a specific question.",
+        "Useful to answer questions that require technical "
+        "knowledge. Ask a specific question."
     )
-    llm: BaseLanguageModel
+    llm: BaseLanguageModel = None
 
     def __init__(self, llm):
+        super().__init__()
         self.llm = llm
 
     def _run(self, query) -> str:

notebooks/lit_search.ipynb

@@ -0,0 +1,121 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/samcox/anaconda3/envs/mda_feb21/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "from mdagent import MDAgent"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#until we update to new version\n",
+    "import nest_asyncio\n",
+    "nest_asyncio.apply()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mda = MDAgent()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompt = \"Are there any studies that show that the use of a mask can reduce the spread of COVID-19?\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\"Masks COVID-19 transmission reduction studies\"\n",
+      "Search: \"Masks COVID-19 transmission reduction studies\"\n",
+      "\n",
+      "Found 14 papers but couldn't load 0\n",
+      "Yes, there are studies that show that the use of a mask can reduce the spread of COVID-19. The review by Howard et al. (2021) indicates that mask-wearing reduces the transmissibility of COVID-19 by limiting the spread of infected respiratory particles. This conclusion is supported by evidence from both laboratory and clinical studies."
+     ]
+    }
+   ],
+   "source": [
+    "answer = mda.run(prompt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'Yes, there are studies that show that the use of a mask can reduce the spread of COVID-19. The review by Howard et al. (2021) indicates that mask-wearing reduces the transmissibility of COVID-19 by limiting the spread of infected respiratory particles. This conclusion is supported by evidence from both laboratory and clinical studies.'"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "answer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "mdagent",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}