| 30 | |
| 31 | class PaperAgent: |
| 32 | def __init__( |
| 33 | self, |
| 34 | user_query: str, |
| 35 | crawler: Agent, # prompt(s) -> response(s) |
| 36 | selector: Agent, # prompt(s) -> score(s) |
| 37 | end_date: str = datetime.now().strftime("%Y%m%d"), |
| 38 | prompts_path: str = "agent_prompt.json", |
| 39 | expand_layers: int = 2, |
| 40 | search_queries: int = 5, |
| 41 | search_papers: int = 10, # per query |
| 42 | expand_papers: int = 20, # per layer |
| 43 | threads_num: int = 20, # number of threads in parallel at the same time |
| 44 | ) -> None: |
| 45 | self.user_query = user_query |
| 46 | self.crawler = crawler |
| 47 | self.selector = selector |
| 48 | self.end_date = end_date |
| 49 | with open(prompts_path, 'r', encoding='utf-8') as f: |
| 50 | self.prompts = json.load(f) |
| 51 | self.root = PaperNode({ |
| 52 | "title": user_query, |
| 53 | "extra": { |
| 54 | "touch_ids": [], |
| 55 | "crawler_recall_papers": [], |
| 56 | "recall_papers": [], |
| 57 | "recall_arxiv_ids": [], |
| 58 | "recall_abstracts": [], |
| 59 | } |
| 60 | }) |
| 61 | |
| 62 | # hyperparameters |
| 63 | self.expand_layers = expand_layers |
| 64 | self.search_queries = search_queries |
| 65 | self.search_papers = search_papers |
| 66 | self.expand_papers = expand_papers |
| 67 | self.threads_num = threads_num |
| 68 | self.papers_queue = [] |
| 69 | self.expand_start = 0 |
| 70 | self.lock = threading.Lock() |
| 71 | self.templates = { |
| 72 | "cite_template": r"~\\cite\{(.*?)\}", |
| 73 | "search_template": r"Search\](.*?)\[", |
| 74 | "expand_template": r"Expand\](.*?)\[" |
| 75 | } |
| 76 | |
| 77 | @staticmethod |
| 78 | def do_parallel(func, args, num): |