H commited on
Commit
3154d16
·
1 Parent(s): c499f9e

Add component arxiv (#1587)

Browse files

### What problem does this PR solve?


### Type of change

- [x] New Feature (non-breaking change which adds functionality)

graph/component/__init__.py CHANGED
@@ -13,6 +13,7 @@ from .baidu import Baidu, BaiduParam
13
  from .duckduckgo import DuckDuckGo, DuckDuckGoParam
14
  from .wikipedia import Wikipedia, WikipediaParam
15
  from .pubmed import PubMed, PubMedParam
 
16
 
17
 
18
  def component_class(class_name):
 
13
  from .duckduckgo import DuckDuckGo, DuckDuckGoParam
14
  from .wikipedia import Wikipedia, WikipediaParam
15
  from .pubmed import PubMed, PubMedParam
16
+ from .arxiv import ArXiv, ArXivParam
17
 
18
 
19
  def component_class(class_name):
graph/component/arxiv.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #
2
+ # Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ #
16
+ import random
17
+ from abc import ABC
18
+ from functools import partial
19
+ import arxiv
20
+ import pandas as pd
21
+ from graph.settings import DEBUG
22
+ from graph.component.base import ComponentBase, ComponentParamBase
23
+
24
+
25
+ class ArXivParam(ComponentParamBase):
26
+ """
27
+ Define the ArXiv component parameters.
28
+ """
29
+
30
+ def __init__(self):
31
+ super().__init__()
32
+ self.top_n = 6
33
+ self.sort_by = 'submittedDate'
34
+
35
+ def check(self):
36
+ self.check_positive_integer(self.top_n, "Top N")
37
+ self.check_valid_value(self.sort_by, "ArXiv Search Sort_by",
38
+ ['submittedDate', 'lastUpdatedDate', 'relevance'])
39
+
40
+
41
+ class ArXiv(ComponentBase, ABC):
42
+ component_name = "ArXiv"
43
+
44
+ def _run(self, history, **kwargs):
45
+ ans = self.get_input()
46
+ ans = " - ".join(ans["content"]) if "content" in ans else ""
47
+ if not ans:
48
+ return ArXiv.be_output("")
49
+
50
+ sort_choices = {"relevance": arxiv.SortCriterion.Relevance,
51
+ "lastUpdatedDate": arxiv.SortCriterion.LastUpdatedDate,
52
+ 'submittedDate': arxiv.SortCriterion.SubmittedDate}
53
+ arxiv_client = arxiv.Client()
54
+ search = arxiv.Search(
55
+ query=ans,
56
+ max_results=self._param.top_n,
57
+ sort_by=sort_choices[self._param.sort_by]
58
+ )
59
+ arxiv_res = [
60
+ {"content": 'Title: ' + i.title + '\nPdf_Url: <a href="' + i.pdf_url + '"></a> \nSummary: ' + i.summary} for
61
+ i in list(arxiv_client.results(search))]
62
+
63
+ if not arxiv_res:
64
+ return ArXiv.be_output("")
65
+
66
+ df = pd.DataFrame(arxiv_res)
67
+ if DEBUG: print(df, ":::::::::::::::::::::::::::::::::")
68
+ return df
requirements.txt CHANGED
@@ -1,3 +1,4 @@
 
1
  Aspose.Slides==24.2.0
2
  BCEmbedding==0.1.3
3
  Bio==1.7.1
 
1
+ arxiv==2.1.3
2
  Aspose.Slides==24.2.0
3
  BCEmbedding==0.1.3
4
  Bio==1.7.1
requirements_arm.txt CHANGED
@@ -152,3 +152,4 @@ google-generativeai==0.7.2
152
  groq==0.9.0
153
  wikipedia==1.4.0
154
  Bio==1.7.1
 
 
152
  groq==0.9.0
153
  wikipedia==1.4.0
154
  Bio==1.7.1
155
+ arxiv==2.1.3
requirements_dev.txt CHANGED
@@ -137,3 +137,4 @@ google-generativeai==0.7.2
137
  groq==0.9.0
138
  wikipedia==1.4.0
139
  Bio==1.7.1
 
 
137
  groq==0.9.0
138
  wikipedia==1.4.0
139
  Bio==1.7.1
140
+ arxiv==2.1.3