Spaces:
Running
Running
File size: 3,371 Bytes
372531f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
from typing import List, Type
from ..config.config import Config
def get_retriever(retriever):
"""
Gets the retriever
Args:
retriever: retriever name
Returns:
retriever: Retriever class
"""
match retriever:
case "google":
from gpt_researcher.retrievers import GoogleSearch
retriever = GoogleSearch
case "searx":
from gpt_researcher.retrievers import SearxSearch
retriever = SearxSearch
case "searchapi":
from gpt_researcher.retrievers import SearchApiSearch
retriever = SearchApiSearch
case "serpapi":
from gpt_researcher.retrievers import SerpApiSearch
retriever = SerpApiSearch
case "serper":
from gpt_researcher.retrievers import SerperSearch
retriever = SerperSearch
case "duckduckgo":
from gpt_researcher.retrievers import Duckduckgo
retriever = Duckduckgo
case "bing":
from gpt_researcher.retrievers import BingSearch
retriever = BingSearch
case "arxiv":
from gpt_researcher.retrievers import ArxivSearch
retriever = ArxivSearch
case "tavily":
from gpt_researcher.retrievers import TavilySearch
retriever = TavilySearch
case "exa":
from gpt_researcher.retrievers import ExaSearch
retriever = ExaSearch
case "semantic_scholar":
from gpt_researcher.retrievers import SemanticScholarSearch
retriever = SemanticScholarSearch
case "pubmed_central":
from gpt_researcher.retrievers import PubMedCentralSearch
retriever = PubMedCentralSearch
case "custom":
from gpt_researcher.retrievers import CustomRetriever
retriever = CustomRetriever
case _:
retriever = None
return retriever
def get_retrievers(headers, cfg):
"""
Determine which retriever(s) to use based on headers, config, or default.
Args:
headers (dict): The headers dictionary
cfg (Config): The configuration object
Returns:
list: A list of retriever classes to be used for searching.
"""
# Check headers first for multiple retrievers
if headers.get("retrievers"):
retrievers = headers.get("retrievers").split(",")
# If not found, check headers for a single retriever
elif headers.get("retriever"):
retrievers = [headers.get("retriever")]
# If not in headers, check config for multiple retrievers
elif cfg.retrievers:
retrievers = cfg.retrievers
# If not found, check config for a single retriever
elif cfg.retriever:
retrievers = [cfg.retriever]
# If still not set, use default retriever
else:
retrievers = [get_default_retriever().__name__]
# Convert retriever names to actual retriever classes
# Use get_default_retriever() as a fallback for any invalid retriever names
return [get_retriever(r) or get_default_retriever() for r in retrievers]
def get_default_retriever(retriever):
from gpt_researcher.retrievers import TavilySearch
return TavilySearch |