Update tools/arxiv_tool.py
Browse files- tools/arxiv_tool.py +10 -39
tools/arxiv_tool.py
CHANGED
@@ -4,25 +4,19 @@ from datetime import datetime, timedelta
|
|
4 |
import json
|
5 |
import os
|
6 |
from typing import List, Dict
|
7 |
-
from smolagents import Tool
|
8 |
|
9 |
class ArxivSearchTool(Tool):
|
10 |
name = "search_arxiv"
|
11 |
description = "Search ArXiv for papers matching the query"
|
12 |
-
|
|
|
|
|
|
|
13 |
output_type = List[Dict]
|
14 |
|
15 |
-
def
|
16 |
-
|
17 |
-
"""Search ArXiv using their API.
|
18 |
-
|
19 |
-
Args:
|
20 |
-
query: Search query string
|
21 |
-
max_results: Maximum number of results to return
|
22 |
-
|
23 |
-
Returns:
|
24 |
-
List[Dict]: List of paper results with metadata
|
25 |
-
"""
|
26 |
try:
|
27 |
# Construct the API URL
|
28 |
base_url = 'http://export.arxiv.org/api/query?'
|
@@ -69,10 +63,10 @@ class ArxivSearchTool(Tool):
|
|
69 |
class LatestPapersTool(Tool):
|
70 |
name = "get_latest_papers"
|
71 |
description = "Get papers from the last N days from saved results"
|
72 |
-
|
73 |
output_type = List[Dict]
|
74 |
|
75 |
-
def
|
76 |
papers = []
|
77 |
base_dir = "daily_papers"
|
78 |
|
@@ -90,27 +84,4 @@ class LatestPapersTool(Tool):
|
|
90 |
day_papers = json.load(f)
|
91 |
papers.extend(day_papers)
|
92 |
|
93 |
-
return papers
|
94 |
-
|
95 |
-
def save_daily_papers(output_dir: str = "daily_papers") -> List[Dict]:
|
96 |
-
"""Helper function to save daily papers - not exposed as a tool"""
|
97 |
-
os.makedirs(output_dir, exist_ok=True)
|
98 |
-
today = datetime.now().strftime("%Y-%m-%d")
|
99 |
-
|
100 |
-
arxiv_tool = ArxivSearchTool()
|
101 |
-
papers = arxiv_tool(
|
102 |
-
query='cat:cs.AI OR cat:cs.LG OR cat:cs.CL OR "artificial intelligence"',
|
103 |
-
max_results=100
|
104 |
-
)
|
105 |
-
|
106 |
-
# Filter for papers published today
|
107 |
-
today_papers = [
|
108 |
-
paper for paper in papers
|
109 |
-
if paper.get('published', '').startswith(today)
|
110 |
-
]
|
111 |
-
|
112 |
-
output_file = os.path.join(output_dir, f"ai_papers_{today}.json")
|
113 |
-
with open(output_file, 'w', encoding='utf-8') as f:
|
114 |
-
json.dump(today_papers, f, indent=2)
|
115 |
-
|
116 |
-
return today_papers
|
|
|
4 |
import json
|
5 |
import os
|
6 |
from typing import List, Dict
|
7 |
+
from smolagents.tools import Tool
|
8 |
|
9 |
class ArxivSearchTool(Tool):
|
10 |
name = "search_arxiv"
|
11 |
description = "Search ArXiv for papers matching the query"
|
12 |
+
inputs = {
|
13 |
+
'query': {'type': 'string', 'description': 'The search query for papers'},
|
14 |
+
'max_results': {'type': 'integer', 'description': 'Maximum number of results to return'}
|
15 |
+
}
|
16 |
output_type = List[Dict]
|
17 |
|
18 |
+
def forward(self, query: str = "artificial intelligence",
|
19 |
+
max_results: int = 50) -> List[Dict]:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
20 |
try:
|
21 |
# Construct the API URL
|
22 |
base_url = 'http://export.arxiv.org/api/query?'
|
|
|
63 |
class LatestPapersTool(Tool):
|
64 |
name = "get_latest_papers"
|
65 |
description = "Get papers from the last N days from saved results"
|
66 |
+
inputs = {'days_back': {'type': 'integer', 'description': 'Number of days to look back'}}
|
67 |
output_type = List[Dict]
|
68 |
|
69 |
+
def forward(self, days_back: int = 1) -> List[Dict]:
|
70 |
papers = []
|
71 |
base_dir = "daily_papers"
|
72 |
|
|
|
84 |
day_papers = json.load(f)
|
85 |
papers.extend(day_papers)
|
86 |
|
87 |
+
return papers
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|