up

crazy_functions/rag_fns/arxiv_fns/arxiv_downloader.py

@@ -0,0 +1,111 @@
+import logging
+import requests
+import tarfile
+from pathlib import Path
+from typing import Optional, Dict
+
+class ArxivDownloader:
+    """用于下载arXiv论文源码的下载器"""
+    
+    def __init__(self, root_dir: str = "./papers", proxies: Optional[Dict[str, str]] = None):
+        """
+        初始化下载器
+        
+        Args:
+            root_dir: 保存下载文件的根目录
+            proxies: 代理服务器设置，例如 {"http": "http://proxy:port", "https": "https://proxy:port"}
+        """
+        self.root_dir = Path(root_dir)
+        self.root_dir.mkdir(exist_ok=True)
+        self.proxies = proxies
+        
+        # 配置日志
+        logging.basicConfig(
+            level=logging.INFO,
+            format='%(asctime)s - %(levelname)s - %(message)s'
+        )
+        
+    def _download_and_extract(self, arxiv_id: str) -> str:
+        """
+        下载并解压arxiv论文源码
+        
+        Args:
+            arxiv_id: arXiv论文ID，例如"2103.00020"
+            
+        Returns:
+            str: 解压后的文件目录路径
+            
+        Raises:
+            RuntimeError: 当下载失败时抛出
+        """
+        paper_dir = self.root_dir / arxiv_id
+        tar_path = paper_dir / f"{arxiv_id}.tar.gz"
+        
+        # 检查缓存
+        if paper_dir.exists() and any(paper_dir.iterdir()):
+            logging.info(f"Using cached version for {arxiv_id}")
+            return str(paper_dir)
+            
+        paper_dir.mkdir(exist_ok=True)
+        
+        urls = [
+            f"https://arxiv.org/src/{arxiv_id}",
+            f"https://arxiv.org/e-print/{arxiv_id}"
+        ]
+        
+        for url in urls:
+            try:
+                logging.info(f"Downloading from {url}")
+                response = requests.get(url, proxies=self.proxies)
+                if response.status_code == 200:
+                    tar_path.write_bytes(response.content)
+                    with tarfile.open(tar_path, 'r:gz') as tar:
+                        tar.extractall(path=paper_dir)
+                    return str(paper_dir)
+            except Exception as e:
+                logging.warning(f"Download failed for {url}: {e}")
+                continue
+                
+        raise RuntimeError(f"Failed to download paper {arxiv_id}")
+        
+    def download_paper(self, arxiv_id: str) -> str:
+        """
+        下载指定的arXiv论文
+        
+        Args:
+            arxiv_id: arXiv论文ID
+            
+        Returns:
+            str: 论文文件所在的目录路径
+        """
+        return self._download_and_extract(arxiv_id)
+
+def main():
+    """测试下载功能"""
+    # 配置代理（如果需要）
+    proxies = {
+        "http": "http://your-proxy:port",
+        "https": "https://your-proxy:port"
+    }
+    
+    # 创建下载器实例（如果不需要代理，可以不传入proxies参数）
+    downloader = ArxivDownloader(root_dir="./downloaded_papers", proxies=None)
+    
+    # 测试下载一篇论文（这里使用一个示例ID）
+    try:
+        paper_id = "2103.00020"  # 这是一个示例ID
+        paper_dir = downloader.download_paper(paper_id)
+        print(f"Successfully downloaded paper to: {paper_dir}")
+        
+        # 检查下载的文件
+        paper_path = Path(paper_dir)
+        if paper_path.exists():
+            print("Downloaded files:")
+            for file in paper_path.rglob("*"):
+                if file.is_file():
+                    print(f"- {file.relative_to(paper_path)}")
+    except Exception as e:
+        print(f"Error downloading paper: {e}")
+
+if __name__ == "__main__":
+    main()