up

crazy_functions/Arxiv_论文对话.py

@@ -46,10 +46,10 @@ class ArxivRagWorker:
     def __init__(self, user_name: str, llm_kwargs: Dict, arxiv_id: str = None):
         self.user_name = user_name
         self.llm_kwargs = llm_kwargs
-        self.max_concurrent_papers = MAX_CONCURRENT_PAPERS
+        self.max_concurrent_papers = MAX_CONCURRENT_PAPERS  # 存储最大并发数
         self.arxiv_id = self._normalize_arxiv_id(arxiv_id) if arxiv_id else None
 
-        # Initialize base storage directory
+        # 初始化基础存储目录
         self.base_dir = Path(get_log_folder(user_name, plugin_name='rag_cache'))
 
         if self.arxiv_id:
@@ -57,6 +57,7 @@ class ArxivRagWorker:
             self.vector_store_dir = self.checkpoint_dir / "vector_store"
             self.fragment_store_dir = self.checkpoint_dir / "fragments"
         else:
+            # 如果没有 arxiv_id,使用基础目录
             self.checkpoint_dir = self.base_dir
             self.vector_store_dir = self.base_dir / "vector_store"
             self.fragment_store_dir = self.base_dir / "fragments"
@@ -75,11 +76,11 @@ class ArxivRagWorker:
         logger.info(f"Vector store directory: {self.vector_store_dir}")
         logger.info(f"Fragment store directory: {self.fragment_store_dir}")
 
-        # Initialize processing queue and thread pool
+        # 初始化处理队列和线程池
         self.processing_queue = {}
         self.thread_pool = ThreadPoolExecutor(max_workers=MAX_WORKERS)
 
-        # Initialize RAG worker
+        # 初始化RAG worker
         self.rag_worker = LlamaIndexRagWorker(
             user_name=user_name,
             llm_kwargs=llm_kwargs,
@@ -87,24 +88,97 @@ class ArxivRagWorker:
             auto_load_checkpoint=True
         )
 
-        # Initialize arxiv splitter
+        # 初始化arxiv splitter
+        # 初始化 arxiv splitter
         self.arxiv_splitter = ArxivSplitter(
             root_dir=str(self.checkpoint_dir / "arxiv_cache")
         )
-    async def _async_get_fragments(self, arxiv_id: str) -> List[Fragment]:
-        """Async helper to get fragments"""
-        return await self.arxiv_splitter.process(arxiv_id)
+        # 初始化处理队列和线程池
+        self._semaphore = None
+        self._loop = None
+
+    @property
+    def loop(self):
+        """获取当前事件循环"""
+        if self._loop is None:
+            self._loop = asyncio.get_event_loop()
+        return self._loop
+
+    @property
+    def semaphore(self):
+        """延迟创建 semaphore"""
+        if self._semaphore is None:
+            self._semaphore = asyncio.Semaphore(self.max_concurrent_papers)
+        return self._semaphore
+
+
+
+    async def _process_fragments(self, fragments: List[Fragment]) -> None:
+        """并行处理论文片段"""
+        if not fragments:
+            logger.warning("No fragments to process")
+            return
+
+        # 首先添加论文概述
+        overview = {
+            "title": fragments[0].title,
+            "abstract": fragments[0].abstract,
+            "arxiv_id": fragments[0].arxiv_id,
+            "section_tree": fragments[0].section_tree,
+        }
+
+        overview_text = (
+            f"Paper Title: {overview['title']}\n"
+            f"ArXiv ID: {overview['arxiv_id']}\n"
+            f"Abstract: {overview['abstract']}\n"
+            f"Section Tree:{overview['section_tree']}\n"
+            f"Type: OVERVIEW"
+        )
 
-    def _get_fragments_sync(self, arxiv_id: str) -> List[Fragment]:
-        """Synchronous wrapper for async fragment retrieval"""
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
         try:
