import os
import re
from typing import Optional

import langchain
import paperqa
import paperscraper
from langchain import SerpAPIWrapper, OpenAI
from langchain.base_language import BaseLanguageModel
from langchain.chains import LLMChain
from langchain.tools import BaseTool
from pydantic import validator
from pypdf.errors import PdfReadError


class LitSearch(BaseTool):
    name = "LiteratureSearch"
    description = (
        "Input a specific question, returns an answer from literature search. "
        "Do not mention any specific molecule names, but use more general features to formulate your questions."
    )
    llm: BaseLanguageModel
    query_chain: Optional[LLMChain] = None
    pdir: str = "query"
    searches: int = 2
    verobse: bool = False
    docs: Optional[paperqa.Docs] = None

    @validator("query_chain", always=True)
    def init_query_chain(cls, v, values):
        if v is None:
            search_prompt = langchain.prompts.PromptTemplate(
                input_variables=["question", "count"],
                template="We want to answer the following question: {question} \n"
                "Provide {count} keyword searches (one search per line) "
                "that will find papers to help answer the question. "
                "Do not use boolean operators. "
                "Make some searches broad and some narrow. "
                "Do not use boolean operators or quotes.\n\n"
                "1. ",
            )

            llm = OpenAI(temperature=0)
            v = LLMChain(llm=llm, prompt=search_prompt)
        return v

    @validator("pdir", always=True)
    def init_pdir(cls, v):
        if not os.path.isdir(v):
            os.mkdir(v)
        return v

    def paper_search(self, search):
        try:
            return paperscraper.search_papers(
                search, pdir=self.pdir, batch_size=6, limit=4, verbose=False
            )
        except KeyError:
            return {}

    def _run(self, query: str) -> str:
        if self.verbose:
            print("\n\nChoosing search terms\n1. ", end="")
        searches = self.query_chain.run(question=query, count=self.searches)
        print("")
        queries = [s for s in searches.split("\n") if len(s) > 3]
        # remove 2., 3. from queries
        queries = [re.sub(r"^\d+\.\s*", "", q) for q in queries]
        # remove quotes
        queries = [re.sub(r"\"", "", q) for q in queries]
        papers = {}
        for q in queries:
            papers.update(self.paper_search(q))
            if self.verbose:
                print(f"retrieved {len(papers)} papers total")

        if len(papers) == 0:
            return "Not enough papers found"
        if self.docs is None:
            self.docs = paperqa.Docs(
                llm=self.llm, summary_llm="gpt-3.5-turbo", memory=True
            )
        not_loaded = 0
        for path, data in papers.items():
            try:
                self.docs.add(path, citation=data["citation"], docname=data["key"])
            except (ValueError, PdfReadError):
                not_loaded += 1

        if not_loaded:
            print(f"\nFound {len(papers.items())} papers, couldn't load {not_loaded}")
        return self.docs.query(query, length_prompt="about 100 words").answer

    async def _arun(self, query: str) -> str:
        """Use the tool asynchronously."""
        raise NotImplementedError()


class WebSearch(BaseTool):
    name = "WebSearch"
    description = (
        "Input search query, returns snippets from web search. "
        "Prefer LitSearch tool over this tool, except for simple questions."
    )
    serpapi: SerpAPIWrapper = None

    def __init__(self, search_engine="google"):
        super(WebSearch, self).__init__()

        self.serpapi = SerpAPIWrapper(
            serpapi_api_key=os.getenv("SERP_API_KEY"), search_engine=search_engine
        )

    def _run(self, query: str) -> str:
        try:
            return self.serpapi.run(query)
        except:
            return "No results, try another search"

    async def _arun(self, query: str) -> str:
        """Use the tool asynchronously."""
        raise NotImplementedError()

langchain.chains.LLMChain langchain.prompts.PromptTemplate langchain.OpenAI
[((810, 847), 'pydantic.validator', 'validator', (['"""query_chain"""'], {'always': '(True)'}), "('query_chain', always=True)\n", (819, 847), False, 'from pydantic import validator\n'), ((1585, 1615), 'pydantic.validator', 'validator', (['"""pdir"""'], {'always': '(True)'}), "('pdir', always=True)\n", (1594, 1615), False, 'from pydantic import validator\n'), ((940, 1318), 'langchain.prompts.PromptTemplate', 'langchain.prompts.PromptTemplate', ([], {'input_variables': "['question', 'count']", 'template': '"""We want to answer the following question: {question} \nProvide {count} keyword searches (one search per line) that will find papers to help answer the question. Do not use boolean operators. Make some searches broad and some narrow. Do not use boolean operators or quotes.\n\n1. """'}), '(input_variables=[\'question\', \'count\'],\n    template=\n    """We want to answer the following question: {question} \nProvide {count} keyword searches (one search per line) that will find papers to help answer the question. Do not use boolean operators. Make some searches broad and some narrow. Do not use boolean operators or quotes.\n\n1. """\n    )\n', (972, 1318), False, 'import langchain\n'), ((1484, 1505), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)'}), '(temperature=0)\n', (1490, 1505), False, 'from langchain import SerpAPIWrapper, OpenAI\n'), ((1522, 1561), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'search_prompt'}), '(llm=llm, prompt=search_prompt)\n', (1530, 1561), False, 'from langchain.chains import LLMChain\n'), ((1658, 1674), 'os.path.isdir', 'os.path.isdir', (['v'], {}), '(v)\n', (1671, 1674), False, 'import os\n'), ((1688, 1699), 'os.mkdir', 'os.mkdir', (['v'], {}), '(v)\n', (1696, 1699), False, 'import os\n'), ((1786, 1878), 'paperscraper.search_papers', 'paperscraper.search_papers', (['search'], {'pdir': 'self.pdir', 'batch_size': '(6)', 'limit': '(4)', 'verbose': '(False)'}), '(search, pdir=self.pdir, batch_size=6, limit=4,\n    verbose=False)\n', (1812, 1878), False, 'import paperscraper\n'), ((2294, 2323), 're.sub', 're.sub', (['"""^\\\\d+\\\\.\\\\s*"""', '""""""', 'q'], {}), "('^\\\\d+\\\\.\\\\s*', '', q)\n", (2300, 2323), False, 'import re\n'), ((2383, 2403), 're.sub', 're.sub', (['"""\\\\\\""""', '""""""', 'q'], {}), '(\'\\\\"\', \'\', q)\n', (2389, 2403), False, 'import re\n'), ((2737, 2805), 'paperqa.Docs', 'paperqa.Docs', ([], {'llm': 'self.llm', 'summary_llm': '"""gpt-3.5-turbo"""', 'memory': '(True)'}), "(llm=self.llm, summary_llm='gpt-3.5-turbo', memory=True)\n", (2749, 2805), False, 'import paperqa\n'), ((3810, 3835), 'os.getenv', 'os.getenv', (['"""SERP_API_KEY"""'], {}), "('SERP_API_KEY')\n", (3819, 3835), False, 'import os\n')]