sycommon-python-lib 0.1.56b11__py3-none-any.whl → 0.1.56b13__py3-none-any.whl

sycommon/llm/embedding.py

@@ -25,15 +25,21 @@ class Embedding(metaclass=SingletonMeta):
 
         # 并发信号量
         self.semaphore = asyncio.Semaphore(self.max_concurrency)
+        # 全局默认超时：永不超时（None）
+        self.default_timeout = aiohttp.ClientTimeout(total=None)
 
     async def _get_embeddings_http_async(
         self,
         input: Union[str, List[str]],
         encoding_format: str = None,
         model: str = None,
+        timeout: aiohttp.ClientTimeout = None,
         **kwargs
     ):
         async with self.semaphore:
+            # 优先使用传入的超时，无则用全局默认
+            request_timeout = timeout or self.default_timeout
+
             # 优先使用传入的模型名，无则用默认值
             target_model = model or self.default_embedding_model
             target_base_url = EmbeddingConfig.from_config(target_model).baseUrl
@@ -46,14 +52,23 @@ class Embedding(metaclass=SingletonMeta):
             }
             request_body.update(kwargs)
 
-            async with aiohttp.ClientSession() as session:
-                async with session.post(url, json=request_body) as response:
-                    if response.status != 200:
-                        error_detail = await response.text()
-                        SYLogger.error(
-                            f"Embedding request failed (model: {target_model}): {error_detail}")
-                        return None
-                    return await response.json()
+            try:
+                async with aiohttp.ClientSession(timeout=request_timeout) as session:
+                    async with session.post(url, json=request_body) as response:
+                        if response.status != 200:
+                            error_detail = await response.text()
+                            SYLogger.error(
+                                f"Embedding request failed (model: {target_model}): {error_detail}")
+                            return None
+                        return await response.json()
+            except asyncio.TimeoutError:
+                SYLogger.error(
+                    f"Embedding request timeout (model: {target_model})")
+                return None
+            except Exception as e:
+                SYLogger.error(
+                    f"Embedding request unexpected error (model: {target_model}): {str(e)}")
+                return None
 
     async def _get_reranker_http_async(
         self,
@@ -64,9 +79,13 @@ class Embedding(metaclass=SingletonMeta):
         max_chunks_per_doc: Optional[int] = None,
         return_documents: Optional[bool] = True,
         return_len: Optional[bool] = True,
+        timeout: aiohttp.ClientTimeout = None,
         **kwargs
     ):
         async with self.semaphore:
+            # 优先使用传入的超时，无则用全局默认
+            request_timeout = timeout or self.default_timeout
+
             # 优先使用传入的模型名，无则用默认值
             target_model = model or self.default_reranker_model
             target_base_url = RerankerConfig.from_config(target_model).baseUrl
@@ -84,19 +103,29 @@ class Embedding(metaclass=SingletonMeta):
             }
             request_body.update(kwargs)
 
-            async with aiohttp.ClientSession() as session:
-                async with session.post(url, json=request_body) as response:
-                    if response.status != 200:
-                        error_detail = await response.text()
-                        SYLogger.error(
-                            f"Rerank request failed (model: {target_model}): {error_detail}")
-                        return None
-                    return await response.json()
+            try:
+                async with aiohttp.ClientSession(timeout=request_timeout) as session:
+                    async with session.post(url, json=request_body) as response:
+                        if response.status != 200:
+                            error_detail = await response.text()
+                            SYLogger.error(
+                                f"Rerank request failed (model: {target_model}): {error_detail}")
+                            return None
+                        return await response.json()
+            except asyncio.TimeoutError:
+                SYLogger.error(
+                    f"Rerank request timeout (model: {target_model})")
+                return None
+            except Exception as e:
+                SYLogger.error(
+                    f"Rerank request unexpected error (model: {target_model}): {str(e)}")
+                return None
 
     async def get_embeddings(
         self,
         corpus: List[str],
-        model: str = None
+        model: str = None,
+        timeout: Optional[Union[int, float]] = None
     ):
         """
         获取语料库的嵌入向量，结果顺序与输入语料库顺序一致
@@ -104,12 +133,24 @@ class Embedding(metaclass=SingletonMeta):
         Args:
             corpus: 待生成嵌入向量的文本列表
             model: 可选，指定使用的embedding模型名称，默认使用bge-large-zh-v1.5
+            timeout: 可选，超时时间（秒）：
+                     - 传int/float：表示总超时时间（秒）
+                     - 不传/None：使用默认永不超时配置
         """
+        request_timeout = None
+        if timeout is not None:
+            if isinstance(timeout, (int, float)):
+                request_timeout = aiohttp.ClientTimeout(total=timeout)
+            else:
+                SYLogger.warning(
+                    f"Invalid timeout type: {type(timeout)}, must be int/float, use default timeout")
+
         SYLogger.info(
-            f"Requesting embeddings for corpus: {corpus} (model: {model or self.default_embedding_model}, max_concurrency: {self.max_concurrency})")
-        # 给每个异步任务传入模型名称
+            f"Requesting embeddings for corpus: {corpus} (model: {model or self.default_embedding_model}, max_concurrency: {self.max_concurrency}, timeout: {timeout or 'None'})")
+
+        # 给每个异步任务传入模型名称和超时配置
         tasks = [self._get_embeddings_http_async(
-            text, model=model) for text in corpus]
+            text, model=model, timeout=request_timeout) for text in corpus]
         results = await asyncio.gather(*tasks)
 
         vectors = []
@@ -131,7 +172,8 @@ class Embedding(metaclass=SingletonMeta):
         self,
         top_results: List[str],
         query: str,
-        model: str = None
+        model: str = None,
+        timeout: Optional[Union[int, float]] = None
     ):
         """
         对搜索结果进行重排序
@@ -140,10 +182,23 @@ class Embedding(metaclass=SingletonMeta):
             top_results: 待重排序的文本列表
             query: 排序参考的查询语句
             model: 可选，指定使用的reranker模型名称，默认使用bge-reranker-large
+            timeout: 可选，超时时间（秒）：
+                     - 传int/float：表示总超时时间（秒）
+                     - 不传/None：使用默认永不超时配置
         """
+        request_timeout = None
+        if timeout is not None:
+            if isinstance(timeout, (int, float)):
+                request_timeout = aiohttp.ClientTimeout(total=timeout)
+            else:
+                SYLogger.warning(
+                    f"Invalid timeout type: {type(timeout)}, must be int/float, use default timeout")
+
         SYLogger.info(
-            f"Requesting reranker for top_results: {top_results} (model: {model or self.default_reranker_model}, max_concurrency: {self.max_concurrency})")
-        data = await self._get_reranker_http_async(top_results, query, model=model)
+            f"Requesting reranker for top_results: {top_results} (model: {model or self.default_reranker_model}, max_concurrency: {self.max_concurrency}, timeout: {timeout or 'None'})")
+
+        data = await self._get_reranker_http_async(
+            top_results, query, model=model, timeout=request_timeout)
         SYLogger.info(
             f"Reranker for top_results: {top_results} completed (model: {model or self.default_reranker_model})")
         return data