-            return loop.run_until_complete(self._async_get_fragments(arxiv_id))
-        finally:
-            loop.close()
-    def _process_single_fragment(self, fragment: Fragment, index: int) -> None:
-        """Process a single paper fragment"""
+            # 同步添加概述
+            self.rag_worker.add_text_to_vector_store(overview_text)
+            logger.info(f"Added paper overview for {overview['arxiv_id']}")
+
+            # 并行处理其余片段
+            tasks = []
+            for i, fragment in enumerate(fragments):
+                tasks.append(self._process_single_fragment(fragment, i))
+            await asyncio.gather(*tasks)
+            logger.info(f"Processed {len(fragments)} fragments successfully")
+
+
+            # 保存到本地文件用于调试
+            save_fragments_to_file(
+                fragments,
+                str(self.fragment_store_dir / f"{overview['arxiv_id']}_fragments.json")
+            )
+
+        except Exception as e:
+            logger.error(f"Error processing fragments: {str(e)}")
+            raise
+
+    async def _process_single_fragment(self, fragment: Fragment, index: int) -> None:
+        """处理单个论文片段（改为异步方法）"""
+        try:
+            text = (
+                f"Paper Title: {fragment.title}\n"
+                f"Abstract: {fragment.abstract}\n"
+                f"ArXiv ID: {fragment.arxiv_id}\n"
+                f"Section: {fragment.current_section}\n"
+                f"Section Tree: {fragment.section_tree}\n"
+                f"Content: {fragment.content}\n"
+                f"Bibliography: {fragment.bibliography}\n"
+                f"Type: FRAGMENT"
+            )
+            logger.info(f"Processing fragment {index} for paper {fragment.arxiv_id}")
+            # 如果 add_text_to_vector_store 是异步的，使用 await
+            self.rag_worker.add_text_to_vector_store(text)
+            logger.info(f"Successfully added fragment {index} to vector store")
+
+        except Exception as e:
+            logger.error(f"Error processing fragment {index}: {str(e)}")
+            raise        """处理单个论文片段"""
         try:
             text = (
                 f"Paper Title: {fragment.title}\n"
@@ -124,62 +198,8 @@ class ArxivRagWorker:
             logger.error(f"Error processing fragment {index}: {str(e)}")
             raise
 
-    def _process_fragments(self, fragments: List[Fragment]) -> None:
-        """Process paper fragments in parallel using thread pool"""
-        if not fragments:
-            logger.warning("No fragments to process")
-            return
-
-        # First add paper overview
-        overview = {
-            "title": fragments[0].title,
-            "abstract": fragments[0].abstract,
-            "arxiv_id": fragments[0].arxiv_id,
-            "section_tree": fragments[0].section_tree,
-        }
-
-        overview_text = (
-            f"Paper Title: {overview['title']}\n"
-            f"ArXiv ID: {overview['arxiv_id']}\n"
-            f"Abstract: {overview['abstract']}\n"
-            f"Section Tree:{overview['section_tree']}\n"
-            f"Type: OVERVIEW"
-        )
-
-        try:
-            # Add overview synchronously
-            self.rag_worker.add_text_to_vector_store(overview_text)
-            logger.info(f"Added paper overview for {overview['arxiv_id']}")
-
-            # Process fragments in parallel using thread pool
-            with ThreadPoolExecutor(max_workers=10) as executor:
-                # Submit all fragments for processing
-                futures = [
-                    executor.submit(self._process_single_fragment, fragment, i)
-                    for i, fragment in enumerate(fragments)
-                ]
-
-                # Wait for all tasks to complete and handle any exceptions
-                for future in futures:
-                    try:
-                        future.result()
-                    except Exception as e:
-                        logger.error(f"Error processing fragment: {str(e)}")
-
-            logger.info(f"Processed {len(fragments)} fragments successfully")
-
-            # Save to local file for debugging
-            save_fragments_to_file(
-                fragments,
-                str(self.fragment_store_dir / f"{overview['arxiv_id']}_fragments.json")
-            )
-
-        except Exception as e:
-            logger.error(f"Error processing fragments: {str(e)}")
-            raise
-
-    def process_paper(self, arxiv_id: str) -> bool:
-        """Process paper main function - mixed sync/async version"""
+    async def process_paper(self, arxiv_id: str) -> bool:
+        """处理论文主函数"""
         try:
             arxiv_id = self._normalize_arxiv_id(arxiv_id)
             logger.info(f"Starting to process paper: {arxiv_id}")
@@ -190,29 +210,30 @@ class ArxivRagWorker:
                 logger.info(f"Paper {arxiv_id} already processed")
                 return True
 
-            # Create processing task
+            # 创建处理任务
             task = ProcessingTask(arxiv_id=arxiv_id)
             self.processing_queue[arxiv_id] = task
             task.status = "processing"
 
-            # Download and split paper using the sync wrapper
-            fragments = self._get_fragments_sync(arxiv_id)
+            async with self.semaphore:
+                # 下载和分割论文
+                fragments = await self.arxiv_splitter.process(arxiv_id)
 
-            if not fragments:
-                raise ValueError(f"No fragments extracted from paper {arxiv_id}")
+                if not fragments:
+                    raise ValueError(f"No fragments extracted from paper {arxiv_id}")
 
-            logger.info(f"Got {len(fragments)} fragments from paper {arxiv_id}")
+                logger.info(f"Got {len(fragments)} fragments from paper {arxiv_id}")
 
-            # Process fragments
-            self._process_fragments(fragments)
+                # 处理片段
+                await self._process_fragments(fragments)
 
-            # Mark as completed
-            paper_path.touch()
-            task.status = "completed"
-            task.fragments = fragments
+                # 标记完成
+                paper_path.touch()
+                task.status = "completed"
+                task.fragments = fragments
 
-            logger.info(f"Successfully processed paper {arxiv_id}")
-            return True
+                logger.info(f"Successfully processed paper {arxiv_id}")
+                return True
 
         except Exception as e:
             logger.error(f"Error processing paper {arxiv_id}: {str(e)}")
@@ -220,8 +241,19 @@ class ArxivRagWorker:
                 self.processing_queue[arxiv_id].status = "failed"
                 self.processing_queue[arxiv_id].error = str(e)
             return False
-    def wait_for_paper(self, arxiv_id: str, timeout: float = 300.0) -> bool:
-        """Wait for paper processing to complete - synchronous version"""
+
+    def _normalize_arxiv_id(self, input_str: str) -> str:
+        """规范化ArXiv ID"""
+        if 'arxiv.org/' in input_str.lower():
+            if '/pdf/' in input_str:
+                arxiv_id = input_str.split('/pdf/')[-1]
+            else:
+                arxiv_id = input_str.split('/abs/')[-1]
+            return arxiv_id.split('v')[0].strip()
+        return input_str.split('v')[0].strip()
+
+    async def wait_for_paper(self, arxiv_id: str, timeout: float = 300.0) -> bool:
+        """等待论文处理完成"""
         try:
             start_time = datetime.now()
             while True:
@@ -235,27 +267,16 @@ class ArxivRagWorker:
                 if task.status == "failed":
                     return False
 
-                # Check timeout
+                # 检查超时
                 if (datetime.now() - start_time).total_seconds() > timeout:
                     logger.error(f"Processing paper {arxiv_id} timed out")
                     return False
 
-                time.sleep(0.1)
+                await asyncio.sleep(0.1)
         except Exception as e:
             logger.error(f"Error waiting for paper {arxiv_id}: {str(e)}")
             return False
 
-    def _normalize_arxiv_id(self, input_str: str) -> str:
-        """Normalize ArXiv ID"""
-        if 'arxiv.org/' in input_str.lower():
-            if '/pdf/' in input_str:
-                arxiv_id = input_str.split('/pdf/')[-1]
-            else:
-                arxiv_id = input_str.split('/abs/')[-1]
-            return arxiv_id.split('v')[0].strip()
-        return input_str.split('v')[0].strip()
-
-
     def retrieve_and_generate(self, query: str) -> str:
         """检索相关内容并生成提示词"""
         try:
@@ -332,10 +353,20 @@ def Arxiv论文对话(txt: str, llm_kwargs: Dict, plugin_kwargs: Dict, chatbot:
         chatbot.append((txt, "正在处理论文，请稍等..."))
         yield from update_ui(chatbot=chatbot, history=history)
 
-        success = worker.process_paper(txt)
-        if success:
-            arxiv_id = worker._normalize_arxiv_id(txt)
-            success = worker.wait_for_paper(arxiv_id)
+        # 创建事件循环来处理异步调用
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            # 运行异步处理函数
+            success = loop.run_until_complete(worker.process_paper(txt))
+            if success:
+                arxiv_id = worker._normalize_arxiv_id(txt)
+                success = loop.run_until_complete(worker.wait_for_paper(arxiv_id))
+                if success:
+                    # 执行自动分析
+                    yield from worker.auto_analyze_paper(chatbot, history, system_prompt)
+        finally:
+            loop.close()
 
         if not success:
             chatbot[-1] = (txt, "论文处理失败，请检查论文ID是否正确或稍后重试。")